diff --git a/.buildkite/pipeline-intel.yaml b/.buildkite/pipeline-intel.yaml
index ff7302ad2..d3a2655bf 100644
--- a/.buildkite/pipeline-intel.yaml
+++ b/.buildkite/pipeline-intel.yaml
@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.24.0"
+          VLLM_VERSION: "v0.25.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false
diff --git a/.buildkite/pipeline-npu.yaml b/.buildkite/pipeline-npu.yaml
index 161b0e944..d035ffc94 100644
--- a/.buildkite/pipeline-npu.yaml
+++ b/.buildkite/pipeline-npu.yaml
@@ -78,8 +78,89 @@ steps:
       echo "--- Image pushed successfully"
       echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
 
+  - label: ":buildkit: Build and Push NPU Test Image - A3"
+    key: image-build-a3
+    agents:
+      queue: "ascend-a2b3"
+      resource_class: "npu-1"
+    plugins:
+      - kubernetes:
+          metadata:
+            annotations:
+              vault.hashicorp.com/agent-init-first: "true"
+              vault.hashicorp.com/agent-inject: "true"
+              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
+              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
+              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
+              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
+              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
+              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
+              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
+              vault.hashicorp.com/agent-pre-populate-only: "true"
+              vault.hashicorp.com/agent-run-as-group: "1000"
+              vault.hashicorp.com/agent-run-as-user: "1000"
+              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
+              vault.hashicorp.com/role: ascend-gha-runners
+              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
+              vault.hashicorp.com/tls-skip-verify: "true"
+          podSpecPatch:
+            volumes:
+            - name: token-vol
+              projected:
+                defaultMode: 420
+                sources:
+                - serviceAccountToken:
+                    audience: api
+                    expirationSeconds: 600
+                    path: token
+    env:
+      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-a3"
+      IMAGE_NAME: "vllm-omni-ci-npu"
+      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
+      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
+    command: |
+      set -ex
+
+      echo "--- Building and pushing NPU Test Image"
+      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+      echo "buildkitd address: $${BUILDKITD_ADDR}"
+
+      if ! command -v buildctl &> /dev/null; then
+        echo "Installing buildctl..."
+        mkdir -p /tmp/buildkit
+        BUILDKIT_VERSION="v0.29.0"
+        wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
+        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
+        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
+      fi
+
+      export DOCKER_CONFIG=/home/user/.docker
+      buildctl \
+        --addr="$${BUILDKITD_ADDR}" \
+        --tlscacert=/home/user/.docker/ca.pem \
+        --tlscert=/home/user/.docker/cert.pem \
+        --tlskey=/home/user/.docker/key.pem \
+        build \
+        --frontend dockerfile.v0 \
+        --local context=. \
+        --local dockerfile=./docker \
+        --opt filename=Dockerfile.npu.ci.a3 \
+        --secret id=dockerconfig,src=/home/user/.docker/config.json \
+        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
+        --progress=plain
+
+      echo "--- Image pushed successfully"
+      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
+
   - label: "Upload Ready Pipeline"
-    depends_on: image-build
+    depends_on:
+      - image-build
+      - image-build-a3
     key: upload-ready-pipeline
     if: >-
       build.branch != "main" && (
@@ -94,7 +175,9 @@ steps:
 
   # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
   - label: "Upload Nightly Pipeline"
-    depends_on: image-build
+    depends_on:
+      - image-build
+      - image-build-a3
     key: upload-nightly-pipeline
     if: >-
       (build.branch == "main" && build.env("NIGHTLY") == "1") ||
diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh
index fa4f00e29..36860522b 100755
--- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh
+++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh
@@ -13,12 +13,33 @@ container_name="xpu_${BUILDKITE_COMMIT}_$(
 )"
 
 cd "${omni_source_dir}"
+
+# The XPU base image is ~37GB; the default gzip layer exporter is single-threaded
+# and takes ~16min to compress it. zstd is multi-threaded (uses all cores) at a
+# similar ratio; set EXPORT_COMPRESSION=uncompressed to skip compression entirely
+# for local-only images. This requires the containerd image store (buildx docker
+# driver), which is the default here.
+EXPORT_COMPRESSION="${EXPORT_COMPRESSION:-zstd}"
+if [ "${EXPORT_COMPRESSION}" = "uncompressed" ]; then
+    export_args=(--output "type=image,name={{IMAGE}},compression=uncompressed")
+else
+    export_args=(--output "type=image,name={{IMAGE}},compression=${EXPORT_COMPRESSION},compression-level=3,force-compression=true")
+fi
+
+docker_build() {
+    # $1 = image name; remaining args passed through to docker build.
+    local image="$1"
+    shift
+    local out=("${export_args[@]/'{{IMAGE}}'/${image}}")
+    docker build "${out[@]}" "$@" -f docker/Dockerfile.xpu .
+}
+
 if [ -z "$(docker images -q "${base_image_name}")" ]; then
-    docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
+    docker_build "${base_image_name}" --target vllm-base --build-arg "VLLM_VERSION=${VLLM_VERSION}"
 fi
 
 # Try building the docker image
-docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
+docker_build "${image_name}" --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}"
 
 # Setup cleanup
 remove_docker_container() {
diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml
index 057a4915b..9e3372f44 100644
--- a/.buildkite/test-merge.yml
+++ b/.buildkite/test-merge.yml
@@ -638,11 +638,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:
diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml
index ddb7960c2..fb94ce5a4 100644
--- a/.buildkite/test-nightly.yml
+++ b/.buildkite/test-nightly.yml
@@ -607,6 +607,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
@@ -627,6 +659,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:
@@ -670,6 +704,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:
@@ -1299,6 +1334,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:
@@ -1428,9 +1501,12 @@ steps:
             set +e
             pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
             EXIT1=$$?
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
+            EXIT2=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT1
+            if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi
+            exit $$EXIT2
         agents:
           queue: "mithril-h100-pool"
         plugins:
diff --git a/.buildkite/test-npu-nightly.yaml b/.buildkite/test-npu-nightly.yaml
index 2c2372f4c..4dbda17f9 100644
--- a/.buildkite/test-npu-nightly.yaml
+++ b/.buildkite/test-npu-nightly.yaml
@@ -22,3 +22,31 @@ steps:
           HF_TOKEN: "${HF_TOKEN}"
         commands:
           - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
+  - group: ":card_index_dividers: TTS Model Test"
+    key: nightly-tts-test-group
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
+    steps:
+      - label: ":full_moon: TTS · Perf Test"
+        agents:
+          queue: "ascend-a3"
+          resource_class: "npu-2"
+        image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a3"
+        plugins:
+          - kubernetes:
+              podSpecPatch:
+                imagePullSecrets:
+                  - name: swr-secret
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+          HF_TOKEN: "${HF_TOKEN}"
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_tts.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml
index d12267fe0..c9b60fe2c 100644
--- a/.buildkite/test-ready.yml
+++ b/.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"
diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md
index 906f3c485..76c894473 100644
--- a/.claude/skills/add-diffusion-model/SKILL.md
+++ b/.claude/skills/add-diffusion-model/SKILL.md
@@ -334,7 +334,7 @@ Classify the model's **CI priority** first:
 | **L1** | `tests/diffusion/models/{slug}/`, `tests/diffusion/cache/`, transformer unit tests | `core_model` + `cpu` | `test-ready.yml` | Weight remap, `_sp_plan`, cache enabler registration, shape contracts |
 | **L2** | `tests/e2e/online_serving/test_{slug}.py` (and offline if the category is offline-first) | **`core_model` + `advanced_model`** (both on baseline smoke) + `diffusion` + `@hardware_test` / `hardware_marks` | `test-ready.yml` | Default deploy smoke — minimal `num_inference_steps`, single prompt |
 | **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` when offline matters | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `diffusion`) | `test-merge.yml` **or** merged into nightly diffusion function job | Real weights, streaming/API paths, LoRA/offload smoke |
-| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py` (+ offline expansion if needed) | `full_model` + `diffusion` | `test-nightly.yml` (X2I / X2V / X2A function groups) | Feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); perf → `tests/dfx/perf/tests/test_{model}_vllm_omni.json` |
+| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py` (+ offline expansion if needed) | `full_model` + `diffusion` | `test-nightly.yml` (X2I / X2V / X2A function groups) | Feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); perf → `tests/dfx/perf/tests/test_{model}_vllm_omni.json` with per-case **`mark`** (`hardware_marks` + `full_model` + `diffusion`) |
 
 **L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`).
 
diff --git a/.claude/skills/vllm-omni-test/SKILL.md b/.claude/skills/vllm-omni-test/SKILL.md
index 40a81f4cf..dcc2506da 100644
--- a/.claude/skills/vllm-omni-test/SKILL.md
+++ b/.claude/skills/vllm-omni-test/SKILL.md
@@ -97,7 +97,7 @@ Existing references: `tests/e2e/offline_inference/test_qwen2_5_omni.py` (L2-styl
 | **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `omni_runner_handler`. **Function (isolation only):** `omni_runner_function` + `omni_runner_handler_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required |
 | **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `openai_client`. **Function (isolation only):** `omni_server_function` + `openai_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` |
 | **Documentation / runnable examples** | `tests/examples/offline_inference/`, `tests/examples/online_serving/` | **Offline docs (preferred):** extract Python/Bash blocks from the doc README (e.g. `ReadmeSnippet.extract_readme_snippets`), `pytest.mark.parametrize` each snippet, run via `example_runner.run` with a stable `output_subfolder`. **Online docs:** copy client/request scripts into dedicated tests and keep them in sync with the doc page. | Usually **L4**: `advanced_model`, often `example` plus hardware marks matching the nightly docs-example job (see `.buildkite/test-nightly.yml`). Full conventions: [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) (introduced in [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910): naming, output directory layout, skip rules, avoid trimming `num_inference_steps` without a strong CI reason). |
-| **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: `full_model` + `benchmark`; wire `test-nightly.yml` Perf steps |
+| **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: JSON `mark` with `full_model` + `omni`/`tts`/`diffusion`; wire `test-nightly.yml` Perf steps |
 | **Invalid parameter / negative HTTP validation** | `tests/dfx/reliability/invalid_param_test/` | Live `omni_server` + low-level `send_*_http_request` with `err_code` / `err_message` | `pytest.mark.slow` + `omni` / `tts` / `diffusion` + `@hardware_marks` (`H100` or `L4`); CI in **`test-weekly.yml`** (not ready/merge/nightly) |
 
 **If the user’s test plan includes invalid parameter validation / invalid params / negative HTTP / 400 validation:** do **not** add those `test_*` functions to `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move or author them** under `tests/dfx/reliability/invalid_param_test/` in the **endpoint-matching script** (see **Invalid parameter validation** below). Success-path e2e and invalid-param dfx tests must stay in separate modules.
@@ -270,7 +270,7 @@ def test_text_to_video_001(omni_server, openai_client) -> None:
 
 *Documentation example tests:* follow the **Preferred Test Strategy** in [l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md): dynamic extraction for offline READMEs; explicit copied client code for online pages until extraction is justified; use the documented **naming**, **output directory** (page folder + case id), and **skipping** rules (e.g. Gradio-only scripts).
 
-*Performance tests:* add or extend entries under `tests/perf/` (and JSON configs where the project uses them), with **explicit baselines** and the same marker/run-level pairing as the CI step that will execute them.
+*Performance tests:* add or extend entries under `tests/dfx/perf/tests/` (and JSON configs where the project uses them), with **explicit baselines**, **`mark`** on each case (`hardware_marks` + `full_model` + type marker), and the same nightly Perf step pattern as in-tree configs.
 
 **3. Cross-cutting rules**
 
@@ -592,27 +592,31 @@ Implementation strategy:
 **When the user explicitly asks for L4 perf (or full L4 including perf):**
 
 1. **Do not** put throughput/latency assertions inside `test_*_expansion.py` — perf uses the **dfx benchmark harness**, not `omni_server` e2e fixtures.
-2. Add or extend a **JSON config** under `tests/dfx/perf/tests/` (mirror in-tree names: `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`).
-3. Each JSON entry: `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]` with explicit **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`, …).
+2. Add or extend a **JSON config** under `tests/dfx/perf/tests/` (mirror in-tree names: `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, `test_qwen3_omni_async_chunk.json`).
+3. Each JSON **case**: `test_name`, optional **`mark`** (`hardware_marks` required when present; `marks`: `full_model` + `omni`/`tts`/`diffusion`), `server_params` (`model`, `serve_args` or `stage_config_name`), `benchmark_params[]` with **`name`** and explicit **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`, …). Diffusion cases also set `server_type` (e.g. `"vllm-omni"`) and usually `benchmark_endpoint`.
 4. Run locally via the matching script (model type → runner):
 
 | Model type | Runner script | Example config |
 |------------|---------------|----------------|
 | **Diffusion X2I/X2V** | `tests/dfx/perf/scripts/run_diffusion_benchmark.py` | `tests/dfx/perf/tests/test_qwen_image_vllm_omni.json` |
 | **TTS** | `tests/dfx/perf/scripts/run_benchmark.py` | `tests/dfx/perf/tests/test_tts.json` (shared) or `test_{slug}.json` (dedicated; use when the model must not join the shared nightly matrix before integration — e.g. `test_voxcpm2.json`, `test_coqui_tts.json`) |
-| **Omni** | (see existing Omni Perf steps — `run_benchmark.py` / dedicated omni configs in-tree) | `tests/dfx/perf/tests/test_omni*.json` |
+| **Omni** | `tests/dfx/perf/scripts/run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` |
 
-5. Wire **`test-nightly.yml` · Perf Test** step for that model (separate from Function Test): export `DIFFUSION_BENCHMARK_DIR` / `BENCHMARK_DIR`, run pytest on the script with `--test-config-file`, **upload artifacts** (`buildkite-agent artifact upload`), often **multi-GPU H100** for diffusion perf.
+5. Wire **`test-nightly.yml` · Perf Test** step for that model (separate from Function Test): export `DIFFUSION_BENCHMARK_DIR` / `BENCHMARK_DIR`, run pytest on the script with **`--test-config-file`** (nightly perf steps do not use `-m`), **upload artifacts** (`buildkite-agent artifact upload`), often **multi-GPU H100** for diffusion perf.
 6. **One benchmark scenario → one `test_name` block** in JSON (same spirit as one `test_*` per function case). Combine server `serve_args` + workload in one entry; use multiple `benchmark_params` rows for size/step sweeps under the same server config.
 
-**Perf local command (diffusion example):**
+**Perf local commands:**
 
 ```bash
 cd tests
 export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
 export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+# Single file (CI-like)
 pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \
   --test-config-file dfx/perf/tests/test_<model>_vllm_omni.json
+# Bulk load + filter by JSON mark
+pytest -sv dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"
+pytest -sv dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
 ```
 
 **Clarify in the test plan** which L4 pillars you deliver: `Function only` | `Function + Perf` | `Function + Accuracy + Perf`.
diff --git a/.claude/skills/vllm-omni-test/references/test-routing.md b/.claude/skills/vllm-omni-test/references/test-routing.md
index 445a78399..f78e1c77d 100644
--- a/.claude/skills/vllm-omni-test/references/test-routing.md
+++ b/.claude/skills/vllm-omni-test/references/test-routing.md
@@ -127,8 +127,12 @@ pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and
 cd tests
 export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
 export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+# CI-like: single JSON file (nightly Perf Test steps)
 pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \
   --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json
+# Local bulk load: all *.json under tests/dfx/perf/tests/, filter by JSON mark
+pytest -sv dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"
+pytest -sv dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
 ```
 
 Broad marker sweep (when no explicit file shard):
@@ -150,7 +154,26 @@ Nightly **L4** for a model is often **multiple jobs** in `test-nightly.yml`, not
 
 **Rule:** “L4 functional cases” → **Function only** by default. State that Perf/Accuracy are separate pillars; add them only when requested.
 
-**Diffusion perf JSON** (`tests/dfx/perf/tests/test_<slug>_vllm_omni.json`): array of objects with `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]`, and per-workload **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`). Copy structure from `test_qwen_image_vllm_omni.json` in-tree.
+**Diffusion perf JSON** (`tests/dfx/perf/tests/test_<slug>_vllm_omni.json`): array of **case objects**, each with:
+
+| Field | Required | Notes |
+|-------|----------|-------|
+| `test_name` | Yes | Unique server/workload id |
+| `mark` | Recommended | When present: **`hardware_marks` required** + optional `marks` (`full_model`, `diffusion`, `local_model`, …). Parsed by `resolve_pytest_marks`; enables local `-m` filtering. |
+| `server_type` | Diffusion | e.g. `"vllm-omni"` — routes case to `run_diffusion_benchmark.py` |
+| `benchmark_endpoint` | Diffusion | e.g. `/v1/videos`, `/v1/images/generations` |
+| `server_params` | Yes | `model`, `serve_args`, … |
+| `benchmark_params[]` | Yes | Each row: **`name`** (pytest id suffix), workload fields, **`baseline`** |
+
+Copy structure from `test_qwen_image_vllm_omni.json` or `test_cosmos3_vllm_omni.json` (2-GPU mark example).
+
+**Omni perf JSON**: `test_qwen3_omni_*.json` — same `mark` shape with `"marks": ["full_model", "omni"]`; run via `run_benchmark.py`.
+
+**Runner split**: `is_diffusion_perf_config()` → diffusion when `server_type` is set or `"diffusion"` ∈ `mark.marks`; otherwise omni/tts → `run_benchmark.py`.
+
+**Param ids**: `{test_name}-{benchmark_params.name}` (e.g. `test_tts-p0`, not `test_tts-0`).
+
+**Result filenames**: runtime hardware from `get_runtime_resource_label()`; `H100` omitted on default CI pool.
 
 **Diffusion perf nightly step** (under **Diffusion X2I(&A&T)**; copy full `kubernetes` plugins from `· Perf Test · Qwen-Image`):
 
@@ -177,9 +200,9 @@ Nightly **L4** for a model is often **multiple jobs** in `test-nightly.yml`, not
 
 | Model type | Perf runner | Config example |
 |------------|-------------|----------------|
-| Diffusion X2I/X2V | `dfx/perf/scripts/run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json` |
-| TTS | `dfx/perf/scripts/run_benchmark.py` | `test_tts.json` (`BENCHMARK_DIR`) |
-| Omni | `run_benchmark.py` / omni configs in-tree | `test_omni*.json` |
+| Diffusion X2I/X2V | `dfx/perf/scripts/run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json` |
+| TTS | `dfx/perf/scripts/run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` |
+| Omni | `dfx/perf/scripts/run_benchmark.py` | `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_no_async_chunk.json`, … |
 
 Do **not** put throughput/latency baselines inside `test_*_expansion.py` — that belongs in the dfx perf JSON + nightly Perf job.
 
diff --git a/.gitignore b/.gitignore
index a175e19e3..224dac867 100644
--- a/.gitignore
+++ b/.gitignore
@@ -208,6 +208,7 @@ cache/
 .cache/
 diffusion_cache/
 kv_cache/
+!vllm_omni/experimental/ar_diffusion/kv_cache/
 
 # Logs
 logs/
@@ -247,6 +248,7 @@ datasets/
 
 # Output files
 outputs/
+!vllm_omni/outputs/
 results/
 generated/
 output_*/
diff --git a/README.md b/README.md
index e300b6b8d..e7fc6912e 100644
--- a/README.md
+++ b/README.md
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
+- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
 - [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
 - [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
 - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
@@ -31,9 +32,9 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 [vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
 
-- **Omni-modality**: Text, image, video, and audio data processing
+- **Omni-modality**: Text, image, audio, video, and action data processing
 - **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
-- **Heterogeneous outputs**: from traditional text generation to multimodal outputs
+- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs
 
 <p align="center">
   <picture>
@@ -57,9 +58,10 @@ vLLM-Omni is flexible and easy to use with:
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
+- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
 - **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
 - **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
 
 ## Getting Started
 
diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py
index 368dc5dcb..f8ad1a71e 100644
--- a/benchmarks/diffusion/backends.py
+++ b/benchmarks/diffusion/backends.py
@@ -28,6 +28,7 @@ class RequestFuncInput:
     slo_ms: float | None = None
     extra_body: dict[str, Any] = field(default_factory=dict)
     image_paths: list[str] | None = None
+    video_paths: list[str] | None = None
     request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
     default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
 
@@ -338,6 +339,11 @@ async def async_request_v1_videos(
         form.add_field(k, str(v))
 
     image_file = None
+    if input.image_paths and input.video_paths:
+        output.error = "Only one of image_paths or video_paths can be provided"
+        output.success = False
+        return output
+
     if input.image_paths and len(input.image_paths) > 0:
         image_path = input.image_paths[0]
         image_file = open(image_path, "rb")
@@ -345,7 +351,16 @@ async def async_request_v1_videos(
             "input_reference",
             image_file,
             filename=os.path.basename(image_path),
-            content_type="application/octet-stream",
+            content_type=_guess_mime_type(image_path),
+        )
+    elif input.video_paths and len(input.video_paths) > 0:
+        video_path = input.video_paths[0]
+        image_file = open(video_path, "rb")
+        form.add_field(
+            "input_reference",
+            image_file,
+            filename=os.path.basename(video_path),
+            content_type=_guess_mime_type(video_path),
         )
 
     job_id = None
diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py
index 7e0bd4ed1..8f7a02858 100644
--- a/benchmarks/diffusion/diffusion_benchmark_serving.py
+++ b/benchmarks/diffusion/diffusion_benchmark_serving.py
@@ -73,6 +73,10 @@ Usage:
             {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
         ]'
 
+    v2v:
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --endpoint /v1/videos --dataset random --task v2v --num-prompts 1 \
+        --height 720 --width 1280 --fps 24 --num-frames 189 --num-inference-steps 35
 
 """
 
@@ -542,11 +546,17 @@ class TraceDataset(BaseDataset):
         if not image_paths:
             single = row.get("image_path")
             image_paths = [single] if single else None
+        video_paths = row.get("video_paths")
+        if not video_paths:
+            single_video = row.get("video_path")
+            video_paths = [single_video] if single_video else None
 
         if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]:
             raise ValueError(
                 f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
             )
+        if not video_paths and self.args.task == "v2v":
+            raise ValueError("Task v2v requires video input, but no video_path or video_paths found in trace row.")
 
         override_w = self.args.width
         override_h = self.args.height
@@ -570,6 +580,7 @@ class TraceDataset(BaseDataset):
             timestamp=timestamp,
             slo_ms=slo_ms,
             image_paths=image_paths,
+            video_paths=video_paths,
             request_id=str(row.get("request_id")) if row.get("request_id") is not None else str(uuid.uuid4()),
         )
 
@@ -589,6 +600,8 @@ class CustomDataset(BaseDataset):
     - seed (optional): Random seed
     - image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
     - image_urls (optional): List of input image URLs (alternative to image_paths)
+    - video_paths (optional): List of input video paths for v2v tasks
+    - video_urls (optional): List of input video URLs (alternative to video_paths)
 
     Example JSONL for ti2i:
     {"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
@@ -660,6 +673,42 @@ class CustomDataset(BaseDataset):
 
         return None
 
+    def _resolve_video_paths(self, item: dict) -> list[str] | None:
+        if "video_paths" in item and item["video_paths"]:
+            paths = item["video_paths"]
+            if isinstance(paths, str):
+                paths = [paths]
+
+            valid_paths = []
+            for path in paths:
+                if os.path.exists(path):
+                    valid_paths.append(path)
+                else:
+                    raise ValueError(f"Video file not found: {path}")
+
+            return valid_paths if valid_paths else None
+
+        if "video_urls" in item and item["video_urls"]:
+            urls = item["video_urls"]
+            if isinstance(urls, str):
+                urls = [urls]
+
+            downloaded_paths = []
+            for url in urls:
+                try:
+                    response = requests.get(url, timeout=30)
+                    response.raise_for_status()
+                    suffix = os.path.splitext(url)[1] or ".mp4"
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                        tmp.write(response.content)
+                        downloaded_paths.append(tmp.name)
+                except Exception as e:
+                    raise ValueError(f"Failed to download video from {url}: {e}")
+
+            return downloaded_paths if downloaded_paths else None
+
+        return None
+
     def __len__(self) -> int:
         return len(self.data)
 
@@ -672,11 +721,22 @@ class CustomDataset(BaseDataset):
         height = item.get("height", self.args.height)
         num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
         seed = item.get("seed", self.args.seed)
-        reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
+        reserved_keys = {
+            "prompt",
+            "width",
+            "height",
+            "num_inference_steps",
+            "seed",
+            "image_paths",
+            "image_urls",
+            "video_paths",
+            "video_urls",
+        }
         extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
 
         # Handle image paths/URLs
         image_paths = self._resolve_image_paths(item)
+        video_paths = self._resolve_video_paths(item)
 
         return RequestFuncInput(
             prompt=prompt,
@@ -684,6 +744,7 @@ class CustomDataset(BaseDataset):
             model=self.model,
             seed=seed,
             image_paths=image_paths,
+            video_paths=video_paths,
             extra_body=extra_body,
             width=width,
             height=height,
@@ -737,6 +798,10 @@ class RandomDataset(BaseDataset):
             self._random_image_path = self._generate_random_image_paths()
         else:
             self._random_image_path = None
+        if self.args.task == "v2v":
+            self._random_video_path = self._generate_random_video_paths()
+        else:
+            self._random_video_path = None
 
     def __len__(self) -> int:
         return self.num_prompts
@@ -756,13 +821,15 @@ class RandomDataset(BaseDataset):
         if self._sampled_requests:
             profile = self._sampled_requests[idx]
             params.update(profile)
+        prompt = str(params.pop("prompt", f"Random prompt {idx} for benchmarking diffusion models"))
         return RequestFuncInput(
-            prompt=f"Random prompt {idx} for benchmarking diffusion models",
+            prompt=prompt,
             api_url=self.api_url,
             model=self.model,
             seed=self.args.seed,
             extra_body=extra_body,
             image_paths=self._random_image_path,
+            video_paths=self._random_video_path,
             **params,
         )
 
@@ -781,6 +848,35 @@ class RandomDataset(BaseDataset):
             image_paths.append(image_path)
         return image_paths
 
+    def _generate_random_video_paths(self) -> list[str]:
+        video_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_video.mp4")
+        try:
+            import cv2
+
+            profile = self.random_request_config[0] if self.random_request_config else {}
+            width = int(self.args.width or profile.get("width") or 1280)
+            height = int(self.args.height or profile.get("height") or 720)
+            num_frames = int(self.args.num_frames or profile.get("num_frames") or 16)
+            fps = float(self.args.fps or profile.get("fps") or 8)
+            writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
+            if not writer.isOpened():
+                raise RuntimeError("cv2.VideoWriter failed to open")
+            for frame_idx in range(num_frames):
+                frame = np.zeros((height, width, 3), dtype=np.uint8)
+                x = int((frame_idx / max(num_frames - 1, 1)) * max(width - 96, 1))
+                y = height // 2
+                frame[:, :, 0] = 24
+                frame[:, :, 1] = 32
+                frame[:, :, 2] = 48
+                frame[max(y - 36, 0) : min(y + 36, height), x : min(x + 96, width), :] = (64, 128, 220)
+                writer.write(frame)
+            writer.release()
+            return [video_path]
+        except Exception as e:
+            raise RuntimeError(
+                "Failed to generate synthetic v2v input video. Install opencv-python or provide a custom dataset."
+            ) from e
+
 
 def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None:
     """Compute expected execution time (ms) based on a base per-step-per-frame unit time.
@@ -1112,7 +1208,7 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
 
 
 def _default_endpoint_for_task(task: str) -> str:
-    if task in {"t2v", "i2v", "ti2v"}:
+    if task in {"t2v", "i2v", "ti2v", "v2v"}:
         return "/v1/videos"
     if task in {"i2i", "ti2i", "it2i"}:
         return "/v1/images/edits"
@@ -1126,7 +1222,7 @@ async def benchmark(args):
     if args.base_url is None:
         args.base_url = f"http://{args.host}:{args.port}"
 
-    VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
+    VIDEO_TASKS = {"t2v", "i2v", "ti2v", "v2v"}
     IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"}
 
     if args.task in VIDEO_TASKS:
@@ -1177,6 +1273,10 @@ async def benchmark(args):
     if args.return_stage_metrics and args.endpoint in _STAGE_METRICS_ENDPOINTS:
         for req in requests_list:
             req.extra_body.setdefault(_RETURN_STAGE_METRICS_FIELD, True)
+    if args.extra_body:
+        extra_body = json.loads(args.extra_body)
+        for req in requests_list:
+            req.extra_body.update(extra_body)
 
     if args.endpoint == "/v1/images/edits":
         for req in requests_list:
@@ -1328,7 +1428,7 @@ if __name__ == "__main__":
         "--task",
         type=str,
         default="t2v",
-        choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
+        choices=["t2v", "i2v", "ti2v", "v2v", "ti2i", "i2i", "it2i", "t2i"],
         help="Task type.",
     )
     parser.add_argument(
@@ -1443,6 +1543,12 @@ if __name__ == "__main__":
             '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
         ),
     )
+    parser.add_argument(
+        "--extra-body",
+        type=str,
+        default=None,
+        help="JSON object merged into every generated request body/form.",
+    )
     parser.add_argument(
         "--num-input-images",
         type=int,
diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci
index 8b19ed678..f4444f8ac 100644
--- a/docker/Dockerfile.ci
+++ b/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}
diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda
index 55d5859a8..4e9075f2a 100644
--- a/docker/Dockerfile.cuda
+++ b/docker/Dockerfile.cuda
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.24.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.25.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app
diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu
index 95774dd25..642b89882 100644
--- a/docker/Dockerfile.npu
+++ b/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.25.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm
diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3
index ea8882a38..4fe1d56fe 100644
--- a/docker/Dockerfile.npu.a3
+++ b/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.25.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm
diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci
index 2c027992e..a3cede3a7 100644
--- a/docker/Dockerfile.npu.ci
+++ b/docker/Dockerfile.npu.ci
@@ -1,11 +1,13 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
+ARG VLLM_ASCEND_TAG=v0.25.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 COPY . .
 
+RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list
+
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         ffmpeg \
diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3
index 4e0de0a1d..1369d8c85 100644
--- a/docker/Dockerfile.npu.ci.a3
+++ b/docker/Dockerfile.npu.ci.a3
@@ -1,11 +1,13 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
+ARG VLLM_ASCEND_TAG=v0.25.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 COPY . .
 
+RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list
+
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         ffmpeg \
diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
index 7c45cb770..49c6f02aa 100644
--- a/docker/Dockerfile.rocm
+++ b/docker/Dockerfile.rocm
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.24.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.25.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu
index 52069f344..b70dd274e 100644
--- a/docker/Dockerfile.xpu
+++ b/docker/Dockerfile.xpu
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.24.0
+ARG VLLM_VERSION=v0.25.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 
diff --git a/docs/.nav.yml b/docs/.nav.yml
index 2e7097efb..6bf161b1b 100644
--- a/docs/.nav.yml
+++ b/docs/.nav.yml
@@ -23,7 +23,7 @@ nav:
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
       - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
       - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
-      - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
+      - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
@@ -38,7 +38,7 @@ nav:
       - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
-      - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
+      - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
       - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
@@ -80,6 +80,7 @@ nav:
       - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
       - FP8 W8A8: user_guide/quantization/fp8.md
       - Int8 W8A8: user_guide/quantization/int8.md
+      - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md
       - MXFP8 W8A8: user_guide/quantization/mxfp8.md
       - MXFP4 W4A4: user_guide/quantization/mxfp4.md
       - ModelOpt: user_guide/quantization/modelopt.md
diff --git a/docs/README.md b/docs/README.md
index 66fc8ef46..f6263483b 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -28,9 +28,9 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 [vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving:
 
-- **Omni-modality**: Text, image, video, and audio data processing
+- **Omni-modality**: Text, image, audio, video, and action data processing
 - **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models
-- **Heterogeneous outputs**: from traditional text generation to multimodal outputs
+- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs
 
 <p align="center">
   <picture>
@@ -55,8 +55,10 @@ vLLM-Omni is flexible and easy to use with:
 
 vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
 
-- Omni-modality models (e.g. Qwen2.5-Omni, Qwen3-Omni)
-- Multi-modality generation models (e.g. Qwen-Image)
+- Omni-modality models (e.g. Qwen3-Omni, Cosmos3, HunyuanImage, BAGEL)
+- TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
+- Diffusion models — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
+- Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
 
 For more information, checkout the following:
 
diff --git a/docs/api/README.md b/docs/api/README.md
index d4c580629..cfea6d175 100644
--- a/docs/api/README.md
+++ b/docs/api/README.md
@@ -46,17 +46,17 @@ Engine classes for offline and online inference.
 - [vllm_omni.engine.PromptEmbedsPayload][]
 - [vllm_omni.engine.arg_utils.OmniEngineArgs][]
 - [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][]
-- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][]
-- [vllm_omni.engine.mm_outputs.MultimodalPayload][]
 - [vllm_omni.engine.orchestrator.Orchestrator][]
 - [vllm_omni.engine.orchestrator.OrchestratorRequestState][]
-- [vllm_omni.engine.output_modality.OutputModality][]
-- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][]
-- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
-- [vllm_omni.engine.output_processor.OmniRequestState][]
 - [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][]
 - [vllm_omni.engine.stage_init_utils.StageMetadata][]
 - [vllm_omni.engine.stage_init_utils.StartedLlmStage][]
+- [vllm_omni.outputs.mm_outputs.MultimodalCompletionOutput][]
+- [vllm_omni.outputs.mm_outputs.MultimodalPayload][]
+- [vllm_omni.outputs.output_modality.OutputModality][]
+- [vllm_omni.outputs.output_modality.TensorAccumulationStrategy][]
+- [vllm_omni.outputs.output_processor.MultimodalOutputProcessor][]
+- [vllm_omni.outputs.output_processor.OmniRequestState][]
 
 ## Core
 
diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg
index 509eb9ef5..cb7300de2 100644
Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ
diff --git a/docs/cli/README.md b/docs/cli/README.md
index 1fcfdb14e..9e3240de4 100644
--- a/docs/cli/README.md
+++ b/docs/cli/README.md
@@ -18,11 +18,17 @@ Specify the port:
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
 ```
 
-If you have custom stage configs file, launch the server with command below
+For a migrated model, load a custom deploy configuration with `--deploy-config`:
+
 ```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni --stage-configs-path /path/to/stage_configs_file
+vllm serve Qwen/Qwen2.5-Omni-7B --omni --deploy-config /path/to/deploy_config.yaml
 ```
 
+The deprecated `--stage-configs-path` flag is retained for models that still use the legacy `stage_args` schema:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --stage-configs-path /path/to/legacy_stage_config.yaml
+```
 
 ## bench
 
diff --git a/docs/cli/bench/serve.md b/docs/cli/bench/serve.md
index d6b3ef393..2a1acd1f0 100644
--- a/docs/cli/bench/serve.md
+++ b/docs/cli/bench/serve.md
@@ -2,7 +2,7 @@
 The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
 
 ## Notes
-We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend.
+vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` serving benchmark backends.
 
 ## Basic Parameter Description
 You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
@@ -10,7 +10,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
   Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
 
 - `--backend`
-  Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
+  Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` to the upstream vLLM backend choices.
 
 - `--model`
   The model identifier to load, filled according to the models supported by vLLM-Omni.
@@ -55,7 +55,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
                     'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", '
                     '"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
 
-- `-print-stage`
+- `--print-stage`
 Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
 
 - `--save-result`
@@ -312,7 +312,7 @@ vllm bench serve \
   --ignore-eos \
   --percentile-metrics ttft,tpot,itl \
   --random-output-len 2 \
-  --extra_body '{"modalities": ["text"]}'
+  --extra-body '{"modalities": ["text"]}'
 ```
 
 If successful, you will see the following output:
@@ -390,7 +390,7 @@ vllm bench serve --omni \
   --ignore-eos \
   --percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \
   --random-output-len 900 \
-  --extra_body '{"modalities": ["text","audio"]}' \
+  --extra-body '{"modalities": ["text","audio"]}' \
   --print-stage
 ```
 
diff --git a/docs/configuration/composable_parallel.md b/docs/configuration/composable_parallel.md
new file mode 100644
index 000000000..a456de843
--- /dev/null
+++ b/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** (`--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` value passed through `--stage-overrides` or a `--tensor-parallel-size` value that conflicts with the strategy's world size cannot slip past the pre-spawn check.
+
+## Worked example
+
+### 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 (`--tensor-parallel-size 2`, `--stage-overrides '{"0": {"tensor_parallel_size": 2}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
+
+!!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`"
+    Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits:
+
+    ```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`.
diff --git a/docs/configuration/pd_disaggregation.md b/docs/configuration/pd_disaggregation.md
index 9196bdb02..7deb0443a 100644
--- a/docs/configuration/pd_disaggregation.md
+++ b/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:
diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md
index 803c74e0d..1f14d21fd 100644
--- a/docs/configuration/stage_configs.md
+++ b/docs/configuration/stage_configs.md
@@ -1,11 +1,11 @@
-# Stage configs for vLLM-Omni
+# Deploy and legacy stage configurations
 
-In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion.
+In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema.
 
 !!! note
     Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. The bundled defaults have been verified on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni. Models that have not yet migrated to the new schema continue to use the legacy `vllm_omni/model_executor/stage_configs/<model>.yaml` files via `--stage-configs-path`.
 
-## New deploy schema reference
+## Deploy configuration schema
 
 The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields:
 
@@ -82,14 +82,14 @@ stages:
     input_connectors:  {from_stage_0: shm}
 ```
 
-### CLI flags introduced in this refactor
+### CLI flags
 
 | Flag | Description |
 |------|-------------|
-| `--deploy-config PATH` | Load a new-schema deploy YAML. Takes precedence over `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
+| `--deploy-config PATH` | Load a new-schema deploy YAML. It is mutually exclusive with `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
 | `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
 | `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
-| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. |
+| `--stage-configs-path` | **Deprecated.** Load a legacy `stage_args` YAML. New-schema YAMLs passed to this flag are auto-detected for compatibility, but should use `--deploy-config`. It is mutually exclusive with `--deploy-config` and will be removed in a future release. |
 
 ### Stage-Based CLI Paradigm
 
@@ -146,7 +146,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 
 From highest to lowest:
 
-1. Per-stage flags (`--stage-overrides` JSON, `--stage-<id>-<key>` if registered)
+1. Per-stage overrides (`--stage-overrides` JSON)
 2. Explicit global CLI flags (`--gpu-memory-utilization 0.85`, etc.)
 3. Platform section (`platforms.npu.stages`, etc.) on top of the base `stages:`
 4. Overlay YAML (via `base_config:`) on top of the base YAML
@@ -210,7 +210,7 @@ Effective config per stage after the merge:
 | 1 | `max_model_len` | `16384` | global CLI |
 | 2 | (all defaults) | — | base YAML (no overrides apply) |
 
-Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions:
+Therefore, as a core part of vLLM-Omni, a model's pipeline and deployment configurations have several main functions:
 
 - Claim partition of stages and their corresponding class implementation in `model_executor/models`.
 - The disaggregated configuration for each stage and the communication topology among them.
@@ -219,16 +219,16 @@ Therefore, as a core part of vLLM-Omni, the stage configs for a model have sever
 - Default input parameters.
 
 To override specific parameters, explicitly inject the customized configuration schema
-in both online and offline instantiation flows. Prioritize the `--deploy-config` flag
-when loading the new-schema deploy YAML schemas, reserving the `--stage-configs-path` parameter
+in both online and offline instantiation flows. Use the `--deploy-config` flag
+when loading a deploy configuration, reserving the `--stage-configs-path` parameter
 exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
 
 Examples:
 
-For offline (Assume necessary dependencies have been imported):
+For offline inference (assuming the necessary dependencies have been imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml")
 ```
 
 For online serving:
@@ -242,11 +242,14 @@ Legacy online serving:
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 !!! important
-    We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
+    We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers.
+
+## Legacy stage configuration reference
+
+The remainder of this page documents the deprecated `stage_args` schema for models that have not migrated to deploy configurations. The following historical Qwen2.5-Omni configuration illustrates that schema; current Qwen2.5-Omni deployments should use `vllm_omni/deploy/qwen2_5_omni.yaml` instead.
 
-Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
 ```python
-# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
+# Historical stage configuration for Qwen2.5-Omni.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration
@@ -338,7 +341,7 @@ runtime:
 
 ```
 
-## Stage Configuration Arguments
+## Legacy stage configuration arguments
 
 Each stage in the `stage_args` list contains the following configuration options:
 
diff --git a/docs/contributing/DOCS_GUIDE.md b/docs/contributing/DOCS_GUIDE.md
index 100bac674..916f5e60f 100644
--- a/docs/contributing/DOCS_GUIDE.md
+++ b/docs/contributing/DOCS_GUIDE.md
@@ -67,6 +67,21 @@ docs/
 └── stylesheets/         # Custom CSS
 ```
 
+## Naming Model Examples
+
+Offline inference and online serving examples for the same model use the same
+directory name and the shared display name in
+`examples/model_display_names.yml`. Use these title forms:
+
+- `# <Model>: Offline inference`
+- `# <Model>: Online serving`
+
+The documentation generator uses the full title for the page H1 and the shared
+display name alone for navigation, and fails the build if a mapped README uses
+a different H1. Keep checkpoint identifiers in commands and prose rather than
+in the display name, and do not rename an existing example directory solely to
+adjust its title because the directory defines its public documentation URL.
+
 ## Publishing Documentation
 
 ### GitHub Pages (Recommended)
diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md
index 17b4a3792..c9d86e1c4 100644
--- a/docs/contributing/ci/CI_5levels.md
+++ b/docs/contributing/ci/CI_5levels.md
@@ -126,11 +126,12 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
         <strong>Interface tests:</strong><br>
         /tests/entrypoints/test_xxx<br>
         <strong>Performance:</strong><br>
-        /tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),<br>
-        and /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)<br>
+        /tests/dfx/perf/tests/test_qwen3_omni_*.json (Omni), test_tts.json (TTS),<br>
+        test_voxcpm2.json, test_higgs_audio_v3.json, and<br>
+        /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)<br>
         <strong>Doc Test:</strong><br>
-        tests/example/online_serving/test_{model_name}.py<br>
-        tests/example/offline_inference/test_{model_name}.py<br>
+        tests/examples/online_serving/test_{model_name}.py<br>
+        tests/examples/offline_inference/test_{model_name}.py<br>
         <strong>Accuracy Test:</strong><br>
         /tests/e2e/accuracy/test_{model_name}.py
       </td>
@@ -708,16 +709,16 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 ### 3.2 Testing Content and Scope
 
 -   ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
--   ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs in the form `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` via `--test-config-file`) to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
+-   ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen3_omni_*.json` (Omni), `test_tts.json` / `test_voxcpm2.json` / `test_higgs_audio_v3.json` (TTS), and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` or `run_diffusion_benchmark.py` via `--test-config-file` in nightly **Perf Test** steps) to drive throughput, latency, and memory benchmarks. Each JSON **case** may declare an optional top-level **`mark`** array: exactly one ``hardware_marks`` object plus pytest marker name strings (`full_model`, `omni` / `tts` / `diffusion`, …). Runners attach those marks to each parametrized `(server, benchmark index)` pair so **local** bulk runs can filter with `-m` (for example `-m "full_model and H100 and diffusion"`). Nightly CI perf jobs select workloads by **`--test-config-file`**, not `-m`. Details are in the Performance Tests example below.
 -   ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
 
 ### 3.3 Test Directory and Execution Files
 
 -   ***Functional Testing***: Same directories as L3.
--   ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`)
+-   ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen3_omni_*.json`, `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`). Optional per-case `mark` for local `-m` filtering is documented in Section 3.4 Performance Tests.
 -   ***Documentation Example Tests***:
--   -   `tests/example/online_serving/test_{model_name}.py`
-    -   `tests/example/offline_inference/test_{model_name}.py`
+    -   `tests/examples/online_serving/test_{model_name}.py`
+    -   `tests/examples/offline_inference/test_{model_name}.py`
 
 ### 3.4 Execution Method and Example
 
diff --git a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md
index f1f3073dc..5c59ba499 100644
--- a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md
+++ b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md
@@ -1,14 +1,41 @@
-When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, or diffusion configs such as `tests/dfx/perf/tests/test_*_vllm_omni.json` (selected via `pytest ... run_benchmark.py --test-config-file <path>`):
+When you want to add L4-level ***performance test*** cases, add entries to JSON files under `tests/dfx/perf/tests/` and run them via `tests/dfx/perf/scripts/run_benchmark.py` (omni / TTS) or `run_diffusion_benchmark.py` (diffusion).
+
+**Config file layout (in-tree examples)**
+
+| Model type | Runner | Example JSON files |
+| ---------- | ------ | ------------------ |
+| Omni | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` |
+| TTS | `run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` |
+| Diffusion | `run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`, `test_wan22_i2v_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, … |
+
+**How runners pick cases**
+
+Without **`--test-config-file`**, each runner scans all `*.json` under `tests/dfx/perf/tests/` but only keeps its own model type:
+
+- **`run_benchmark.py`**: omni and TTS cases only (skips diffusion JSON).
+- **`run_diffusion_benchmark.py`**: diffusion cases only (skips omni / TTS JSON).
+
+Diffusion cases are detected when the JSON has `server_type` (typically `"vllm-omni"`) or `"diffusion"` in the `mark` array. Omni / TTS JSON has neither.
+
+**Running perf cases**
+
+Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bulk scan above and filter with pytest **`-m`** on each case's JSON `mark` (for example `pytest … run_benchmark.py -m "full_model and tts and H100"`).
 
 ```JSON
 {
-    "test_name": "test_qwen3_omni",
+    "test_name": "test_qwen3_omni_async_chunk",
+    "mark": [
+        {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
+        "full_model",
+        "omni"
+    ],
     "server_params": {
         "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
         "stage_config_name": "qwen3_omni.yaml"
     },
     "benchmark_params": [
         {
+            "name": "random_10p",
             "dataset_name": "random",
             "num_prompts": [10, 20],
             "max_concurrency": [1, 4],
@@ -33,8 +60,95 @@ When you want to add L4-level ***performance test*** cases, you can refer to the
 | Field            | Required | Description                                                     |
 | ---------------- | -------- | --------------------------------------------------------------- |
 | test_name        | Yes      | Unique identifier for the test case                             |
+| mark             | No       | Pytest marks for this case (see **`mark` field** below). Omit only for configs not meant to be filtered by `-m`. |
 | server_params    | Yes      | Server-side configuration parameters                            |
 | benchmark_params | Yes      | Benchmark running parameters (supports multiple configurations) |
+| server_type      | Diffusion only | When set (e.g. `"vllm-omni"`), routes the case to `run_diffusion_benchmark.py`. |
+| benchmark_endpoint | Diffusion only | API path for the benchmark client (e.g. `/v1/videos`, `/v1/images/generations`). |
+
+**`mark` field**
+
+Optional top-level field on each perf JSON **case object** (one per `test_name`). `run_benchmark.py` and `run_diffusion_benchmark.py` read it via `tests.dfx.conftest.resolve_pytest_marks` and attach the marks to the corresponding `pytest.param`, so you can filter locally with `-m`.
+
+When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, `diffusion`, or `local_model`.
+
+Supported form:
+
+| Form | Example | Effect |
+| ---- | ------- | ------ |
+| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` > 1 adds `distributed_*` (+ CUDA `skipif_cuda` when applicable). String entries become extra pytest marks. |
+| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325", "npu": "A2"}, "num_cards": {"cuda": 2, "rocm": 2, "npu": 1}}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** in one case. Filter examples: `-m "full_model and H100 and cuda"`, `-m "full_model and MI325 and rocm"`. |
+
+Recommended for L4 perf cases:
+
+```JSON
+{
+    "test_name": "test_bagel_single_device_single_stage_t2i",
+    "mark": [
+        {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+        "full_model",
+        "diffusion"
+    ],
+    "server_type": "vllm-omni",
+    "server_params": { "...": "..." },
+    "benchmark_params": [ { "name": "1024x1024_steps20", "...": "..." } ]
+}
+```
+
+Multi-GPU diffusion (example: Cosmos3 with `cfg-parallel-size=2`):
+
+```JSON
+{
+    "test_name": "test_cosmos3_t2i_official_demo_2gpu",
+    "mark": [
+        {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
+        "full_model",
+        "diffusion"
+    ],
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/images/generations",
+    "server_params": { "...": "..." },
+    "benchmark_params": [ { "name": "1024x1024_steps4", "...": "..." } ]
+}
+```
+
+- Omni perf: include `"omni"` in the `mark` array
+- TTS perf: include `"tts"` in the `mark` array
+- Diffusion perf: include `"diffusion"` in the `mark` array
+- HunyuanImage local-weight cases: add `"local_model"` to the `mark` array
+
+**Parametrization IDs**
+
+Each `(server, benchmark index)` pair becomes one `pytest.param` with id `{test_name}-{suffix}`. The suffix comes from `benchmark_params[].name` when set (for example `test_tts-p0`, `test_omni-p1`); otherwise it is derived from fields like `task` / `eval_phase`.
+
+**Benchmark result filenames**
+
+Result files use the **runtime** hardware label from `get_runtime_resource_label()` (detected GPU/NPU name on the machine that ran the job), **not** `mark.hardware_marks.res`. On the default H100 CI pool, `H100` is **omitted** from filenames (`resource_label_for_filename`).
+
+Examples:
+
+- Omni/TTS: `result_{test_name}_{optional_hw}_{dataset}_....json` under `BENCHMARK_DIR`
+- Diffusion: one aggregate `diffusion_result_{config_stem}_{optional_hw}_{timestamp}.json` per source JSON file (array of all runs from that file)
+
+**Local commands**
+
+```bash
+# Bulk load + filter by JSON mark
+pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and H100 and diffusion"
+pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
+
+# Single file (same as nightly CI Perf steps)
+pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
+  --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
+pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
+  --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
+
+# Optional baseline assertion (default off)
+pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --assert-baseline \
+  --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+```
+
+See also [Markers for Tests](./tests_markers.md) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
 
 **`server_params` Configuration**
 
@@ -79,10 +193,12 @@ You can add any benchmark running parameters you need here. For all optional par
 1.  Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
 2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
 3.  Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
-4.  The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
+4.  Set `"name"` on each `benchmark_params` entry for stable pytest ids and readable result keys.
+5.  The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
 
 | Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
 | --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| name            | string      | No       | `"1024x1024_steps4"` | Stable suffix for pytest param id and logs. |
 | num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
 | request_rate    | float / array | No  | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
 | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |
diff --git a/docs/contributing/ci/test_guide.md b/docs/contributing/ci/test_guide.md
index 0db814e08..ad6a63963 100644
--- a/docs/contributing/ci/test_guide.md
+++ b/docs/contributing/ci/test_guide.md
@@ -12,12 +12,37 @@ When running test cases, you may need to install the following dependencies:
 
 ```bash
 uv pip install ".[dev]"
-apt-get install -y ffmpeg
+apt-get install -y espeak-ng jq
 ```
 
 ## Running Tests
 Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
 
+### CI job runners (L2–L4): logs, timing, and timeouts
+
+[`tools/run_ready_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_ready_jobs.sh), [`tools/run_merge_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_merge_jobs.sh), and [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh) share the same run-time behavior (via [`tools/run_jobs_common.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_jobs_common.sh)):
+
+**Log layout** (default under `logs/ready_jobs`, `logs/merge_jobs`, or `logs/nightly_jobs`; override with `--log-dir`):
+
+| Path | Description |
+|------|-------------|
+| `jobs/*.sh` | Generated per-step bash wrappers |
+| `jobs/.job_timeouts` | Job key → `timeout_in_minutes` from Buildkite YAML (when set) |
+| `<job_key>.log` | Tee output for each job |
+| `timing_summary.log` | Per-job duration and total wall time after the run completes |
+
+**Timing:** while each job runs, stderr prints `finished in …` or `failed after …`. When all jobs finish, a summary is printed and written to `timing_summary.log`, for example:
+
+```
+=== Job timing summary ===
+  TTS_VoxCPM2_Test  25m 40s  OK
+  Omni_Qwen3-Omni_Test  8m 05s  TIMED OUT
+Total wall time: 33m 45s (2 jobs)
+Failed jobs: 1/2
+```
+
+**Timeouts:** if a Buildkite step defines `timeout_in_minutes`, the runner wraps that job with `timeout ${N}m` (aligned with CI). Steps without `timeout_in_minutes` rely on any inline `timeout …` already present in the extracted pytest command. On timeout the process exits with code `124` and the summary marks the job as `TIMED OUT`.
+
 === "L1 level"
 
     ```bash
@@ -28,60 +53,114 @@ Our test scripts use the pytest framework. First, please use `git clone https://
 
 === "L2 level"
 
+    **Recommended:** run CI-aligned jobs from the repo root via [`tools/run_ready_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_ready_jobs.sh). The script reads [`.buildkite/test-ready.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml), generates per-step bash wrappers, runs pytest, and tees logs under `logs/ready_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts) for timing and timeout behavior). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
+    ```bash
+    # All L2 and L1 jobs (default: every pytest step in test-ready.yml)
+    bash tools/run_ready_jobs.sh
+
+    # Skip L1-style "Simple Test" steps in test-ready.yml (group or "Simple ·" labels)
+    bash tools/run_ready_jobs.sh --skip-simple
+    bash tools/run_ready_jobs.sh --skip-simple --model-type diffusion
+
+    # Preview extracted commands without running
+    bash tools/run_ready_jobs.sh --dry-run
+
+    # Filter by model area (OR semantics): omni | tts | diffusion | all
+    bash tools/run_ready_jobs.sh --model-type omni
+    bash tools/run_ready_jobs.sh --model-type tts,diffusion
+
+    # Match a Buildkite label substring (e.g. one model suite)
+    bash tools/run_ready_jobs.sh --model-type tts --label-substr "VoxCPM2"
+
+    ```
+
+    **Ad hoc:** run a single test file or marker expression manually:
+
     ```bash
     cd tests
-    pytest -s -v -m "core_model and not cpu" --run-level=core_model
-    ```
-    If you only want to run a specific test case, you can use:
-    ```bash
     pytest -s -v test_xxxx.py --run-level=core_model
+    pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
     ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
-    ```
-    The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
 
 === "L3 level"
 
+    **Recommended:** run CI-aligned jobs from the repo root via [`tools/run_merge_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_merge_jobs.sh). The script reads [`.buildkite/test-merge.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml), generates per-step bash wrappers, runs pytest, and tees logs under `logs/merge_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts)). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
     ```bash
-    pytest -s -v -m "advanced_model" --run-level=advanced_model
+    # All L3 and L1 jobs (default: every pytest step in test-merge.yml)
+    bash tools/run_merge_jobs.sh
+
+    # Skip L1-style "Simple Test" steps in test-merge.yml (group or "Simple ·" labels)
+    bash tools/run_merge_jobs.sh --skip-simple
+    bash tools/run_merge_jobs.sh --skip-simple --model-type omni
+
+    # Preview extracted commands without running
+    bash tools/run_merge_jobs.sh --dry-run
+
+    # Filter by model area (OR semantics): omni | tts | diffusion | all
+    bash tools/run_merge_jobs.sh --model-type diffusion
+    bash tools/run_merge_jobs.sh --model-type omni,tts
+
+    # Match a Buildkite label substring
+    bash tools/run_merge_jobs.sh --model-type diffusion --label-substr "Wan22"
+
     ```
-    If you only want to run a specific test case, you can use:
+
+    **Ad hoc:** run a single test file or marker expression manually:
+
     ```bash
     pytest -s -v test_xxxx.py --run-level=advanced_model
+    pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
     ```
-    If you only want to run specific test cases on a particular platform, you can use:
-    ```bash
-    pytest -s -v -m "advanced_model and distributed_cuda and L4"  --run-level=advanced_model
-    ```
-    The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
 
 
 === "L4 level"
 
+    **Recommended:** run CI-aligned nightly jobs from the repo root via [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh). The script reads [`.buildkite/test-nightly.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml), generates per-step bash wrappers, runs pytest (perf jobs first, then aggregates perf JSON into Excel), and tees logs under `logs/nightly_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts); `generate_nightly_perf_excel` is included in the timing summary). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`).
+
+    ```bash
+    # All nightly jobs (default: test-type all, model-type all)
+    bash tools/nightly/run_nightly_jobs.sh
+
+    # Preview extracted commands without running
+    bash tools/nightly/run_nightly_jobs.sh --dry-run
+
+    # Test kind (--test-type, OR semantics): perf | acc | function | stability | local | all
+    bash tools/nightly/run_nightly_jobs.sh --test-type function
+    bash tools/nightly/run_nightly_jobs.sh --test-type perf,acc
+
+    # Model area (--model-type, OR semantics): omni | tts | diffusion | all
+    bash tools/nightly/run_nightly_jobs.sh --test-type function --model-type omni
+    bash tools/nightly/run_nightly_jobs.sh --test-type perf --model-type diffusion
+
+    # Stability scripts (tests/dfx/stability/scripts/) and local marker tests
+    bash tools/nightly/run_nightly_jobs.sh --test-type stability --model-type omni
+    bash tools/nightly/run_nightly_jobs.sh --test-type local --model-type tts
+    ```
+
+    Perf steps (label contains ``Perf Test``) run before other jobs; Excel output is written under ``logs/``. Diffusion nightly steps may also appear in [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml) on CI.
+
+    **Ad hoc:** run a single test file or marker expression manually:
+
     ```bash
     cd tests
-    pytest -s -v -m "full_model" --run-level=full_model
-    ```
-    If you only want to run a specific test case, you can use:
-    ```bash
     pytest -s -v test_xxxx.py --run-level=full_model
+    pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
+    pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
     ```
     If you only want to run specific test cases on a particular platform, you can use:
     ```bash
     pytest -s -v -m "full_model and distributed_cuda and L4"  --run-level=full_model
     ```
-    Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
+    Note: ``run_benchmark.py`` and ``run_diffusion_benchmark.py`` accept an optional ``--test-config-file``. If omitted, each loads every ``*.json`` under ``tests/dfx/perf/tests/`` (omni/tts vs diffusion split by ``is_diffusion_perf_config``) and pytest ``-m`` filters by each case's JSON ``mark``:
     ```bash
-    pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
-    ```
-    The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example:
-
-    ```bash
-    cd tests
-    pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
+    pytest -sv tests/dfx/perf/scripts/run_benchmark.py -m "full_model and tts and H100"
+    pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"
+    pytest -sv tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json
+    pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
     ```
+    Nightly **Perf Test** jobs in ``test-nightly.yml`` use ``--test-config-file`` only (no ``-m``). E2e L4 function tests still use ``full_model`` + ``--run-level full_model`` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml)). Example:
 
 === "L5 level"
 
diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md
index cc89cee47..0a0a974e6 100644
--- a/docs/contributing/ci/tests_markers.md
+++ b/docs/contributing/ci/tests_markers.md
@@ -33,7 +33,8 @@ Defined in `pyproject.toml`:
 | `skipif_rocm`      | Skip if the num of ROCm cards is less than the required * |
 | `skipif_npu`       | Skip if the num of NPU cards is less than the required *  |
 | `slow`             | Slow tests (may skip in quick CI)                         |
-| `benchmark`        | Benchmark tests                                           |
+| `benchmark`        | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) |
+| `local_model`      | Tests requiring local / non-HF-hub model weights          |
 
 \* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
 
@@ -94,7 +95,7 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
         num_cards=2,
     )
     ```
-- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), NPU (A2/A3)
+- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), XPU (B60), MUSA (S5000), NPU (A2/A3)
 - `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
 - Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
 - Filtering examples:
@@ -120,6 +121,32 @@ def test_feature(omni_server):
     ...
 ```
 
+### JSON `mark` field (L4 perf configs)
+
+Perf JSON under `tests/dfx/perf/tests/` can attach pytest marks per **case** (`test_name` block). Parsed by `tests.dfx.conftest.resolve_pytest_marks` and applied to each parametrized run in `run_benchmark.py` / `run_diffusion_benchmark.py`.
+
+When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same semantics as `hardware_marks()` above), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, or `diffusion`.
+
+```json
+{
+  "test_name": "test_cosmos3_t2i_official_demo_2gpu",
+  "mark": [
+    {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
+    "full_model",
+    "diffusion"
+  ],
+  "server_type": "vllm-omni",
+  "server_params": { "...": "..." },
+  "benchmark_params": [{ "name": "1024x1024_steps4", "...": "..." }]
+}
+```
+
+- Local bulk load: `pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"`
+- Nightly CI perf steps: `--test-config-file tests/dfx/perf/tests/test_<model>_vllm_omni.json` (file selects cases; no `-m`)
+- Result filenames use **runtime** GPU detection (`get_runtime_resource_label`); `H100` is omitted on the default CI pool
+
+See [L4 performance test examples](./test_examples/l4_performance_tests.inc.md) for the full schema.
+
 ## Add Support for a New Platform
 
 If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md
index 7ff13fd1b..b38805112 100644
--- a/docs/contributing/model/adding_diffusion_model.md
+++ b/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
@@ -638,7 +638,7 @@ mkdir -p examples/online_serving/your_model_name
 
 - Script: `examples/offline_inference/your_model_name/end2end.py`
   - Parse args like BAGEL (`--model`, `--modality`, optional `--image-path`, `--steps`, etc.)
-  - Use `from vllm_omni.entrypoints.omni import Omni` (or `OmniDiffusion` if your model is diffusion-only)
+  - Use `from vllm_omni.entrypoints.omni import Omni` for both multi-stage and diffusion-only models
   - Save outputs (images/audio/video/text) with deterministic filenames (e.g., `output_0_0.png`)
 - Doc: `examples/offline_inference/your_model_name/README.md`
   - Include at least one runnable command, e.g.:
diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md
index 1eaff1059..c20e94eda 100644
--- a/docs/contributing/model/adding_omni_model.md
+++ b/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).
diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md
index 1c38ba671..20c99fd9c 100644
--- a/docs/design/architecture_overview.md
+++ b/docs/design/architecture_overview.md
@@ -13,7 +13,7 @@ This document outlines the architectural design for vLLM-Omni.
 
 The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support:
 
-* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text.
+* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, video, and action trajectories, alongside text.
 * **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation.
 * **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable.
 * **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats.
diff --git a/docs/design/feature/cfg_parallel.md b/docs/design/feature/cfg_parallel.md
index c73a87749..053360494 100644
--- a/docs/design/feature/cfg_parallel.md
+++ b/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(
diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md
index 953e11f7f..02d50181b 100644
--- a/docs/design/module/dit_module.md
+++ b/docs/design/module/dit_module.md
@@ -470,7 +470,7 @@ class Attention(nn.Module):
                  role="self", role_category=None, ...):
         # Resolve backend for this role from the active OmniDiffusionConfig
         config = get_current_diffusion_config_or_none()
-        attention_config = config.attention_config if config is not None else None
+        attention_config = config.diffusion_attention_config if config is not None else None
         attn_backend_cls, spec = get_attn_backend_for_role(
             role=role,
             head_size=head_size,
@@ -493,7 +493,7 @@ These backends provide the **kernel implementations** for attention computation.
 
 #### Backend Selection Mechanism
 
-Selection is driven by `AttentionConfig` on `OmniDiffusionConfig`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
+Selection is driven by `OmniDiffusionConfig.diffusion_attention_config`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order:
 
 ```python
 def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None):
@@ -909,7 +909,7 @@ def initialize_model_parallel(
 
 ```
 1. User Request
-   └─> OmniDiffusion.generate(prompt)
+   └─> Omni.generate(prompt)
        └─> Prepare OmniDiffusionRequest
            └─> DiffusionEngine.step(requests)
 
diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md
index dc8b645e4..78d680f26 100644
--- a/docs/features/custom_pipeline.md
+++ b/docs/features/custom_pipeline.md
@@ -10,7 +10,7 @@ This guide demonstrates how to use the newly added features for extending vLLM-O
 Three main features enable custom pipeline extension:
 
 1. **`WorkerWrapperBase`**: A wrapper class that enables dynamic worker extension with custom functionality
-2. **`load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
+2. **`diffusion_load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
 3. **`CustomPipelineWorkerExtension`**: An extension class that enables pipeline re-initialization with custom implementations
 
 ## Features
@@ -27,15 +27,15 @@ Three main features enable custom pipeline extension:
 
 **Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
 
-### load_format Parameter
+### diffusion_load_format parameter
 
-The `load_format` parameter controls how diffusion models are loaded. It supports the following values:
+The `diffusion_load_format` parameter controls the initial diffusion model load. The relevant values are:
 
 - **`"default"`**: Standard model loading using the model registry (default behavior)
-- **`"custom_pipeline"`**: Load a custom pipeline class specified by `custom_pipeline_name`
-- **`"dummy"`**: Skip model loading (useful for testing or when pipeline will be initialized separately)
+- **`"dummy"`**: Skip the initial model load; use this with `custom_pipeline_args["pipeline_class"]` for the custom pipeline workflow below
+- **`"diffusers"`**: Load through the Hugging Face Diffusers adapter
 
-**Location:** `vllm_omni/diffusion/model_loader/diffusers_loader.py`
+**Location:** `vllm_omni/diffusion/data.py` (`OmniDiffusionConfig`)
 
 ### CustomPipelineWorkerExtension
 
@@ -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)
 
diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md
index 4ef35398c..e21a11cee 100644
--- a/docs/getting_started/installation/gpu.md
+++ b/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"
 
diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md
index 83cbcd68d..d4808139e 100644
--- a/docs/getting_started/installation/gpu/cuda.inc.md
+++ b/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
 ```
 
diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md
index af88e1e0c..4a64e8723 100644
--- a/docs/getting_started/installation/gpu/rocm.inc.md
+++ b/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]
diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md
index c5b13dd73..7381f5b10 100644
--- a/docs/getting_started/installation/npu/npu.inc.md
+++ b/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.25.0
+# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.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.25.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.25.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -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.25.0
+# Atlas A3: quay.io/ascend/vllm-omni:v0.25.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.25.0-a3 bash
 ```
 
 !!! tip
diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md
index 2c7f278ca..914f1f7b4 100644
--- a/docs/getting_started/quickstart.md
+++ b/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
 
diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py
index 81f2bac24..6cb601b96 100644
--- a/docs/mkdocs/hooks/generate_examples.py
+++ b/docs/mkdocs/hooks/generate_examples.py
@@ -16,8 +16,35 @@ ROOT_DIR_RELATIVE = "../../../../.."
 EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
 NAV_FILE = ROOT_DIR / "docs/.nav.yml"
+MODEL_DISPLAY_NAMES_FILE = EXAMPLE_DIR / "model_display_names.yml"
 MAX_INLINE_MATERIAL_SIZE = 8 * 1024
 
+SERVING_MODE_TITLES = {
+    "offline_inference": "Offline inference",
+    "online_serving": "Online serving",
+}
+
+
+def load_model_display_names() -> dict[str, str]:
+    try:
+        with open(MODEL_DISPLAY_NAMES_FILE, encoding="utf-8") as f:
+            model_display_names = yaml.safe_load(f)
+    except (OSError, yaml.YAMLError) as exc:
+        raise ValueError(f"Failed to load MODEL_DISPLAY_NAMES_FILE at {MODEL_DISPLAY_NAMES_FILE}: {exc}") from exc
+
+    if not isinstance(model_display_names, dict) or not all(
+        isinstance(key, str) and isinstance(value, str) for key, value in model_display_names.items()
+    ):
+        raise ValueError(
+            "MODEL_DISPLAY_NAMES_FILE "
+            f"at {MODEL_DISPLAY_NAMES_FILE} must contain a YAML mapping "
+            "of string model IDs to string display names"
+        )
+    return model_display_names
+
+
+MODEL_DISPLAY_NAMES = load_model_display_names()
+
 
 def fix_case(text: str) -> str:
     subs = {
@@ -54,6 +81,7 @@ class Example:
         main_file (Path): The main file in the directory.
         other_files (list[Path]): list of other files in the directory.
         title (str): The title of the document.
+        nav_title (str): The concise title used in navigation.
 
     Methods:
         __post_init__(): Initializes the main_file, other_files, and title attributes.
@@ -78,6 +106,16 @@ class Example:
     def is_code(self) -> bool:
         return self.main_file.suffix != ".md"
 
+    @property
+    def model_display_name(self) -> str | None:
+        if self.category not in SERVING_MODE_TITLES:
+            return None
+        return MODEL_DISPLAY_NAMES.get(self.path.stem)
+
+    @property
+    def nav_title(self) -> str:
+        return self.model_display_name or self.title
+
     def determine_main_file(self) -> Path:
         """
         Determines the main file in the given path.
@@ -151,13 +189,25 @@ class Example:
         return [file for file in self.path.rglob("*") if is_other_file(file)]
 
     def determine_title(self) -> str:
+        model_display_name = self.model_display_name
         if not self.is_code:
             # Specify encoding for building on Windows
             with open(self.main_file, encoding="utf-8") as f:
                 first_line = f.readline().strip()
             match = re.match(r"^#\s+(?P<title>.+)$", first_line)
+            if model_display_name:
+                expected_title = f"{model_display_name}: {SERVING_MODE_TITLES[self.category]}"
+                actual_title = match.group("title") if match else first_line
+                if actual_title != expected_title:
+                    raise ValueError(
+                        f"Model example title mismatch in {self.main_file}: "
+                        f"expected '# {expected_title}', got {first_line!r}"
+                    )
+                return expected_title
             if match:
                 return match.group("title")
+        elif model_display_name:
+            raise ValueError(f"Mapped model example must use a Markdown README: {self.path}")
         return fix_case(self.path.stem.replace("_", " ").title())
 
     def fix_relative_links(self, content: str) -> str:
@@ -300,7 +350,7 @@ def update_nav_file(examples: list[Example]):
         for example in category_examples:
             doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
             rel_path = doc_path.relative_to(ROOT_DIR / "docs")
-            category_items.append({example.title: rel_path.as_posix()})
+            category_items.append({example.nav_title: rel_path.as_posix()})
 
         if category_items:
             # Format category name (e.g., "offline_inference" -> "Offline Inference")
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 95f93f5c7..99b17d529 100644
--- a/docs/models/supported_models.md
+++ b/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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
diff --git a/docs/serving/audio_generate_api.md b/docs/serving/audio_generate_api.md
index afe693ae4..3109446c1 100644
--- a/docs/serving/audio_generate_api.md
+++ b/docs/serving/audio_generate_api.md
@@ -66,7 +66,7 @@ Content-Type: application/json
 |-----------|------|---------|-------------|
 | `input` | string | **required** | Text prompt describing the audio to generate |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
 
 #### Diffusion Parameters
@@ -90,7 +90,6 @@ Returns binary audio data with the appropriate `Content-Type` header:
 | `mp3` | `audio/mpeg` |
 | `flac` | `audio/flac` |
 | `pcm` | `audio/pcm` |
-| `aac` | `audio/aac` |
 | `opus` | `audio/opus` |
 
 ## Examples
@@ -285,7 +284,7 @@ Pydantic validation failure (e.g. invalid `response_format`, `speed` out of rang
     "detail": [
         {
             "type": "literal_error",
-            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3' or 'opus'",
             ...
         }
     ]
diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md
index f64d37f44..fef28142e 100644
--- a/docs/serving/image_edit_api.md
+++ b/docs/serving/image_edit_api.md
@@ -252,7 +252,7 @@ Validation errors (missing required fields):
 
 ```bash
 # Check if server is responding
-curl -X http://localhost:8000/v1/images/edit \
+curl http://localhost:8000/v1/images/edits \
   -F "prompt='test'"
 ```
 
diff --git a/docs/serving/image_generation_api.md b/docs/serving/image_generation_api.md
index 4702a19c1..12e2935ea 100644
--- a/docs/serving/image_generation_api.md
+++ b/docs/serving/image_generation_api.md
@@ -138,7 +138,7 @@ Content-Type: application/json
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
 | `n` | integer | 1 | Number of images to generate (1-10) |
 | `size` | string | model defaults | Image dimensions in WxH format (e.g., "1024x1024", "512x512") |
-| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
+| `response_format` | string | "b64_json" | Response format (`"b64_json"` or `"file"`) |
 | `user` | string | null | User identifier for tracking |
 
 #### vllm-omni Extension Parameters
diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md
index 27dfda95f..2bf630476 100644
--- a/docs/serving/speech_api.md
+++ b/docs/serving/speech_api.md
@@ -109,7 +109,7 @@ Content-Type: application/json
 | `input` | string | **required** | The text to synthesize into speech |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
 | `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25-4.0) |
 
 #### vLLM-Omni Extension Parameters
diff --git a/docs/source/architecture/omni-modality-model-architecture.png b/docs/source/architecture/omni-modality-model-architecture.png
index 53978b323..7bbd5b9b3 100644
Binary files a/docs/source/architecture/omni-modality-model-architecture.png and b/docs/source/architecture/omni-modality-model-architecture.png differ
diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md
index abc63ebed..50e02971f 100644
--- a/docs/user_guide/diffusion/attention_backends.md
+++ b/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
@@ -81,7 +83,7 @@ When constructing `OmniDiffusionConfig` directly:
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
 
 config = OmniDiffusionConfig(
-    attention_config=AttentionConfig(
+    diffusion_attention_config=AttentionConfig(
         default=AttentionSpec(backend="FLASH_ATTN"),
         per_role={
             "cross": AttentionSpec(backend="TORCH_SDPA"),
@@ -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)
diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md
index f827e79b3..d6fcd8944 100644
--- a/docs/user_guide/diffusion/cpu_offload_diffusion.md
+++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md
@@ -33,7 +33,7 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
 
 **CLI:**
 ```bash
-vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
 ```
 
 ### To Support a Model
@@ -82,6 +82,33 @@ mutual exclusion: when one group runs, the other moves to CPU.
 - Support single GPU only for now
 
 
+### Component offloading for split models (e.g. Cosmos3)
+
+Some models split their transformer into mutually-exclusive *components* that run
+in different phases of a single forward pass rather than as separate pipeline
+components -- e.g. Cosmos3's understanding (reasoner) component runs once per
+generation while the generation (generator) component runs every denoising step.
+Such models have no separate text encoder to swap against, so the transformer
+owns a small model-local offload path and wraps each phase with
+`with self._offload_context(name):`
+
+```python
+class Cosmos3VFMTransformer(nn.Module):
+    def forward(self, ...):
+        with self._offload_context("reasoner"):
+            ...  # understanding pass, runs once
+        with self._offload_context("generator"):
+            ...  # denoising pass, runs every step
+```
+
+Model-level offloading then keeps exactly one component GPU-resident at a time
+(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The
+pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the
+transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading
+works for these models too -- each component declares its own block container via
+`_layerwise_offload_blocks_attrs`.
+
+
 ## Layerwise (Blockwise) Offloading
 
 ### How It Works
@@ -119,10 +146,10 @@ m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True
 **CLI:**
 ```bash
 # Text-to-video
-vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
 
 # Or image-to-video
-vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
 ```
 
 ### To Support a Model
@@ -145,9 +172,9 @@ class Flux2Transformer2DModel(nn.Module):
 ```
 
 ### Limitations
-- Cold start latency increases because of
-    1) components are loaded to CPU first at the very first during initialization,
-    2) weight consolidation and pinning
+- Cold start latency increases because offloaded components must be moved to CPU
+  during setup; layerwise offload may add extra weight consolidation and pinning
+  work.
 - Performance depends on compute cost and H2D bandwidth as well
 - Support single GPU only for now
 
@@ -178,11 +205,18 @@ Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelH
 
 ```
 OffloadBackend (base class)
-├── ModelLevelOffloadBackend → uses SequentialOffloadHook
+├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
+│                              (delegates to a pipeline's enable_omni_model_cpu_offload
+│                               for split models like Cosmos3)
 └── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
 ```
 
-Factory function `get_offload_backend()` selects the appropriate backend based on configuration.
+Factory function `get_offload_backend()` selects the appropriate backend based on
+configuration.
+
+For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's
+`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its
+reasoner/generator components inside the model forward pass.
 
 
 ## Supported Models
diff --git a/docs/user_guide/diffusion/frame_interpolation.md b/docs/user_guide/diffusion/frame_interpolation.md
index ab03363cf..817ae67f7 100644
--- a/docs/user_guide/diffusion/frame_interpolation.md
+++ b/docs/user_guide/diffusion/frame_interpolation.md
@@ -47,7 +47,8 @@ For supported Wan2.2 pipelines, the execution order is:
 1. Diffusion worker finishes denoising and decodes the raw video tensor.
 2. Worker-side model-specific post-processing runs.
 3. If frame interpolation is enabled, RIFE interpolates the decoded video
-   tensor on the worker side and records a FPS multiplier in `custom_output`.
+   tensor on the worker side and records a FPS multiplier in
+   `metadata.video.video_fps_multiplier`.
 4. The API server receives the already-interpolated video and only performs
    MP4 export.
 
diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md
index 6b923a793..d3ed3bd67 100644
--- a/docs/user_guide/diffusion_features.md
+++ b/docs/user_guide/diffusion_features.md
@@ -105,7 +105,7 @@ The following tables show which models support each feature:
 
 > Notes:
 
-> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
+> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md).
 > 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
 
 ### ImageGen
@@ -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)      |       ✅        |        ✅         |
@@ -143,7 +144,8 @@ The following tables show which models support each feature:
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
 > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
-> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
+> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
+> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 
 ### VideoGen
 
diff --git a/docs/user_guide/examples/offline_inference/bagel.md b/docs/user_guide/examples/offline_inference/bagel.md
index 0d3498b28..391c41b35 100644
--- a/docs/user_guide/examples/offline_inference/bagel.md
+++ b/docs/user_guide/examples/offline_inference/bagel.md
@@ -257,7 +257,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
 | :------- | :--- | :------ | :---------- |
 | `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
 | `--stage-configs-path` | string | `None` | [Deprecated] Legacy path to `stage_args` YAML; prefer `--deploy-config` |
-| `--worker-backend` | choice | `process` | `process` or `ray` |
+| `--worker-backend` | choice | `multi_process` | `multi_process` or `ray` |
 | `--ray-address` | string | `None` | Ray cluster address |
 | `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
 | `--log-stats` | flag | `False` | Enable statistics logging |
diff --git a/docs/user_guide/examples/offline_inference/glm_image.md b/docs/user_guide/examples/offline_inference/glm_image.md
index c6ac6e33f..09d927ff7 100644
--- a/docs/user_guide/examples/offline_inference/glm_image.md
+++ b/docs/user_guide/examples/offline_inference/glm_image.md
@@ -26,18 +26,22 @@ Stage 0 (AR Model)                Stage 1 (Diffusion)
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 if __name__ == "__main__":
     omni = Omni(model="zai-org/GLM-Image")
     outputs = omni.generate(
         "A photorealistic mountain landscape at sunset",
-        sampling_params={
-            "height": 1024,
-            "width": 1024,
-            "num_inference_steps": 50,
-            "guidance_scale": 1.5,
-            "seed": 42,
-        },
+        sampling_params_list=[
+            omni.default_sampling_params_list[0],
+            OmniDiffusionSamplingParams(
+                height=1024,
+                width=1024,
+                num_inference_steps=50,
+                guidance_scale=1.5,
+                seed=42,
+            ),
+        ],
     )
     outputs[0].request_output.images[0].save("output.png")
 ```
@@ -46,6 +50,7 @@ if __name__ == "__main__":
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 if __name__ == "__main__":
     omni = Omni(model="zai-org/GLM-Image")
@@ -56,13 +61,16 @@ if __name__ == "__main__":
                 "image": "input.png",
             },
         },
-        sampling_params={
-            "height": 1024,
-            "width": 1024,
-            "num_inference_steps": 50,
-            "guidance_scale": 1.5,
-            "seed": 42,
-        },
+        sampling_params_list=[
+            omni.default_sampling_params_list[0],
+            OmniDiffusionSamplingParams(
+                height=1024,
+                width=1024,
+                num_inference_steps=50,
+                guidance_scale=1.5,
+                seed=42,
+            ),
+        ],
     )
     outputs[0].request_output.images[0].save("output.png")
 ```
diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md
index 6e105741a..dc1d3d6b2 100644
--- a/docs/user_guide/examples/offline_inference/image_to_video.md
+++ b/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.
diff --git a/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md b/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md
index 6ab31298c..60f1cbe5f 100644
--- a/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md
+++ b/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
 ```
diff --git a/docs/user_guide/examples/offline_inference/mimo_audio.md b/docs/user_guide/examples/offline_inference/mimo_audio.md
index 1cba2f77d..a43d8e408 100644
--- a/docs/user_guide/examples/offline_inference/mimo_audio.md
+++ b/docs/user_guide/examples/offline_inference/mimo_audio.md
@@ -1,4 +1,4 @@
-# MiMo-Audio Offline Inference
+# MiMo-Audio: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/mimo_audio>.
 
diff --git a/docs/user_guide/examples/offline_inference/qwen2_5_omni.md b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md
index c54976b54..efd14bc1c 100644
--- a/docs/user_guide/examples/offline_inference/qwen2_5_omni.md
+++ b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md
@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>.
 
diff --git a/docs/user_guide/examples/offline_inference/qwen3_omni.md b/docs/user_guide/examples/offline_inference/qwen3_omni.md
index 2d856f738..3781c1cc8 100644
--- a/docs/user_guide/examples/offline_inference/qwen3_omni.md
+++ b/docs/user_guide/examples/offline_inference/qwen3_omni.md
@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Offline inference
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_omni>.
 
diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md
index a09dbfc97..5d89eac34 100644
--- a/docs/user_guide/examples/offline_inference/text_to_video.md
+++ b/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.
 
diff --git a/docs/user_guide/examples/online_serving/aura_omni.md b/docs/user_guide/examples/online_serving/aura_omni.md
index 9c66d0df9..8b7f56ac5 100644
--- a/docs/user_guide/examples/online_serving/aura_omni.md
+++ b/docs/user_guide/examples/online_serving/aura_omni.md
@@ -1,4 +1,4 @@
-# AURA Omni Pipeline
+# AURA Omni: Online serving
 
 `aura_omni` wires ASR, AURA, and Qwen3-TTS into one vLLM-Omni pipeline:
 
diff --git a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md
index 1adc61193..c6be3c014 100644
--- a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md
+++ b/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
diff --git a/docs/user_guide/examples/online_serving/mimo_audio.md b/docs/user_guide/examples/online_serving/mimo_audio.md
index c8752f578..9db59f9e8 100644
--- a/docs/user_guide/examples/online_serving/mimo_audio.md
+++ b/docs/user_guide/examples/online_serving/mimo_audio.md
@@ -1,4 +1,4 @@
-# Online serving Example of vLLM-Omni for MiMo-Audio
+# MiMo-Audio: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/mimo_audio>.
 
diff --git a/docs/user_guide/examples/online_serving/qwen2_5_omni.md b/docs/user_guide/examples/online_serving/qwen2_5_omni.md
index b3a2c9f2a..d8c48a403 100644
--- a/docs/user_guide/examples/online_serving/qwen2_5_omni.md
+++ b/docs/user_guide/examples/online_serving/qwen2_5_omni.md
@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen2_5_omni>.
 
diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md
index 1973b198b..ec509d597 100644
--- a/docs/user_guide/examples/online_serving/qwen3_omni.md
+++ b/docs/user_guide/examples/online_serving/qwen3_omni.md
@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Online serving
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni>.
 
diff --git a/docs/user_guide/examples/online_serving/text_to_audio.md b/docs/user_guide/examples/online_serving/text_to_audio.md
index b1ec33a20..a4833aef6 100644
--- a/docs/user_guide/examples/online_serving/text_to_audio.md
+++ b/docs/user_guide/examples/online_serving/text_to_audio.md
@@ -139,7 +139,7 @@ Content-Type: application/json
 |-----------|------|---------|-------------|
 | `input` | string | **required** | Text prompt describing the audio to generate |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, opus |
 | `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
 | `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
 | `audio_start` | float | 0.0 | Audio start time in seconds |
diff --git a/docs/user_guide/quantization/autoround.md b/docs/user_guide/quantization/autoround.md
index 3afc8b062..bafadec9c 100644
--- a/docs/user_guide/quantization/autoround.md
+++ b/docs/user_guide/quantization/autoround.md
@@ -3,7 +3,7 @@
 ## Overview
 
 [AutoRound](https://github.com/intel/auto-round) produces pre-quantized
-checkpoints for LLMs, VLMs, and diffusion models. vLLM-Omni reads the
+checkpoints for LLMs, VLMs, diffusion models, and world models. vLLM-Omni reads the
 checkpoint's `config.json` and auto-detects
 `quantization_config.quant_method = "auto-round"`.
 
@@ -36,6 +36,13 @@ guide. AutoRound is Intel-supported.
 | Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
 | Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
 
+### World Model (Cosmos3)
+
+| Model | Checkpoint | Scope | Scheme | Backend |
+|-------|------------|-------|--------|---------|
+| Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+
 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
 
 | Model | Checkpoint | Scope | Scheme | Backend |
diff --git a/docs/user_guide/quantization/bitsandbytes.md b/docs/user_guide/quantization/bitsandbytes.md
new file mode 100644
index 000000000..37ccd2471
--- /dev/null
+++ b/docs/user_guide/quantization/bitsandbytes.md
@@ -0,0 +1,120 @@
+# BitsAndBytes W4 Quantization
+
+## Overview
+
+BitsAndBytes 4-bit quantization supports weight-only NF4/FP4 diffusion transformer
+inference on CUDA GPUs. It quantizes BF16/FP16 weights at load time; activations
+stay in BF16/FP16. No pre-quantized checkpoint is required.
+
+This is an online (dynamic) path: load the normal HuggingFace checkpoint and
+quantize during model loading via the `bitsandbytes` CUDA kernels.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA CUDA GPU (SM 75+) | ✅ |
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ❌ |
+| Intel XPU | ❌ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported. Non-CUDA platforms raise in
+`get_quant_method()`; CUDA requires compute capability 7.5+ (`SM 75+`).
+
+Requires the optional `bitsandbytes` package (`pip install bitsandbytes`).
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF models | CUDA | Mode | Recommendation |
+|-------|-----------|:----:|------|----------------|
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Online W4 weight-only | All heavy linear layers; sensitive embedders stay BF16 |
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Not validated | Online W4 weight-only | Compare vs BF16 before enabling |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Online W4 weight-only | Validate before enabling in docs |
+
+Other diffusion models may work if their transformer uses supported linear
+layers, but they are not validated in this guide.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker language-model stage | Not validated | Prefer checkpoint-supported ModelOpt FP8 or AutoRound paths |
+| Qwen3-TTS | TTS language-model stage | Not validated | No BitsAndBytes TTS stage support is documented |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer or DiT module | Not validated | Requires explicit stage routing |
+| GLM-Image | Stage-specific transformer or DiT module | Not validated | Requires quality comparison with BF16 |
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="Tongyi-MAI/Z-Image-Turbo", quantization="bitsandbytes")
+
+omni_with_skips = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    quantization_config={
+        "method": "bitsandbytes",
+        "ignored_layers": ["to_out"],
+    },
+)
+
+outputs = omni.generate(
+    "A cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI:
+
+```bash
+python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes
+python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes --ignored-layers "to_out"
+vllm serve Tongyi-MAI/Z-Image-Turbo --omni --quantization bitsandbytes
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"bitsandbytes"`) |
+| `quant_type` | str | `"nf4"` | 4-bit data type: `"nf4"` (recommended) or `"fp4"` |
+| `compress_statistics` | bool | `True` | Double-quantize block scaling statistics for better accuracy |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
+
+## Validation and Notes
+
+On Z-Image-Turbo (single GPU, 1024×1024, 50 steps), BitsAndBytes W4 typically
+reduces peak VRAM from roughly 24.5 GiB (BF16) to roughly 17 GiB. Compare output
+quality against a BF16 baseline before enabling on new models.
+
+Multi-GPU tensor parallelism (`tensor_parallel_size` > 1) is not validated for
+BitsAndBytes in diffusion models; each rank quantizes its own weight shard
+independently.
+
+If quality regresses, use `ignored_layers` to keep sensitive projections in BF16
+(for example `to_out` or `w2`).
+
+```python
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    quantization="bitsandbytes",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```
+
+Only add a new model to the supported table after comparing the BitsAndBytes
+output against a BF16 baseline and documenting any required `ignored_layers`.
diff --git a/docs/user_guide/quantization/gguf.md b/docs/user_guide/quantization/gguf.md
index e70d2d26e..206e25319 100644
--- a/docs/user_guide/quantization/gguf.md
+++ b/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 |
diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md
index 401a55ad4..b18b3e7b3 100644
--- a/docs/user_guide/quantization/mxfp4.md
+++ b/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) | ✅ |
 
diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md
index 38d027f12..3a4a45573 100644
--- a/docs/user_guide/quantization/overview.md
+++ b/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 | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
 
diff --git a/examples/model_display_names.yml b/examples/model_display_names.yml
new file mode 100644
index 000000000..be3f8e381
--- /dev/null
+++ b/examples/model_display_names.yml
@@ -0,0 +1,12 @@
+# Canonical display names for model examples that have both offline and online pages.
+# The docs generator combines these names with the serving mode for each page H1.
+aura_omni: AURA Omni
+covo_audio: Covo-Audio-Chat
+dynin_omni: Dynin-Omni
+lance: Lance
+mimo_audio: MiMo-Audio
+ming_flash_omni: Ming-flash-omni 2.0
+qwen2_5_omni: Qwen2.5-Omni
+qwen3_omni: Qwen3-Omni
+sensenova_u1: SenseNova-U1
+step_audio2: Step-Audio2
diff --git a/examples/offline_inference/aura_omni/README.md b/examples/offline_inference/aura_omni/README.md
index 37fabf1ba..b73976a15 100644
--- a/examples/offline_inference/aura_omni/README.md
+++ b/examples/offline_inference/aura_omni/README.md
@@ -1,4 +1,4 @@
-# AURA Omni
+# AURA Omni: Offline inference
 
 This example runs the native AURA Omni pipeline offline:
 
diff --git a/examples/offline_inference/covo_audio/README.md b/examples/offline_inference/covo_audio/README.md
index d13783d39..80a822fa2 100644
--- a/examples/offline_inference/covo_audio/README.md
+++ b/examples/offline_inference/covo_audio/README.md
@@ -1,4 +1,4 @@
-# Covo-Audio-Chat (Offline Inference)
+# Covo-Audio-Chat: Offline inference
 
 ## Setup
 
diff --git a/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py
index b1c74e2b3..fb14d9d33 100644
--- a/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py
+++ b/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)
diff --git a/examples/offline_inference/dynin_omni/README.md b/examples/offline_inference/dynin_omni/README.md
index 0b1cc22f2..19fa8c69e 100644
--- a/examples/offline_inference/dynin_omni/README.md
+++ b/examples/offline_inference/dynin_omni/README.md
@@ -1,4 +1,4 @@
-# Dynin-Omni Offline End2End Example
+# Dynin-Omni: Offline inference
 
 This folder contains a unified offline inference entrypoint:
 
diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py
index 299efb33b..12bd978a3 100644
--- a/examples/offline_inference/image_to_image/image_edit.py
+++ b/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")
diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md
index 68bddca06..c161022a5 100644
--- a/examples/offline_inference/image_to_video/README.md
+++ b/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.
diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py
index cc8f1ff2a..2d1e70697 100644
--- a/examples/offline_inference/image_to_video/image_to_video.py
+++ b/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)
diff --git a/examples/offline_inference/lance/README.md b/examples/offline_inference/lance/README.md
index d029abdea..f3a422a3b 100644
--- a/examples/offline_inference/lance/README.md
+++ b/examples/offline_inference/lance/README.md
@@ -1,4 +1,4 @@
-# Lance (ByteDance) — offline inference
+# Lance: Offline inference
 
 [Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
 autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is
diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py
index 504dceec6..d26910161 100644
--- a/examples/offline_inference/lance/end2end.py
+++ b/examples/offline_inference/lance/end2end.py
@@ -293,16 +293,19 @@ def main():
 
     if args.modality in ("img2text", "video2text"):
         for i, req_output in enumerate(outputs):
-            text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
-            if text is None and hasattr(req_output, "custom_output"):
-                text = (req_output.custom_output or {}).get("text_output")
+            multimodal_output = getattr(req_output, "multimodal_output", {}) or {}
+            metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+            text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+            text = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or (
+                text_metadata.get("text_output") if isinstance(text_metadata, dict) else None
+            )
+            text = text or getattr(req_output, "output", None) or getattr(req_output, "text", None)
             print(f"[Output {i}] {text}")
         return
 
     if args.modality in ("text2video", "video2video", "image2video"):
         for i, req_output in enumerate(outputs):
-            custom = getattr(req_output, "custom_output", None) or {}
-            frames = custom.get("video_frames")
+            frames = getattr(req_output, "images", None)
             if not frames:
                 print(f"[Output {i}] no video frames returned")
                 continue
diff --git a/examples/offline_inference/lance/gradio_demo.py b/examples/offline_inference/lance/gradio_demo.py
index 7dfeb46c9..756b006cb 100644
--- a/examples/offline_inference/lance/gradio_demo.py
+++ b/examples/offline_inference/lance/gradio_demo.py
@@ -835,12 +835,15 @@ def run_task(
     print(f"[unified] task={task} {width}x{height} frames={num_frames} cfg={cfg} steps={steps} seed={seed}", flush=True)
 
     outputs = list(omni.generate(prompts=formatted, sampling_params_list=params))
-    custom = getattr(outputs[0], "custom_output", None) or {}
-    # Lance/Bagel return text at outputs[0].output (string) or
-    # custom_output["text_output"].
+    multimodal_output = getattr(outputs[0], "multimodal_output", {}) or {}
+    metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+    text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
     text_out = None
     if task in {"x2t_image", "x2t_video"}:
-        text_out = custom.get("text_output") or getattr(outputs[0], "output", None)
+        text_out = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or (
+            text_metadata.get("text_output") if isinstance(text_metadata, dict) else None
+        )
+        text_out = text_out or getattr(outputs[0], "output", None)
         if not isinstance(text_out, str):
             text_out = None
 
@@ -856,13 +859,11 @@ def run_task(
             gr.update(value=str(text_out or "(empty)"), visible=True),
         )
     if task in {"t2v", "i2v", "video_edit"}:
-        frames = custom.get("video_frames")
+        frames = getattr(outputs[0], "images", None)
         if frames is None:
             raise gr.Error("No video frames in output.")
         return (gr.update(visible=False), gr.update(value=_save_video(frames), visible=True), gr.update(visible=False))
-    # image tasks — orchestrator strips ``outputs[0].output`` across the
-    # IPC boundary but preserves ``custom_output``, so prefer the latter.
-    img_out = custom.get("image") or getattr(outputs[0], "output", None)
+    img_out = (getattr(outputs[0], "images", None) or [None])[0] or getattr(outputs[0], "output", None)
     if img_out is None:
         raise gr.Error("No image in output.")
     return (gr.update(value=_save_image(img_out), visible=True), gr.update(visible=False), gr.update(visible=False))
diff --git a/examples/offline_inference/mammothmodal2_preview/README.md b/examples/offline_inference/mammothmodal2_preview/README.md
index 16937680a..bb68f3845 100644
--- a/examples/offline_inference/mammothmodal2_preview/README.md
+++ b/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
 ```
diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py
index ca87b9e9a..7a22803c9 100644
--- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py
+++ b/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,
     )
diff --git a/examples/offline_inference/mimo_audio/README.md b/examples/offline_inference/mimo_audio/README.md
index 5615dea51..60f433a4f 100644
--- a/examples/offline_inference/mimo_audio/README.md
+++ b/examples/offline_inference/mimo_audio/README.md
@@ -1,4 +1,4 @@
-# MiMo-Audio Offline Inference
+# MiMo-Audio: Offline inference
 
 This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally.
 
diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md
index 7e7ab0a8f..ee7e93fd0 100644
--- a/examples/offline_inference/ming_flash_omni/README.md
+++ b/examples/offline_inference/ming_flash_omni/README.md
@@ -1,4 +1,4 @@
-# Ming-flash-omni 2.0
+# Ming-flash-omni 2.0: Offline inference
 
 [Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
 
@@ -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
 
diff --git a/examples/offline_inference/qwen2_5_omni/README.md b/examples/offline_inference/qwen2_5_omni/README.md
index e2eae8a96..059c10226 100644
--- a/examples/offline_inference/qwen2_5_omni/README.md
+++ b/examples/offline_inference/qwen2_5_omni/README.md
@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Offline inference
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
@@ -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"}}'
+```
diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py
index 8290d27b5..d8769b5f7 100644
--- a/examples/offline_inference/qwen2_5_omni/end2end.py
+++ b/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()
 
 
diff --git a/examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml b/examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml
new file mode 100644
index 000000000..f0e2d9bf5
--- /dev/null
+++ b/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
diff --git a/examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml b/examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml
new file mode 100644
index 000000000..02d773952
--- /dev/null
+++ b/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
diff --git a/examples/offline_inference/qwen3_omni/README.md b/examples/offline_inference/qwen3_omni/README.md
index 0710faa13..2f4ee11cc 100644
--- a/examples/offline_inference/qwen3_omni/README.md
+++ b/examples/offline_inference/qwen3_omni/README.md
@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Offline inference
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
diff --git a/examples/offline_inference/sensenova_u1/README.md b/examples/offline_inference/sensenova_u1/README.md
index 1e656c22d..9fdc23a00 100644
--- a/examples/offline_inference/sensenova_u1/README.md
+++ b/examples/offline_inference/sensenova_u1/README.md
@@ -1,4 +1,4 @@
-# SenseNova-U1-8B-MoT
+# SenseNova-U1: Offline inference
 
 ## Architecture
 
diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py
index 7c727449e..9547c40b9 100644
--- a/examples/offline_inference/sensenova_u1/end2end.py
+++ b/examples/offline_inference/sensenova_u1/end2end.py
@@ -278,12 +278,17 @@ def main():
     )
 
     for req_output in outputs:
-        custom = getattr(req_output, "_custom_output", {}) or {}
-        if args.print_think and custom.get("think_text"):
-            print(f"[Think]\n{custom['think_text']}\n")
+        multimodal_output = getattr(req_output, "multimodal_output", {}) or {}
+        metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+        text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+        think_text = text_metadata.get("think_text") if isinstance(text_metadata, dict) else None
+        if args.print_think and think_text:
+            print(f"[Think]\n{think_text}\n")
 
         if is_text_output:
-            text = custom.get("text_output", "")
+            text = multimodal_output.get("text", "") if isinstance(multimodal_output, dict) else ""
+            if not text:
+                text = text_metadata.get("text_output", "") if isinstance(text_metadata, dict) else ""
             if not text:
                 text = getattr(req_output, "text", "") or ""
             print(f"[Response]\n{text}")
diff --git a/examples/offline_inference/step_audio2/README.md b/examples/offline_inference/step_audio2/README.md
index 2ba69ceb5..84d0ecdde 100644
--- a/examples/offline_inference/step_audio2/README.md
+++ b/examples/offline_inference/step_audio2/README.md
@@ -1,4 +1,4 @@
-# Step-Audio2 Offline Inference Examples
+# Step-Audio2: Offline inference
 
 This directory contains examples for running offline inference with Step-Audio2 using vLLM-Omni.
 
diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md
index 60ca0fb8e..985c4ff7d 100644
--- a/examples/offline_inference/text_to_image/README.md
+++ b/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`
diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py
index 69f9e4fcb..61dd29261 100644
--- a/examples/offline_inference/text_to_image/text_to_image.py
+++ b/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,12 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "int8", "gguf"],
+        choices=["fp8", "int8", "bitsandbytes"],
         help="Quantization method for the transformer. "
-        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), "
+        "'int8' (Int8 W8A8), 'bitsandbytes' (NF4 4-bit weight-only, CUDA SM 75+). "
         "Default: None (no quantization, uses BF16).",
     )
-    parser.add_argument(
-        "--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 +230,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 +397,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 +417,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 +430,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 +462,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 +534,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.
diff --git a/examples/offline_inference/text_to_speech/omnivoice/end2end.py b/examples/offline_inference/text_to_speech/omnivoice/end2end.py
index 619f8160b..a6bdc7b4e 100644
--- a/examples/offline_inference/text_to_speech/omnivoice/end2end.py
+++ b/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,
     )
 
diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md
index a2634d5ae..edc75bb26 100644
--- a/examples/offline_inference/text_to_video/text_to_video.md
+++ b/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).
 
diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py
index 1323985e5..0f9bb8c8e 100644
--- a/examples/offline_inference/text_to_video/text_to_video.py
+++ b/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)
diff --git a/examples/online_serving/aura_omni/README.md b/examples/online_serving/aura_omni/README.md
index 6abd0ef47..5f246c336 100644
--- a/examples/online_serving/aura_omni/README.md
+++ b/examples/online_serving/aura_omni/README.md
@@ -1,4 +1,4 @@
-# AURA Omni Native Pipeline
+# AURA Omni: Online serving
 
 `aura_omni` serves AURA as a native multi-stage vLLM-Omni pipeline:
 
diff --git a/examples/online_serving/covo_audio/README.md b/examples/online_serving/covo_audio/README.md
index 6894507ef..b6a4ad74a 100644
--- a/examples/online_serving/covo_audio/README.md
+++ b/examples/online_serving/covo_audio/README.md
@@ -1,4 +1,4 @@
-# Covo-Audio-Chat
+# Covo-Audio-Chat: Online serving
 
 ## Installation
 
diff --git a/examples/online_serving/diffusers_pipeline_adapter/README.md b/examples/online_serving/diffusers_pipeline_adapter/README.md
index d6f41f632..37bf48aa0 100644
--- a/examples/online_serving/diffusers_pipeline_adapter/README.md
+++ b/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
diff --git a/examples/online_serving/dynin_omni/README.md b/examples/online_serving/dynin_omni/README.md
index 50e2f3f80..6e78a5c1f 100644
--- a/examples/online_serving/dynin_omni/README.md
+++ b/examples/online_serving/dynin_omni/README.md
@@ -1,4 +1,4 @@
-# Dynin-Omni Online Serving Example
+# Dynin-Omni: Online serving
 
 ## Installation
 
diff --git a/examples/online_serving/lance/README.md b/examples/online_serving/lance/README.md
index 5f2febce9..614f4990c 100644
--- a/examples/online_serving/lance/README.md
+++ b/examples/online_serving/lance/README.md
@@ -1,4 +1,4 @@
-# Lance — online serving
+# Lance: Online serving
 
 OpenAI-compatible chat completions API for Lance, served by `vllm-omni`.
 
diff --git a/examples/online_serving/mimo_audio/README.md b/examples/online_serving/mimo_audio/README.md
index 9c1be7f21..7c5e2db0f 100644
--- a/examples/online_serving/mimo_audio/README.md
+++ b/examples/online_serving/mimo_audio/README.md
@@ -1,4 +1,4 @@
-# Online serving Example of vLLM-Omni for MiMo-Audio
+# MiMo-Audio: Online serving
 
 ## 🛠️ Installation
 
diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md
index d60ba8156..377acc566 100644
--- a/examples/online_serving/ming_flash_omni/README.md
+++ b/examples/online_serving/ming_flash_omni/README.md
@@ -1,4 +1,4 @@
-# Ming-flash-omni 2.0
+# Ming-flash-omni 2.0: Online serving
 
 ## Installation
 
@@ -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
 
diff --git a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh
deleted file mode 100755
index 92bb32c47..000000000
--- a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh
+++ /dev/null
@@ -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')"
diff --git a/examples/online_serving/qwen2_5_omni/README.md b/examples/online_serving/qwen2_5_omni/README.md
index c52873206..48406003d 100644
--- a/examples/online_serving/qwen2_5_omni/README.md
+++ b/examples/online_serving/qwen2_5_omni/README.md
@@ -1,4 +1,4 @@
-# Qwen2.5-Omni
+# Qwen2.5-Omni: Online serving
 
 ## 🛠️ Installation
 
diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md
index 5e2e24a1f..3c3ad2b16 100644
--- a/examples/online_serving/qwen3_omni/README.md
+++ b/examples/online_serving/qwen3_omni/README.md
@@ -1,4 +1,4 @@
-# Qwen3-Omni
+# Qwen3-Omni: Online serving
 
 ## 🛠️ Installation
 
diff --git a/examples/online_serving/replica_data_parallel/README.md b/examples/online_serving/replica_data_parallel/README.md
new file mode 100644
index 000000000..4bc98d6f7
--- /dev/null
+++ b/examples/online_serving/replica_data_parallel/README.md
@@ -0,0 +1,91 @@
+# Replica Data Parallelism for Video Generation
+
+A runnable recipe + benchmark for **replica data parallelism** (replica-DP) on a
+video DiT. Replica-DP runs `N` independent diffusion engine replicas, one per GPU,
+and routes each request to a single replica. It scales **request throughput**
+near-linearly; it does **not** change single-request latency (for that, use an
+intra-request axis such as tensor/sequence/CFG parallelism).
+
+This example uses `Wan-AI/Wan2.2-TI2V-5B-Diffusers` with one GPU per replica.
+
+## Files
+
+| File | Purpose |
+|------|---------|
+| `wan2_2_ti2v_dp.yaml` | Diffusion stage config; `runtime.num_replicas` + `runtime.devices` drive the fan-out. |
+| `run_server.sh` | Substitutes `NUM_REPLICAS` / `DEVICES` into the config and serves. |
+| `bench_replica_dp.py` | Replica-agnostic load driver; reports throughput, latency, and a per-input isolation check. |
+
+## Run
+
+Start the server with the replica count you want (one GPU per replica):
+
+```bash
+# baseline
+NUM_REPLICAS=1 DEVICES=0        ./run_server.sh
+# 2 replicas
+NUM_REPLICAS=2 DEVICES=0,1      ./run_server.sh
+# 4 replicas
+NUM_REPLICAS=4 DEVICES=0,1,2,3  ./run_server.sh
+```
+
+In another shell, drive a fixed workload and read the throughput:
+
+```bash
+python bench_replica_dp.py --url http://127.0.0.1:8098 \
+    --num-requests 28 --concurrency 8 --label "replicas=4"
+```
+
+Re-run the server at `NUM_REPLICAS=1,2,4` (same client workload each time) and
+compare `videos/min`. Pass `--baseline-tpm <N=1 value>` to print a scaling factor.
+
+### Isolation check
+
+Identical outputs under a *shared* prompt+seed prove nothing -- every request
+would be identical regardless of isolation. Instead each request gets a distinct
+input (a prompt from the pool + a unique seed `base_seed+i`), and we check that a
+request's output is unchanged whether it runs alone or concurrently amid other
+distinct requests. Capture a per-input baseline on `N=1`, then verify it under
+`N>=2` concurrency:
+
+```bash
+# 1) against NUM_REPLICAS=1 (serial), record each input's output hash
+python bench_replica_dp.py --num-requests 8 --concurrency 1 \
+    --write-baseline baseline.json
+
+# 2) against NUM_REPLICAS>=2, replay the SAME inputs concurrently and verify
+python bench_replica_dp.py --num-requests 8 --concurrency 8 \
+    --check-baseline baseline.json
+```
+
+Step 2 exits non-zero if any request's output differs from its own baseline (a
+mismatch means cross-request state bleed or misrouting). Keep `--num-requests`,
+`--base-seed`, and the prompt source identical between the two runs so the input
+sets match.
+
+## Measured scaling
+
+Wan2.2-TI2V-5B (832×480, 33 frames, 30 steps), 4× A800-80GB (NVLink), one GPU per replica:
+
+| Replicas | Throughput (videos/min) | Scaling | Efficiency |
+|----------|-------------------------|---------|------------|
+| 1 | 4.71 | 1.00× | — |
+| 2 | 9.20 | 1.95× | 98% |
+| 4 | 18.02 | 3.83× | 96% |
+
+Per-request latency stayed flat (~13–14 s) across all replica counts. Isolation
+is checked per-input via the baseline/verify protocol above: each request's
+output under `N>=2` concurrency must match its own single-replica baseline.
+
+## Notes on the config surface
+
+- `runtime.num_replicas: N` fans out `N` replicas; `runtime.devices` assigns their
+  GPUs. With `tensor_parallel_size = 1`, list one GPU per replica
+  (`num_replicas * tensor_parallel_size` entries, pool mode).
+- Replica fan-out here comes from the stage config, not a CLI flag. (The headless /
+  multi-runtime launch path uses the process-local `--omni-dp-size-local`, which
+  requires `--stage-id`.)
+- This recipe serves via `--stage-configs-path`. Expressing replica-DP under the
+  newer `--deploy-config` schema is part of the ongoing serving-config alignment
+  (see the diffusion parallelism guide and the `#4707` discussion); `num_replicas`
+  is accepted by both paths.
diff --git a/examples/online_serving/replica_data_parallel/bench_replica_dp.py b/examples/online_serving/replica_data_parallel/bench_replica_dp.py
new file mode 100755
index 000000000..977f475ac
--- /dev/null
+++ b/examples/online_serving/replica_data_parallel/bench_replica_dp.py
@@ -0,0 +1,241 @@
+#!/usr/bin/env python3
+"""Throughput-scaling + isolation benchmark for diffusion replica data parallelism.
+
+Sends `--num-requests` text-to-video requests (up to `--concurrency` in flight)
+to a running vLLM-Omni video server via `POST /v1/videos/sync` (blocks until the
+clip is produced) and reports:
+
+    completed / failed, makespan (wall clock), throughput (videos/min),
+    per-request latency p50/p90/max (nearest-rank), and -- when a baseline
+    manifest is supplied -- a per-request isolation check.
+
+Each request gets a *distinct* input (a prompt from the pool + a unique seed),
+so an output is a deterministic function of its own input and nothing else.
+
+Measuring replica scaling
+-------------------------
+This client is replica-agnostic. Vary the *server's* replica count
+(`runtime.num_replicas` in the stage config; see run_server.sh) and keep the
+client workload fixed, then compare throughput:
+
+    N=1 -> baseline
+    N=2 -> expect ~2x
+    N=4 -> expect ~4x   (near-linear is the goal)
+
+Pass the single-replica videos/min via `--baseline-tpm` to print a scaling
+factor for the current run.
+
+Isolation / correctness
+-----------------------
+Byte-identical outputs under a *shared* prompt+seed prove nothing -- every
+request would be identical regardless of isolation. Instead, give each request a
+distinct input and check that its output is unchanged whether it runs alone or
+concurrently amid other distinct requests:
+
+    1. Against an N=1 server, capture a per-input baseline:
+         bench_replica_dp.py --num-requests 8 --concurrency 1 \
+             --write-baseline baseline.json
+    2. Against an N>=2 server, replay the *same* inputs concurrently and verify
+       each output matches its own baseline:
+         bench_replica_dp.py --num-requests 8 --concurrency 8 \
+             --check-baseline baseline.json
+
+A mismatch means a request's output changed under concurrency -- cross-request
+state bleed or misrouting -- and the run exits non-zero.
+
+Requires: requests (`pip install requests`).
+"""
+
+import argparse
+import concurrent.futures as cf
+import hashlib
+import json
+import math
+import sys
+import time
+from pathlib import Path
+
+import requests
+
+# Distinct prompts, cycled across requests; combined with a per-request unique
+# seed this makes every request a distinct input (see input_key()).
+DEFAULT_PROMPTS = [
+    "A cat playing piano, cinematic, high detail",
+    "A drone shot over a snowy mountain range at sunrise",
+    "A red sports car driving through a neon-lit city at night",
+    "Ocean waves crashing on a rocky shore, slow motion",
+    "A hot air balloon floating over green rolling hills",
+    "Timelapse of clouds moving over a desert canyon at dusk",
+    "A steam train crossing a stone viaduct in autumn",
+    "A jellyfish drifting through deep blue water, backlit",
+]
+
+
+def input_key(prompt, seed, args):
+    """Stable identity of a request's input; matches baseline<->verify by content, not order."""
+    return "\x1f".join([prompt, str(seed), args.size, str(args.num_frames), str(args.steps)])
+
+
+def nearest_rank(xs, q):
+    """Nearest-rank percentile; q in [0, 1]. Correct for tiny samples (unlike int(q*n)-1)."""
+    if not xs:
+        return float("nan")
+    s = sorted(xs)
+    rank = max(1, math.ceil(q * len(s)))
+    return s[rank - 1]
+
+
+def one_request(args, spec, idx, save):
+    """Fire one /v1/videos/sync request; return (ok, latency, nbytes, md5, key, err)."""
+    prompt, seed = spec["prompt"], spec["seed"]
+    key = input_key(prompt, seed, args)
+    # multipart form fields match POST /v1/videos/sync in the omni API server
+    form = {
+        "prompt": (None, prompt),
+        "size": (None, args.size),  # e.g. "832x480"
+        "num_frames": (None, str(args.num_frames)),
+        "num_inference_steps": (None, str(args.steps)),
+        "seed": (None, str(seed)),  # per-request unique -> distinct, reproducible input
+    }
+    if args.model:
+        form["model"] = (None, args.model)
+    if args.negative_prompt:
+        form["negative_prompt"] = (None, args.negative_prompt)
+
+    t0 = time.perf_counter()
+    try:
+        r = requests.post(f"{args.url}/v1/videos/sync", files=form, timeout=args.timeout)
+        dt = time.perf_counter() - t0
+        if r.status_code != 200:
+            return (False, dt, 0, "", key, f"HTTP {r.status_code}: {r.text[:160]}")
+        body = r.content
+        md5 = hashlib.md5(body).hexdigest()
+        if save and args.save_dir:
+            Path(args.save_dir).mkdir(parents=True, exist_ok=True)
+            (Path(args.save_dir) / f"req{idx:03d}_seed{seed}.mp4").write_bytes(body)
+        return (True, dt, len(body), md5, key, "")
+    except Exception as e:  # noqa: BLE001 - report any client-side failure
+        return (False, time.perf_counter() - t0, 0, "", key, repr(e))
+
+
+def build_specs(args):
+    """Deterministic list of distinct (prompt, seed) inputs; identical across baseline/verify runs."""
+    prompts = DEFAULT_PROMPTS
+    if args.prompts_file:
+        prompts = [ln.strip() for ln in Path(args.prompts_file).read_text().splitlines() if ln.strip()]
+    return [{"prompt": prompts[i % len(prompts)], "seed": args.base_seed + i} for i in range(args.num_requests)]
+
+
+def run_batch(args, specs, save):
+    """Fire specs concurrently; return (results, makespan). results = list of one_request tuples."""
+    results = []
+    wall0 = time.perf_counter()
+    with cf.ThreadPoolExecutor(max_workers=args.concurrency) as ex:
+        futs = {ex.submit(one_request, args, spec, i, save): i for i, spec in enumerate(specs)}
+        for done, f in enumerate(cf.as_completed(futs), 1):
+            res = f.result()
+            results.append(res)
+            good, dt, nbytes, _md5, _key, err = res
+            print(f"  [{done}/{len(specs)}] {'ok' if good else 'FAIL'} {dt:6.1f}s {(nbytes / 1e6):5.1f}MB {err}")
+    return results, time.perf_counter() - wall0
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Diffusion replica-DP throughput + isolation benchmark")
+    ap.add_argument("--url", default="http://127.0.0.1:8098", help="server base URL (no path)")
+    ap.add_argument("--num-requests", type=int, default=16, help="total requests to send")
+    ap.add_argument("--concurrency", type=int, default=4, help="max in-flight (set >= replica count)")
+    ap.add_argument("--warmup", type=int, default=-1, help="warmup requests before timing (default: =concurrency)")
+    ap.add_argument("--base-seed", type=int, default=42, help="request i uses seed base_seed+i (distinct inputs)")
+    ap.add_argument("--prompts-file", default="", help="optional file, one prompt per line (else a built-in pool)")
+    ap.add_argument("--negative-prompt", default="")
+    ap.add_argument("--size", default="832x480", help="WIDTHxHEIGHT")
+    ap.add_argument("--num-frames", type=int, default=33)
+    ap.add_argument("--steps", type=int, default=30, help="num_inference_steps")
+    ap.add_argument("--model", default="", help="optional; defaults to the served model")
+    ap.add_argument("--timeout", type=float, default=1800, help="per-request timeout (s)")
+    ap.add_argument("--save-dir", default="", help="save returned clips")
+    ap.add_argument("--label", default="", help="run label, e.g. 'replicas=2'")
+    ap.add_argument(
+        "--baseline-tpm", type=float, default=0.0, help="single-replica videos/min, to print a scaling factor"
+    )
+    ap.add_argument("--write-baseline", default="", help="run against N=1: write per-input md5 manifest to this path")
+    ap.add_argument("--check-baseline", default="", help="run against N>=2: verify each output matches this manifest")
+    args = ap.parse_args()
+
+    if args.write_baseline and args.check_baseline:
+        ap.error("--write-baseline and --check-baseline are mutually exclusive")
+
+    specs = build_specs(args)
+    warmup = args.concurrency if args.warmup < 0 else args.warmup
+
+    print(
+        f"-> {args.url}/v1/videos/sync | {args.num_requests} reqs "
+        f"| concurrency {args.concurrency} | warmup {warmup} | {args.size} {args.num_frames}f {args.steps}steps "
+        f"| {args.label}"
+    )
+
+    # Warmup: absorb each replica's first-call compile/init cost so it does not
+    # skew the measured makespan. Results are discarded and never saved.
+    if warmup > 0:
+        print(f"warming up ({warmup} reqs, not timed)...")
+        run_batch(args, [specs[i % len(specs)] for i in range(warmup)], save=False)
+
+    results, makespan = run_batch(args, specs, save=True)
+
+    ok = sum(1 for r in results if r[0])
+    fail = len(results) - ok
+    lat = [r[1] for r in results if r[0]]
+    errs = [r[5] for r in results if not r[0]]
+    got = {r[4]: r[3] for r in results if r[0]}  # input_key -> md5
+
+    print(f"\n===== result {args.label} =====")
+    print(f"completed {ok} / failed {fail} | makespan {makespan:.1f}s")
+    if ok:
+        tpm = ok / makespan * 60.0
+        print(f"throughput: {tpm:.2f} videos/min  ({ok / makespan:.4f} videos/s)")
+        print(
+            f"latency (s): p50={nearest_rank(lat, 0.50):.1f} "
+            f"p90={nearest_rank(lat, 0.90):.1f} "
+            f"min={min(lat):.1f} max={max(lat):.1f}"
+        )
+        if args.baseline_tpm > 0:
+            print(f"scaling vs baseline ({args.baseline_tpm:.2f}/min): {tpm / args.baseline_tpm:.2f}x")
+    if errs:
+        print("sample failures:", errs[:3])
+
+    # Isolation check: compare each request's output against its own single-replica baseline.
+    isolation_failed = False
+    if args.write_baseline:
+        Path(args.write_baseline).write_text(json.dumps(got, indent=2))
+        print(f"baseline: wrote {len(got)} per-input hashes -> {args.write_baseline}")
+    elif args.check_baseline:
+        baseline = json.loads(Path(args.check_baseline).read_text())
+        matched = mismatched = unknown = 0
+        for key, md5 in got.items():
+            if key not in baseline:
+                unknown += 1
+                print(f"  isolation: NO BASELINE for input {key!r} (not in the baseline manifest)")
+            elif baseline[key] == md5:
+                matched += 1
+            else:
+                mismatched += 1
+                print(f"  isolation: MISMATCH for input {key!r} (output changed under concurrency)")
+        # PASS requires the run to cover the *whole* baseline: a subset (fewer requests, or
+        # some failed) must not pass just because the inputs it did run happened to match.
+        unverified = sorted(set(baseline) - set(got))
+        for key in unverified:
+            print(f"  isolation: NOT VERIFIED for baseline input {key!r} (absent from this run)")
+        isolation_failed = mismatched > 0 or unknown > 0 or len(unverified) > 0
+        verdict = "PASS (isolated)" if not isolation_failed else "FAIL"
+        print(
+            f"isolation: {matched} matched / {mismatched} mismatched / {unknown} unknown / "
+            f"{len(unverified)} unverified vs baseline({len(baseline)}) -> {verdict}"
+        )
+
+    if fail > 0 or ok == 0 or isolation_failed:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/online_serving/replica_data_parallel/run_server.sh b/examples/online_serving/replica_data_parallel/run_server.sh
new file mode 100755
index 000000000..915732eee
--- /dev/null
+++ b/examples/online_serving/replica_data_parallel/run_server.sh
@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+# Serve N independent Wan2.2 TI2V-5B replicas (replica data parallelism).
+#
+# Usage:
+#   NUM_REPLICAS=1 DEVICES=0        ./run_server.sh      # baseline
+#   NUM_REPLICAS=2 DEVICES=0,1      ./run_server.sh
+#   NUM_REPLICAS=4 DEVICES=0,1,2,3  ./run_server.sh
+#
+# Then drive load with bench_replica_dp.py (see README.md) and compare
+# throughput across replica counts.
+set -euo pipefail
+
+MODEL="${MODEL:-Wan-AI/Wan2.2-TI2V-5B-Diffusers}"
+HOST="${HOST:-127.0.0.1}"
+PORT="${PORT:-8098}"
+NUM_REPLICAS="${NUM_REPLICAS:-1}"
+DEVICES="${DEVICES:-0}"   # one GPU per replica for tensor_parallel_size=1
+
+HERE="$(cd "$(dirname "$0")" && pwd)"
+TEMPLATE="$HERE/wan2_2_ti2v_dp.yaml"
+CFG="$(mktemp --suffix=_wan2_2_dp.yaml)"
+trap 'rm -f "$CFG"' EXIT
+
+sed -e "s/__NUM_REPLICAS__/${NUM_REPLICAS}/" \
+    -e "s/__DEVICES__/${DEVICES}/" \
+    "$TEMPLATE" > "$CFG"
+
+echo "Serving ${MODEL}: num_replicas=${NUM_REPLICAS} devices=${DEVICES} (config ${CFG})"
+
+vllm serve "$MODEL" --omni \
+    --host "$HOST" \
+    --port "$PORT" \
+    --stage-configs-path "$CFG"
diff --git a/examples/online_serving/replica_data_parallel/wan2_2_ti2v_dp.yaml b/examples/online_serving/replica_data_parallel/wan2_2_ti2v_dp.yaml
new file mode 100644
index 000000000..b472d7e2d
--- /dev/null
+++ b/examples/online_serving/replica_data_parallel/wan2_2_ti2v_dp.yaml
@@ -0,0 +1,43 @@
+# Replica data-parallel recipe for Wan2.2 TI2V-5B (video DiT).
+#
+# Fans out `num_replicas` independent DiffusionEngine replicas, one per GPU,
+# and lets the serving layer route each request to a single replica. This
+# scales request throughput near-linearly; it does not change single-request
+# latency. See ./README.md for measured numbers and the `run_server.sh` driver.
+#
+# run_server.sh substitutes __NUM_REPLICAS__ / __DEVICES__ from env before serving.
+# For tensor_parallel_size = 1, `devices` lists one GPU per replica (pool mode):
+# num_replicas=4 -> devices "0,1,2,3".
+
+stage_args:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      num_replicas: __NUM_REPLICAS__
+      devices: "__DEVICES__"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      # `model_class_name` takes the diffusion registry KEY (WanPipeline), which
+      # resolves to the Wan22Pipeline class. Wan2.2 TI2V-5B is served by the base
+      # WanPipeline, which auto-detects TI2V-5B mode from model_index.json.
+      model_class_name: WanPipeline
+      max_num_seqs: 1
+      enforce_eager: true
+      trust_remote_code: true
+      distributed_executor_backend: "mp"
+      vae_use_tiling: true
+      parallel_config:
+        tensor_parallel_size: 1
+
+    final_output: true
+    final_output_type: video
+    is_comprehension: false
+    default_sampling_params:
+      seed: 42
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
diff --git a/examples/online_serving/sensenova_u1/README.md b/examples/online_serving/sensenova_u1/README.md
index d66a3506a..d529eedad 100644
--- a/examples/online_serving/sensenova_u1/README.md
+++ b/examples/online_serving/sensenova_u1/README.md
@@ -1,4 +1,4 @@
-# SenseNova-U1 Online Serving
+# SenseNova-U1: Online serving
 
 ## Launch the Server
 
diff --git a/examples/online_serving/step_audio2/README.md b/examples/online_serving/step_audio2/README.md
index 3a04f4565..e93f58ccc 100644
--- a/examples/online_serving/step_audio2/README.md
+++ b/examples/online_serving/step_audio2/README.md
@@ -1,4 +1,4 @@
-# Step-Audio2 Online Serving
+# Step-Audio2: Online serving
 
 This directory contains examples for running Step-Audio2 with vLLM-Omni's online serving API.
 
diff --git a/pyproject.toml b/pyproject.toml
index 2a8fc9c8f..0da2d7b5f 100644
--- a/pyproject.toml
+++ b/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",
diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md
index 269ea6425..7a4f1507c 100644
--- a/recipes/LTX/LTX-2.3.md
+++ b/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_vllm_omni.json` with
+`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. Each case in that file includes a JSON **`mark`** (`hardware_marks` + `full_model` + `diffusion`) for local `-m` filtering. That config captures the
+single-device eager baseline and CFG-parallel=2 case for a small 384x512,
+25-frame, 20-step workload.
diff --git a/recipes/MammothModa2/MammothModa2-Preview.md b/recipes/MammothModa2/MammothModa2-Preview.md
index 289efece0..0733b0630 100644
--- a/recipes/MammothModa2/MammothModa2-Preview.md
+++ b/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 \
diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md
index efc1699b3..8b3d2909d 100644
--- a/recipes/OpenBMB/MiniCPM-o-4_5.md
+++ b/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
diff --git a/recipes/README.md b/recipes/README.md
index 77a78ef7d..d335183ca 100644
--- a/recipes/README.md
+++ b/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 |
diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md
index e58a28623..c3282ba51 100644
--- a/recipes/cosmos3/Cosmos3-Nano.md
+++ b/recipes/cosmos3/Cosmos3-Nano.md
@@ -422,3 +422,185 @@ ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_
   `generate_sound`/`sound_duration`, `guardrails`, `action_*`, ...) are declared
   once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
   unknown keys for the model are dropped.
+
+## NPU
+
+### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving
+
+#### Environment
+
+- OS: Linux (aarch64)
+- Python: 3.12+
+- Driver / runtime: CANN 8.5.1 + NNAL + Ascend 910B / 910C
+- vLLM version: match the repository requirements from your current checkout
+- vLLM-Ascend version: match the repository requirements from your current checkout
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Command
+
+Requires the `vllm-omni` package (or the `quay.io/atlas-ci/vllm-ascend` A2 / A3 container),
+which provides the `vllm serve … --omni` entrypoint used below.
+
+Safety guardrails are **on by default** (NVIDIA Open Model License). They load
+the **gated** `nvidia/Cosmos-1.0-Guardrail` model, so to keep them on you must:
+
+1. `pip install cosmos-guardrail`
+2. Accept the license at <https://huggingface.co/nvidia/Cosmos-1.0-Guardrail>
+3. Export a token with access: `export HF_TOKEN=hf_...`
+
+Then launch the recommended server:
+
+```bash
+vllm serve nvidia/Cosmos3-Nano \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --init-timeout 1800
+```
+
+To run **without** guardrails (you are responsible for license compliance),
+add `--no-guardrails` (no token/`cosmos-guardrail` needed). For tensor parallel
+add `--tensor-parallel-size 8`. `--quantization fp8` and
+`--enable-layerwise-offload` are not supported on NPU.
+The pipeline auto-resolves from `model_index.json`; pass
+`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.
+
+#### Verification
+
+Best quality uses the JSON-upsampled prompts from `assets/` (download with
+`hf download nvidia/Cosmos3-Nano assets/ --local-dir Cosmos3-Nano`). Minimal
+self-contained examples:
+
+```bash
+curl http://localhost:8000/v1/models
+
+# Text-to-image -> /v1/images/generations  (1024x1024, 10 steps; base64 PNG)
+curl -sS -X POST http://localhost:8000/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "nvidia/Cosmos3-Nano",
+    "prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.",
+    "negative_prompt": "blurry, distorted, low quality",
+    "size": "1024x1024", "n": 1, "response_format": "b64_json",
+    "num_inference_steps": 10, "guidance_scale": 7.0, "seed": 42
+  }' | python -c "import sys,json,base64; open('cosmos3_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))"
+
+# Text-to-video -> /v1/videos/sync  (720p, 49 frames @ 24fps)
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=A robot arm is cleaning a plate in the kitchen" \
+  -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
+  -F "size=1280x720" -F "num_frames=49" -F "fps=24" \
+  -F "num_inference_steps=20" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F "seed=123" \
+  -o cosmos3_t2v.mp4
+
+# Image-to-video -> /v1/videos/sync with an uploaded reference image
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=The scene comes to life with smooth, natural motion." \
+  -F "size=1280x720" -F "num_frames=25" -F "fps=8" \
+  -F "num_inference_steps=10" -F "guidance_scale=6.0" \
+  -F "seed=42" \
+  -F "input_reference=@reference.jpg;type=image/jpeg" \
+  -o cosmos3_i2v.mp4
+
+# Video-to-video -> /v1/videos/sync with an uploaded reference video
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
+  -F "size=1280x720" -F "num_frames=17" -F "fps=5" \
+  -F "num_inference_steps=10" -F "guidance_scale=6.0" \
+  -F "seed=42" \
+  -F "input_reference=@reference.mp4;type=video/mp4" \
+  -o cosmos3_v2v.mp4
+```
+
+#### Notes
+
+- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):**
+  - T2I 1024² — 10 steps → ~8 s
+  - T2V 1280×720 @ 20 steps — 49 frames → ~55 s
+  - I2V 1280×720 @ 10 steps — 25 frames → ~25 s
+  - V2V 480×320 @ 10 steps — 17 frames → ~12 s
+- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU;
+  full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
+  ~33 GB on disk.
+- **Determinism:** identical seed reproduces identical output on the same
+  hardware; outputs are not bit-identical across different GPU/NPU types.
+- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
+  3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
+  1280×720, 35 steps, guidance 6.0, `flow_shift=10.0`.
+- **Key flags / params:** `--no-guardrails` (optional, to disable guardrails), `--init-timeout 1800`
+  (for model loading), `--tensor-parallel-size 8` for multi-NPU, and
+  `--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class.
+- **Known limitations:**
+  - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
+    resolution-parsing bug; basic V2V without transfer hints works.
+  - FP8 online quantization and layerwise offload are not supported on NPU.
+
+### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation
+
+#### Environment
+
+- OS: Linux (aarch64)
+- Python: 3.12+
+- Driver / runtime: CANN 8.5.1 + NNAL + Ascend 910B / 910C
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Command
+
+The same offline task examples run on NPU; pass model-specific knobs via
+`--extra-body`. Guardrails are on by default — pass `"guardrails": false` for a
+quick local run (install `cosmos-guardrail` + accept the gated repo to enable them).
+
+```bash
+# Text-to-image -> examples/offline_inference/text_to_image
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model nvidia/Cosmos3-Nano \
+  --prompt "A photorealistic red sports car at golden hour, cinematic lighting." \
+  --negative-prompt "blurry, distorted, low quality" \
+  --height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \
+  --extra-body '{"flow_shift": 3.0, "guardrails": false}' \
+  --output cosmos3_t2i.png
+
+# Text-to-video -> examples/offline_inference/text_to_video
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model nvidia/Cosmos3-Nano \
+  --prompt "A robot arm is cleaning a plate in the kitchen." \
+  --negative-prompt "blurry, distorted, low quality, jittery, deformed" \
+  --height 720 --width 1280 --num-frames 189 --fps 24 \
+  --num-inference-steps 35 --guidance-scale 6.0 \
+  --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
+                 "use_resolution_template": false, "use_duration_template": false}' \
+  --output cosmos3_t2v.mp4
+
+# Image-to-video -> examples/offline_inference/image_to_video
+# (Cosmos3 bundles example frames under assets/; any RGB image works too.)
+python examples/offline_inference/image_to_video/image_to_video.py \
+  --model nvidia/Cosmos3-Nano \
+  --image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
+  --prompt "The scene comes to life with smooth, natural motion." \
+  --height 720 --width 1280 --num-frames 189 --fps 24 \
+  --num-inference-steps 35 --guidance-scale 6.0 \
+  --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
+  --output cosmos3_i2v.mp4
+```
+
+#### Verification
+
+```bash
+python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)"
+ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_t2v.mp4
+```
+
+#### Notes
+
+- Guardrails are on by default on NPU with `cosmos-guardrail` installed. Pass
+  `"guardrails": false` in `--extra-body` to disable them (there is no
+  `--no-guardrails` flag for the offline scripts).
+- Video at the 189-frame default takes ~15 min/clip on 1 NPU; reduce
+  `--num-frames` for faster iteration.
diff --git a/recipes/inclusionAI/Ming-flash-omni-2.0.md b/recipes/inclusionAI/Ming-flash-omni-2.0.md
index 4a3a86669..b9c05606d 100644
--- a/recipes/inclusionAI/Ming-flash-omni-2.0.md
+++ b/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.
diff --git a/recipes/krea/Krea-2.md b/recipes/krea/Krea-2.md
new file mode 100644
index 000000000..80e5105dd
--- /dev/null
+++ b/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.
diff --git a/requirements/common.txt b/requirements/common.txt
index bfc7093f6..243a6ce7e 100644
--- a/requirements/common.txt
+++ b/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
diff --git a/tests/ar_diffusion/__init__.py b/tests/ar_diffusion/__init__.py
new file mode 100644
index 000000000..988131360
--- /dev/null
+++ b/tests/ar_diffusion/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py
new file mode 100644
index 000000000..fb55e01e9
--- /dev/null
+++ b/tests/ar_diffusion/test_engine_routing.py
@@ -0,0 +1,150 @@
+# 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)
+
+
+def test_execute_model_accepts_kv_prefetch_jobs_kwarg():
+    """Regression: the worker passes ``kv_prefetch_jobs`` to every runner
+    (added with the KV prefetch path, #4448); the AR-Diffusion override must
+    accept and forward it or DreamZero crashes on the warm-up dummy run."""
+    import inspect
+
+    from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+    from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
+
+    for cls in (DiffusionModelRunner, ARDiffusionModelRunner):
+        params = inspect.signature(cls.execute_model).parameters
+        assert "kv_prefetch_jobs" in params, cls
diff --git a/tests/ar_diffusion/test_kv_cache.py b/tests/ar_diffusion/test_kv_cache.py
new file mode 100644
index 000000000..beb575a0b
--- /dev/null
+++ b/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)
diff --git a/tests/ar_diffusion/test_kv_orchestration.py b/tests/ar_diffusion/test_kv_orchestration.py
new file mode 100644
index 000000000..4ba5cb49b
--- /dev/null
+++ b/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
diff --git a/tests/ar_diffusion/test_paged_attention.py b/tests/ar_diffusion/test_paged_attention.py
new file mode 100644
index 000000000..1836656f9
--- /dev/null
+++ b/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}"
diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py
new file mode 100644
index 000000000..28c3c959d
--- /dev/null
+++ b/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, kv_prefetch_jobs=None):
+        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
diff --git a/tests/ar_diffusion/test_slot_mapping.py b/tests/ar_diffusion/test_slot_mapping.py
new file mode 100644
index 000000000..061df4d9c
--- /dev/null
+++ b/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
diff --git a/tests/config/composable_parallel/test_apply.py b/tests/config/composable_parallel/test_apply.py
new file mode 100644
index 000000000..044b2cb02
--- /dev/null
+++ b/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")],
+            },
+        )
diff --git a/tests/config/composable_parallel/test_factory_integration.py b/tests/config/composable_parallel/test_factory_integration.py
new file mode 100644
index 000000000..ea685541c
--- /dev/null
+++ b/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)
diff --git a/tests/config/composable_parallel/test_spec_types.py b/tests/config/composable_parallel/test_spec_types.py
new file mode 100644
index 000000000..90b00cf95
--- /dev/null
+++ b/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
diff --git a/tests/config/composable_parallel/test_strategy_loader.py b/tests/config/composable_parallel/test_strategy_loader.py
new file mode 100644
index 000000000..01e202d05
--- /dev/null
+++ b/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"
diff --git a/tests/config/composable_parallel/test_translator.py b/tests/config/composable_parallel/test_translator.py
new file mode 100644
index 000000000..b55c10e9e
--- /dev/null
+++ b/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}"
diff --git a/tests/config/test_endpoint_policy.py b/tests/config/test_endpoint_policy.py
new file mode 100644
index 000000000..aeb1ffa51
--- /dev/null
+++ b/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"
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index 8f15347d1..bf2dba701 100644
--- a/tests/config/test_omni_config.py
+++ b/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
diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py
new file mode 100644
index 000000000..27087d253
--- /dev/null
+++ b/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]
diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py
index a22f760b2..2eb88bbd6 100644
--- a/tests/dfx/conftest.py
+++ b/tests/dfx/conftest.py
@@ -11,11 +11,180 @@ from typing import Any
 import pytest
 
 from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw
+from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import modify_stage_config
 from vllm_omni.platforms import current_omni_platform
 
 
+def _named_pytest_marks(names: list[str]) -> list[pytest.MarkDecorator]:
+    marks: list[pytest.MarkDecorator] = []
+    for name in names:
+        name = name.strip()
+        if not name:
+            raise ValueError("mark name must be a non-empty string")
+        marks.append(getattr(pytest.mark, name))
+    return marks
+
+
+def _hardware_marks_from_dict(hw: Any) -> list[pytest.MarkDecorator]:
+    if not isinstance(hw, dict):
+        raise ValueError(f"mark.hardware_marks must be a dict, got {type(hw).__name__}")
+    res = hw.get("res")
+    if not isinstance(res, dict):
+        raise ValueError(f"mark.hardware_marks.res must be a dict, got {type(res).__name__}")
+    num_cards = hw.get("num_cards", 1)
+    return list(hardware_marks(res=res, num_cards=num_cards))
+
+
+def resolve_pytest_marks(mark_field: Any) -> list[pytest.MarkDecorator]:
+    """Convert a JSON ``mark`` field into pytest mark decorators.
+
+    Supported form (per test-case object in perf/stability JSON)::
+
+        "mark": [
+            {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
+            "full_model",
+            "diffusion"
+        ]
+
+    Exactly one ``hardware_marks`` object is required when ``mark`` is present.
+    ``hardware_marks`` delegates to :func:`tests.helpers.mark.hardware_marks`
+    (same shape as ``@hardware_test``). Additional array entries are registered
+    pytest marker names (strings).
+    """
+    if mark_field is None:
+        return []
+
+    if isinstance(mark_field, list):
+        marks: list[pytest.MarkDecorator] = []
+        hw_seen = False
+        for item in mark_field:
+            if isinstance(item, dict) and "hardware_marks" in item:
+                if hw_seen:
+                    raise ValueError("mark array must contain at most one hardware_marks object")
+                hw_seen = True
+                marks.extend(_hardware_marks_from_dict(item["hardware_marks"]))
+                unknown_keys = set(item) - {"hardware_marks"}
+                if unknown_keys:
+                    raise ValueError(
+                        f"mark hardware_marks object only allows hardware_marks; unknown keys: {sorted(unknown_keys)}"
+                    )
+            elif isinstance(item, str):
+                item = item.strip()
+                if not item:
+                    raise ValueError("mark name must be a non-empty string")
+                marks.extend(_named_pytest_marks([item]))
+            else:
+                raise ValueError(
+                    f"mark array entries must be hardware_marks objects or marker name strings; got {type(item).__name__}"
+                )
+        if not hw_seen:
+            raise ValueError("mark array must contain a hardware_marks object")
+        return marks
+
+    raise ValueError(f"mark must be a list; got {type(mark_field).__name__}")
+
+
+def _mark_names(mark_field: Any) -> set[str]:
+    if isinstance(mark_field, list):
+        return {str(item) for item in mark_field if isinstance(item, str)}
+    return set()
+
+
+def is_diffusion_perf_config(cfg: dict[str, Any]) -> bool:
+    """True for perf JSON cases intended for ``run_diffusion_benchmark.py``."""
+    if cfg.get("server_type") is not None:
+        return True
+    return "diffusion" in _mark_names(cfg.get("mark"))
+
+
+def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]:
+    return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs}
+
+
+def create_unique_server_pytest_params(
+    configs: list[dict[str, Any]],
+    stage_configs_dir: Path,
+) -> list[Any]:
+    """Like :func:`create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks."""
+    marks_by_name = _marks_by_test_name(configs)
+    return [
+        pytest.param(
+            row,
+            marks=marks_by_name.get(row[0], []),
+            id=row[0],
+        )
+        for row in create_unique_server_params(configs, stage_configs_dir)
+    ]
+
+
+def create_benchmark_pytest_params(
+    benchmark_configs: list[dict[str, Any]],
+    server_to_benchmark_mapping: dict[str, dict],
+) -> list[Any]:
+    """Like :func:`create_benchmark_indices`, but wrap each index in ``pytest.param`` with JSON marks."""
+    marks_by_name = _marks_by_test_name(benchmark_configs)
+    params: list[Any] = []
+    seen: set[str] = set()
+    for config in benchmark_configs:
+        test_name = config["test_name"]
+        if test_name in seen:
+            continue
+        seen.add(test_name)
+        params_list = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
+        id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
+        for idx, id_suffix in enumerate(id_suffixes):
+            params.append(
+                pytest.param(
+                    (test_name, idx),
+                    marks=marks_by_name.get(test_name, []),
+                    id=f"{test_name}-{id_suffix}",
+                )
+            )
+    return params
+
+
+def create_paired_benchmark_pytest_params(
+    server_entries: list[tuple[Any, str]],
+    params_by_test_name: dict[str, list[Any]],
+    marks_by_name: dict[str, list[pytest.MarkDecorator]],
+) -> list[Any]:
+    """One ``pytest.param`` per ``(server entry, benchmark index)`` pair.
+
+    Pass ``server_entry`` and ``(test_name, idx)`` as separate ``pytest.param``
+    arguments for ``@pytest.mark.parametrize("server_fixture,benchmark_params", ...)``.
+    """
+    pairs: list[Any] = []
+    for server_entry, test_name in server_entries:
+        params_list = params_by_test_name.get(test_name, [])
+        id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
+        for idx, id_suffix in enumerate(id_suffixes):
+            pairs.append(
+                pytest.param(
+                    server_entry,
+                    (test_name, idx),
+                    marks=marks_by_name.get(test_name, []),
+                    id=f"{test_name}-{id_suffix}",
+                )
+            )
+    return pairs
+
+
+def create_paired_omni_benchmark_pytest_params(
+    configs: list[dict[str, Any]],
+    stage_configs_dir: Path,
+) -> list[Any]:
+    """Paired params for ``run_benchmark.py`` (omni/tts)."""
+    mapping = create_test_parameter_mapping(configs)
+    marks_by_name = _marks_by_test_name(configs)
+    server_entries = [(row, row[0]) for row in create_unique_server_params(configs, stage_configs_dir)]
+    params_by_test_name = {
+        test_name: get_benchmark_params_for_server(test_name, mapping) for _, test_name in server_entries
+    }
+    return create_paired_benchmark_pytest_params(server_entries, params_by_test_name, marks_by_name)
+
+
 def load_configs(config_path: str) -> list[dict[str, Any]]:
     try:
         abs_path = Path(config_path).resolve()
@@ -32,6 +201,29 @@ def load_configs(config_path: str) -> list[dict[str, Any]]:
         raise RuntimeError(f"Failed to load configuration file: {str(e)}")
 
 
+def load_benchmark_configs(
+    config_path: str | None = None,
+    *,
+    config_dir: Path | None = None,
+) -> list[dict[str, Any]]:
+    """Load one benchmark JSON file, or merge all ``*.json`` under *config_dir*.
+
+    When *config_path* is omitted, every ``*.json`` in *config_dir* is loaded and
+    concatenated. Pytest ``-m`` expressions then select cases via each entry's
+    ``mark`` field (e.g. ``-m tts`` after bulk load).
+    """
+    if config_path is not None:
+        return load_configs(config_path)
+    if config_dir is None:
+        raise ValueError("load_benchmark_configs requires config_path or config_dir")
+    configs: list[dict[str, Any]] = []
+    for path in sorted(config_dir.glob("*.json")):
+        configs.extend(load_configs(str(path)))
+    if not configs:
+        raise ValueError(f"No benchmark JSON files found under {config_dir}")
+    return configs
+
+
 def modify_stage(default_path: str, updates: dict[str, Any] | None, deletes: dict[str, Any] | None) -> str:
     kwargs: dict[str, Any] = {}
     if updates is not None:
@@ -293,6 +485,158 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
     return s if s else default
 
 
+def _benchmark_param_id_suffix(param: dict[str, Any], *, idx: int) -> str:
+    """Derive a readable pytest id suffix from one ``benchmark_params`` entry."""
+    name = param.get("name")
+    if isinstance(name, str) and name.strip():
+        return _safe_filename_token(name.strip())
+
+    parts: list[str] = []
+    for key in ("task", "eval_phase", "dataset_name", "dataset"):
+        value = param.get(key)
+        if value is None or value == "":
+            continue
+        token = _safe_filename_token(str(value))
+        if token != "na":
+            parts.append(token)
+    if parts:
+        return "_".join(parts)
+
+    return f"case{idx}"
+
+
+def _unique_benchmark_param_id_suffixes(params_list: list[dict[str, Any]]) -> list[str]:
+    """Return unique pytest id suffixes for a server's benchmark param list."""
+    raw = [_benchmark_param_id_suffix(param, idx=idx) for idx, param in enumerate(params_list)]
+    seen: dict[str, int] = {}
+    unique: list[str] = []
+    for suffix in raw:
+        count = seen.get(suffix, 0)
+        seen[suffix] = count + 1
+        if count == 0:
+            unique.append(suffix)
+        else:
+            unique.append(f"{suffix}_{count}")
+    return unique
+
+
+def extract_mark_resource_label(mark_field: Any) -> str:
+    """Return a filename-safe hardware label from ``mark.hardware_marks.res`` values.
+
+    Example: ``{"cuda": "H100"}`` -> ``"H100"``; multiple platforms join with ``-``.
+
+    Prefer :func:`get_runtime_resource_label` for perf result filenames so labels
+    reflect the machine that actually ran the benchmark.
+    """
+    if isinstance(mark_field, list):
+        for item in mark_field:
+            if isinstance(item, dict) and "hardware_marks" in item:
+                return extract_mark_resource_label(item)
+        return "na"
+    if not isinstance(mark_field, dict):
+        return "na"
+    hw = mark_field.get("hardware_marks")
+    if not isinstance(hw, dict):
+        return "na"
+    res = hw.get("res")
+    if not isinstance(res, dict) or not res:
+        return "na"
+    labels = [_safe_filename_token(value) for value in res.values()]
+    return "-".join(labels) if labels else "na"
+
+
+_KNOWN_RUNTIME_RESOURCE_TOKENS: tuple[str, ...] = (
+    "H100",
+    "H800",
+    "H200",
+    "H20",
+    "L40S",
+    "L40",
+    "L4",
+    "A100",
+    "A800",
+    "A10G",
+    "A10",
+    "A30",
+    "MI325",
+    "MI300",
+    "MI250",
+    "B60",
+    "S5000",
+    "910B4",
+    "910B",
+    "910",
+    "310P",
+    "A2",
+    "A3",
+)
+_RUNTIME_RESOURCE_LABEL: str | None = None
+
+
+def _normalize_runtime_device_label(raw: str) -> str:
+    """Map a platform device name to a short filename-safe resource token."""
+    if not raw or not str(raw).strip():
+        return "na"
+    upper = str(raw).upper()
+    for token in _KNOWN_RUNTIME_RESOURCE_TOKENS:
+        if token.upper() in upper:
+            return _safe_filename_token(token)
+    compact = re.sub(r"[^a-zA-Z0-9]+", "", str(raw))
+    for prefix in ("NVIDIA", "AMD", "ASCEND", "HUAWEI"):
+        if compact.upper().startswith(prefix):
+            compact = compact[len(prefix) :]
+            break
+    return _safe_filename_token(compact[:48]) if compact else "na"
+
+
+def _read_runtime_device_name(*, device_id: int = 0) -> str | None:
+    """Device name from the active Omni platform."""
+    if current_omni_platform.device_count() <= device_id:
+        return None
+    get_name = getattr(current_omni_platform, "get_device_name", None)
+    if not callable(get_name):
+        return None
+    raw = get_name(device_id)
+    if isinstance(raw, str) and raw.strip():
+        return raw.strip()
+    return None
+
+
+def get_runtime_resource_label(*, device_id: int = 0, refresh: bool = False) -> str:
+    """Return a filename-safe hardware label detected on the running machine."""
+    global _RUNTIME_RESOURCE_LABEL
+    if not refresh and _RUNTIME_RESOURCE_LABEL is not None:
+        return _RUNTIME_RESOURCE_LABEL
+    raw = _read_runtime_device_name(device_id=device_id)
+    label = _normalize_runtime_device_label(raw) if raw else "na"
+    if not refresh:
+        _RUNTIME_RESOURCE_LABEL = label
+    return label
+
+
+_FILENAME_OMIT_RESOURCE_LABELS = frozenset({"H100"})
+
+
+def hardware_json_value(resource_label: str | None) -> str:
+    """Hardware token stored in perf result JSON (empty when unknown)."""
+    token = _safe_filename_token(resource_label)
+    return "" if token == "na" else token
+
+
+def resource_label_for_filename(resource_label: str | None) -> str:
+    """Hardware token embedded in result filenames (H100 omitted on default CI pool)."""
+    token = _safe_filename_token(resource_label)
+    if token in _FILENAME_OMIT_RESOURCE_LABELS:
+        return ""
+    return token
+
+
+def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str:
+    """Return runtime hardware label for perf result filenames."""
+    del configs
+    return get_runtime_resource_label()
+
+
 def resolve_baseline_value(
     baseline_raw: Any,
     *,
@@ -357,6 +701,7 @@ def run_benchmark(
     max_concurrency: Any | None = None,
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
+    resource_label: str | None = None,
 ) -> dict[str, Any]:
     """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
 
@@ -367,7 +712,11 @@ def run_benchmark(
     current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
     ri = _safe_filename_token(random_input_len)
     ro = _safe_filename_token(random_output_len)
-    result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
+    hw = resource_label_for_filename(resource_label)
+    if hw:
+        result_filename = f"result_{test_name}_{hw}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
+    else:
+        result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
     if "--result-filename" in args:
         print(f"The result file will be overwritten by {result_filename}")
     command = (
@@ -437,6 +786,7 @@ def run_benchmark(
         result["random_input_len"] = random_input_len
     if random_output_len is not None:
         result["random_output_len"] = random_output_len
+    result["Hardware"] = hardware_json_value(resource_label)
     with open(result_path, "w", encoding="utf-8") as f:
         json.dump(result, f, ensure_ascii=False, indent=2)
     return result
diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py
index 7ca376323..21167cb07 100644
--- a/tests/dfx/perf/scripts/run_benchmark.py
+++ b/tests/dfx/perf/scripts/run_benchmark.py
@@ -7,21 +7,24 @@ from typing import Any
 import pytest
 
 from tests.dfx.conftest import (
-    create_benchmark_indices,
+    create_paired_omni_benchmark_pytest_params,
     create_test_parameter_mapping,
-    create_unique_server_params,
     get_benchmark_params_for_server,
-    load_configs,
+    get_runtime_resource_label,
+    is_diffusion_perf_config,
+    load_benchmark_configs,
     resolve_baseline_value,
     run_benchmark,
 )
 from tests.helpers.runtime import OmniServer
 
-pytestmark = [pytest.mark.full_model]
-
 # Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline``
 # (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are
 # defined in the same module.
+#
+# Optional JSON field ``mark`` is applied as pytest marks via
+# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
+# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).
 
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -40,22 +43,22 @@ def _get_config_file_from_argv() -> str | None:
 
 
 _PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"
-_DEFAULT_CONFIG_FILE = str(_PERF_TESTS_DIR / "test_qwen_omni.json")
 
 CONFIG_FILE_PATH = _get_config_file_from_argv()
 if CONFIG_FILE_PATH is None:
+    _all_configs = load_benchmark_configs(config_dir=_PERF_TESTS_DIR)
+    BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if not is_diffusion_perf_config(cfg)]
     print(
-        "No --test-config-file in argv, using default: tests/dfx/perf/tests/test_qwen_omni.json "
-        "(override with e.g. --test-config-file tests/dfx/perf/tests/test_tts.json)"
+        f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} omni/tts case(s) from "
+        f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} diffusion; "
+        f"use -m to filter, e.g. -m tts)"
     )
-    CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE
-
-BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-
+else:
+    BENCHMARK_CONFIGS = load_benchmark_configs(CONFIG_FILE_PATH)
 
 DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy"
-test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR)
 server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS)
+paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR)
 
 _omni_server_lock = threading.Lock()
 
@@ -91,17 +94,11 @@ def omni_server(request):
         print("OmniServer stopped")
 
 
-benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
-
-
-@pytest.fixture(params=benchmark_indices)
-def benchmark_params(request, omni_server):
-    """Benchmark parameters fixture with proper parametrization"""
+@pytest.fixture
+def benchmark_params(request):
+    """Benchmark parameters fixture; paired with ``omni_server`` via parametrization."""
     test_name, param_index = request.param
 
-    if test_name != omni_server.test_name:
-        pytest.skip(f"Skipping parameter for {test_name} - current server is {omni_server.test_name}")
-
     all_params = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
 
     if not all_params:
@@ -153,8 +150,11 @@ def assert_result(
 
 
 @pytest.mark.benchmark
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
+@pytest.mark.parametrize(
+    "omni_server,benchmark_params",
+    paired_benchmark_params,
+    indirect=["omni_server", "benchmark_params"],
+)
 def test_performance_benchmark(omni_server, benchmark_params, request):
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
@@ -168,6 +168,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
     print(f"Benchmark parameters: {benchmark_params}")
 
     assert_baseline = request.config.getoption("--assert-baseline", default=False)
+    resource_label = get_runtime_resource_label()
 
     def to_list(value, default=None):
         if value is None:
@@ -239,6 +240,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
             max_concurrency=None,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
+            resource_label=resource_label,
         )
         assert_result(
             result,
@@ -264,6 +266,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
             max_concurrency=concurrency,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
+            resource_label=resource_label,
         )
         assert_result(
             result,
diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py
index 98462743d..d19c0900e 100644
--- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py
+++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py
@@ -8,15 +8,21 @@ This runner separates two concepts:
 2. ``benchmark_endpoint``: which serving API the benchmark client calls.
    Examples: ``/v1/chat/completions`` and ``/v1/videos``.
 
-A config JSON file is REQUIRED via --test-config-file:
+A config JSON file may be passed via --test-config-file. If omitted, every ``*.json`` under
+``tests/dfx/perf/tests/`` is loaded and pytest ``-m`` filters by each case's ``mark``:
+  pytest run_diffusion_benchmark.py -m "diffusion"
   pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
 
 Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off).
 
-All benchmark results for a session are consolidated into a single JSON file under
-BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable).
-Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params,
-timestamp) together with the raw metrics returned by the benchmark script.
+Optional JSON field ``mark`` is applied as pytest marks on that case via
+``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``).
+
+All benchmark results are written under BENCHMARK_RESULT_DIR (override via the
+DIFFUSION_BENCHMARK_DIR environment variable). Each source JSON file gets one
+aggregated ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` (JSON array
+of all runs from cases in that file). Bulk load without ``--test-config-file`` uses
+the same per-file aggregation; ``-m`` only selects which cases run.
 """
 
 import json
@@ -35,8 +41,15 @@ import psutil
 import pytest
 
 from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+from tests.dfx.conftest import (
+    create_paired_benchmark_pytest_params,
+    get_runtime_resource_label,
+    hardware_json_value,
+    is_diffusion_perf_config,
+    resolve_pytest_marks,
+    resource_label_for_filename,
+)
+from tests.helpers.runtime import get_open_port
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
@@ -126,6 +139,10 @@ _BRANCHPOINT_COMMIT_SHA: str | None = None
 DIFFUSION_RESULT_TEMPLATE_PATH = Path(__file__).parent / "diffusion_result_template.json"
 
 
+_DIFFUSION_SOURCE_CONFIG_KEY = "_source_config_file"
+_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"
+
+
 def _get_config_file_from_argv() -> str | None:
     """Read --test-config-file from sys.argv at import time so pytest parametrize can use it.
 
@@ -142,9 +159,6 @@ def _get_config_file_from_argv() -> str | None:
 
 
 CONFIG_FILE_PATH = _get_config_file_from_argv()
-if CONFIG_FILE_PATH is None:
-    print("No config file provided, using default config file: tests/dfx/perf/tests/test_qwen_image_vllm_omni.json")
-    CONFIG_FILE_PATH = "tests/dfx/perf/tests/test_qwen_image_vllm_omni.json"
 
 # ---------------------------------------------------------------------------
 # Config loading
@@ -184,28 +198,85 @@ def load_configs(config_path: str) -> list[dict[str, Any]]:
         raise RuntimeError(f"Failed to load configuration file: {str(e)}")
 
 
-BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
-
-_config_stem = Path(CONFIG_FILE_PATH).stem  # e.g. "test_qwen_image_vllm_omni"
-AGGREGATED_RESULT_FILE = BENCHMARK_RESULT_DIR / f"diffusion_result_{_config_stem}_{_SESSION_TIMESTAMP}.json"
+def load_diffusion_benchmark_configs(
+    config_path: str | None = None,
+    *,
+    config_dir: Path | None = None,
+) -> list[dict[str, Any]]:
+    """Load one diffusion benchmark JSON, or merge all ``*.json`` under *config_dir*."""
+    if config_path is not None:
+        configs = load_configs(config_path)
+        source = str(Path(config_path).resolve())
+        for cfg in configs:
+            cfg.setdefault(_DIFFUSION_SOURCE_CONFIG_KEY, source)
+        return configs
+    if config_dir is None:
+        raise ValueError("load_diffusion_benchmark_configs requires config_path or config_dir")
+    configs: list[dict[str, Any]] = []
+    for path in sorted(config_dir.glob("*.json")):
+        source = str(path.resolve())
+        for cfg in load_configs(str(path)):
+            cfg[_DIFFUSION_SOURCE_CONFIG_KEY] = source
+            configs.append(cfg)
+    if not configs:
+        raise ValueError(f"No benchmark JSON files found under {config_dir}")
+    return configs
 
 
-def _append_to_aggregated_file(record: dict[str, Any]) -> None:
-    """Thread-safe append of *record* to the session-level aggregated JSON file.
+if CONFIG_FILE_PATH is None:
+    _all_configs = load_diffusion_benchmark_configs(config_dir=_PERF_TESTS_DIR)
+    BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if is_diffusion_perf_config(cfg)]
+    print(
+        f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} diffusion case(s) from "
+        f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} omni/tts; "
+        f"use -m to filter, e.g. -m diffusion)"
+    )
+else:
+    BENCHMARK_CONFIGS = load_diffusion_benchmark_configs(CONFIG_FILE_PATH)
 
-    The file contains a JSON array; each call loads the existing array (or
-    starts a new one), appends the record, and writes the file back atomically.
-    """
+_AGGREGATED_RESULT_FILES_BY_SOURCE: dict[str, Path] = {}
+
+
+def _normalized_source_path(source_file: str) -> str:
+    return str(Path(source_file).resolve())
+
+
+def _aggregated_result_file_for_source(source_file: str) -> Path:
+    """One session aggregate per source JSON (same naming as single ``--test-config-file``)."""
+    key = _normalized_source_path(source_file)
+    if key not in _AGGREGATED_RESULT_FILES_BY_SOURCE:
+        stem = Path(key).stem
+        resource = resource_label_for_filename(get_runtime_resource_label())
+        if resource:
+            result_name = f"diffusion_result_{stem}_{resource}_{_SESSION_TIMESTAMP}.json"
+        else:
+            result_name = f"diffusion_result_{stem}_{_SESSION_TIMESTAMP}.json"
+        _AGGREGATED_RESULT_FILES_BY_SOURCE[key] = BENCHMARK_RESULT_DIR / result_name
+    return _AGGREGATED_RESULT_FILES_BY_SOURCE[key]
+
+
+def _write_result_record(record: dict[str, Any]) -> Path:
+    """Append one benchmark record to the aggregate file for its source JSON."""
+    source_file = record.get("source_file") or CONFIG_FILE_PATH
+    if not source_file:
+        raise ValueError("benchmark record missing source_file")
+    target = _aggregated_result_file_for_source(str(source_file))
     with _RESULT_LOCK:
         BENCHMARK_RESULT_DIR.mkdir(parents=True, exist_ok=True)
-        if AGGREGATED_RESULT_FILE.exists():
-            with open(AGGREGATED_RESULT_FILE, encoding="utf-8") as f:
+        if target.exists():
+            with open(target, encoding="utf-8") as f:
                 records: list[dict] = json.load(f)
         else:
             records = []
         records.append(record)
-        with open(AGGREGATED_RESULT_FILE, "w", encoding="utf-8") as f:
+        with open(target, "w", encoding="utf-8") as f:
             json.dump(records, f, indent=2, ensure_ascii=False)
+    return target
+
+
+def _append_to_aggregated_file(record: dict[str, Any]) -> None:
+    """Backward-compatible wrapper; prefer :func:`_write_result_record`."""
+    _write_result_record(record)
 
 
 _server_lock = threading.Lock()
@@ -215,14 +286,6 @@ _server_lock = threading.Lock()
 # ---------------------------------------------------------------------------
 
 
-def _get_open_port() -> int:
-    """Return an available TCP port on localhost."""
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("", 0))
-        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        return s.getsockname()[1]
-
-
 def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None:
     """Block until the given host:port accepts connections or timeout expires.
 
@@ -312,7 +375,7 @@ class DiffusionServer:
         self.model = server_cfg["model"]
         self.serve_args = server_cfg["serve_args"]
         self.host = "127.0.0.1"
-        self.port = port if port is not None else _get_open_port()
+        self.port = port if port is not None else get_open_port(self.host)
         self.proc: subprocess.Popen | None = None
         self.test_name: str = ""
 
@@ -482,6 +545,7 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]
                 "serve_args": _build_serve_args(serve_args_dict),
                 "benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")),
                 "server_params": cfg["server_params"],
+                "mark": cfg.get("mark"),
             }
         )
     return result
@@ -496,19 +560,29 @@ def _test_param_mapping(configs: list[dict[str, Any]]) -> dict[str, list[dict]]:
     return mapping
 
 
-def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer:
-    """Factory: return a vLLM-Omni diffusion server instance for the config."""
-    return DiffusionServer(server_cfg=server_cfg)
+def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]:
+    return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs}
+
+
+def _paired_diffusion_benchmark_pytest_params(configs: list[dict[str, Any]]) -> list[Any]:
+    """Paired params for ``run_diffusion_benchmark.py``; same shape as omni runner."""
+    test_param_map = _test_param_mapping(configs)
+    server_entries = [(cfg, cfg["test_name"]) for cfg in _unique_server_params(configs)]
+    return create_paired_benchmark_pytest_params(server_entries, test_param_map, _marks_by_test_name(configs))
 
 
 # ---------------------------------------------------------------------------
 # Parametrize data
 # ---------------------------------------------------------------------------
 
-server_params = _unique_server_params(BENCHMARK_CONFIGS)
 test_param_map = _test_param_mapping(BENCHMARK_CONFIGS)
+paired_benchmark_params = _paired_diffusion_benchmark_pytest_params(BENCHMARK_CONFIGS)
+
+
+def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer:
+    """Factory: return a vLLM-Omni diffusion server instance for the config."""
+    return DiffusionServer(server_cfg=server_cfg)
 
-benchmark_indices: list[int] = list(range(max(len(v) for v in test_param_map.values())))
 
 # ---------------------------------------------------------------------------
 # Pytest fixtures
@@ -534,16 +608,13 @@ def diffusion_server(request):
 
 
 @pytest.fixture
-def benchmark_params(request, diffusion_server):
-    """Yield the benchmark params dict for the current (server, index) pair."""
-    param_index: int = request.param
-    test_name = diffusion_server.test_name
+def benchmark_params(request):
+    """Benchmark params for the paired server/index parametrization."""
+    test_name, param_index = request.param
 
     params_list = test_param_map.get(test_name, [])
     if not params_list:
         raise ValueError(f"No benchmark params for test: {test_name}")
-    if param_index >= len(params_list):
-        pytest.skip(f"Param index {param_index} out of range for {test_name} (has {len(params_list)} params)")
 
     current = param_index + 1
     total = len(params_list)
@@ -575,8 +646,8 @@ def run_benchmark(
     The raw metrics are written to a temporary file by the subprocess.  After
     the run completes the metrics are merged with full metadata (test_name,
     endpoint, benchmark_params, timestamp, flat reporting fields) and appended
-    to the session-wide aggregated JSON file (AGGREGATED_RESULT_FILE).  The
-    temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
+    to ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` for the
+    source JSON file. The temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
     to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in
     the record.
     """
@@ -586,7 +657,12 @@ def run_benchmark(
     log_dir = BENCHMARK_RESULT_DIR / "logs"
     log_dir.mkdir(parents=True, exist_ok=True)
     endpoint_label = endpoint_filename_token(endpoint)
-    log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
+    resource_label = get_runtime_resource_label()
+    hw_for_filename = resource_label_for_filename(resource_label)
+    if hw_for_filename:
+        log_file = log_dir / f"{test_name}_{hw_for_filename}_{endpoint_label}_{timestamp}.log"
+    else:
+        log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
 
     with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
         tmp_result_file = Path(tmp.name)
@@ -696,7 +772,7 @@ def run_benchmark(
         "Model": model,
         "Framework": server_type,
         "API Endpoint": endpoint,
-        "Hardware": "",
+        "Hardware": hardware_json_value(resource_label),
         "Deployment": "",
         "Task": params.get("task", "t2i"),
         "Dataset": params.get("dataset", "random"),
@@ -728,8 +804,8 @@ def run_benchmark(
         "build_url": os.environ.get("BUILDKITE_BUILD_URL", ""),
         "source_file": source_file,
     }
-    _append_to_aggregated_file(record)
-    print(f"\n  Result appended to: {AGGREGATED_RESULT_FILE}")
+    result_path = _write_result_record(record)
+    print(f"\n  Result saved to: {result_path}")
     print(f"  Log saved to:       {log_file}")
 
     return metrics
@@ -807,7 +883,7 @@ def assert_result(
 
 def _default_benchmark_endpoint_for_task(task: str) -> str:
     """Return the default client-side benchmark endpoint for a diffusion task."""
-    if task in {"t2v", "i2v", "ti2v"}:
+    if task in {"t2v", "i2v", "ti2v", "v2v"}:
         return "/v1/videos"
     if task in {"t2i", "i2i", "ti2i"}:
         return "/v1/chat/completions"
@@ -932,12 +1008,10 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
 # ---------------------------------------------------------------------------
 @pytest.mark.benchmark
 @pytest.mark.parametrize(
-    "diffusion_server",
-    server_params,
-    ids=[p["test_name"] for p in server_params],
-    indirect=True,
+    "diffusion_server,benchmark_params",
+    paired_benchmark_params,
+    indirect=["diffusion_server", "benchmark_params"],
 )
-@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
 def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
     """Run the diffusion performance benchmark and verify request completion.
 
@@ -962,7 +1036,10 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
             test_name=test_name,
             endpoint=endpoint,
             server_cfg=server_cfg,
-            source_file=cast(str, CONFIG_FILE_PATH),
+            source_file=server_cfg.get(
+                _DIFFUSION_SOURCE_CONFIG_KEY,
+                CONFIG_FILE_PATH or f"{_PERF_TESTS_DIR}/*.json",
+            ),
         )
 
         print(f"\n{'=' * 60}")
@@ -980,7 +1057,9 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
             if key in result:
                 print(f"  {key}: {result[key]:.4f}")
 
-        print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
+        source = server_cfg.get(_DIFFUSION_SOURCE_CONFIG_KEY) or CONFIG_FILE_PATH
+        if source:
+            print(f"\n  Aggregated results: {_aggregated_result_file_for_source(str(source))}")
         print("=" * 60)
 
         assert_baseline = request.config.getoption("--assert-baseline", default=False)
diff --git a/tests/dfx/perf/tests/test_bagel_vllm_omni.json b/tests/dfx/perf/tests/test_bagel_vllm_omni.json
index 09fc96030..303839387 100644
--- a/tests/dfx/perf/tests/test_bagel_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_bagel_vllm_omni.json
@@ -1,6 +1,18 @@
 [
   {
     "test_name": "test_bagel_single_device_single_stage_t2i",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
     "description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -38,6 +50,18 @@
   },
   {
     "test_name": "test_bagel_single_device_single_stage_i2i",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
     "description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -76,6 +100,18 @@
   },
   {
     "test_name": "test_bagel_single_device_multi_stage_t2i",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
     "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -113,6 +149,18 @@
   },
   {
     "test_name": "test_bagel_single_device_multi_stage_i2i",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
     "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
     "server_type": "vllm-omni",
     "server_params": {
diff --git a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
new file mode 100644
index 000000000..9bf31658d
--- /dev/null
+++ b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
@@ -0,0 +1,261 @@
+[
+  {
+    "test_name": "test_cosmos3_t2i_official_demo_2gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Cosmos3-Nano text-to-image official demo workload on 2 GPUs",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/images/generations",
+    "server_params": {
+      "model": "nvidia/Cosmos3-Nano",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "use-hsdp": true,
+        "hsdp-shard-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "model-class-name": "Cosmos3OmniDiffusersPipeline",
+        "no-guardrails": true,
+        "stage-init-timeout": 900,
+        "init-timeout": 1200,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1024x1024_steps4",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 4,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "seed": 42,
+        "extra-body": {
+          "guidance_scale": 7.0,
+          "flow_shift": 3.0,
+          "guardrails": false
+        },
+        "baseline": {
+          "throughput_qps": 1.107,
+          "latency_mean": 0.8929,
+          "peak_memory_mb_mean": 80000.0
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_cosmos3_t2v_official_demo_2gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Cosmos3-Nano text-to-video official demo workload on 2 GPUs",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "nvidia/Cosmos3-Nano",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "use-hsdp": true,
+        "hsdp-shard-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "model-class-name": "Cosmos3OmniDiffusersPipeline",
+        "no-guardrails": true,
+        "stage-init-timeout": 900,
+        "init-timeout": 1200,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1280x720_frames189_steps4",
+        "dataset": "random",
+        "task": "t2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1280,
+            "height": 720,
+            "prompt": "A robot arm is cleaning a plate in the kitchen.",
+            "num_inference_steps": 4,
+            "num_frames": 189,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 10.0,
+          "max_sequence_length": 4096,
+          "guardrails": false,
+          "use_resolution_template": false,
+          "use_duration_template": false
+        },
+        "baseline": {
+          "throughput_qps": 0.027,
+          "latency_mean": 36.6824,
+          "peak_memory_mb_mean": 30155.4
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_cosmos3_i2v_official_demo_2gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Cosmos3-Nano image-to-video official demo workload on 2 GPUs",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "nvidia/Cosmos3-Nano",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "use-hsdp": true,
+        "hsdp-shard-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "model-class-name": "Cosmos3OmniDiffusersPipeline",
+        "no-guardrails": true,
+        "stage-init-timeout": 900,
+        "init-timeout": 1200,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1280x720_frames189_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1280,
+            "height": 720,
+            "prompt": "The scene comes to life with smooth, natural motion.",
+            "num_inference_steps": 4,
+            "num_frames": 189,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 10.0,
+          "max_sequence_length": 4096,
+          "guardrails": false
+        },
+        "baseline": {
+          "throughput_qps": 0.027,
+          "latency_mean": 36.968,
+          "peak_memory_mb_mean": 30443.6
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_cosmos3_v2v_official_demo_2gpu",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Cosmos3-Nano video-to-video official demo workload on 2 GPUs",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/videos",
+    "server_params": {
+      "model": "nvidia/Cosmos3-Nano",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "use-hsdp": true,
+        "hsdp-shard-size": 2,
+        "vae-patch-parallel-size": 2,
+        "vae-use-tiling": true,
+        "model-class-name": "Cosmos3OmniDiffusersPipeline",
+        "no-guardrails": true,
+        "stage-init-timeout": 900,
+        "init-timeout": 1200,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1280x720_frames189_steps4",
+        "dataset": "random",
+        "task": "v2v",
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "seed": 42,
+        "random-request-config": [
+          {
+            "width": 1280,
+            "height": 720,
+            "prompt": "Continue the same scene with smooth natural motion.",
+            "num_inference_steps": 4,
+            "num_frames": 189,
+            "fps": 24,
+            "weight": 1
+          }
+        ],
+        "extra-body": {
+          "guidance_scale": 6.0,
+          "flow_shift": 10.0,
+          "max_sequence_length": 4096,
+          "guardrails": false,
+          "condition_frame_indexes_vision": [
+            0,
+            1
+          ],
+          "condition_video_keep": "first"
+        },
+        "baseline": {
+          "throughput_qps": 0.027,
+          "latency_mean": 37.2117,
+          "peak_memory_mb_mean": 30135.6
+        }
+      }
+    ]
+  }
+]
diff --git a/tests/dfx/perf/tests/test_higgs_audio_v3.json b/tests/dfx/perf/tests/test_higgs_audio_v3.json
index 6fdaec580..9b7b2eb5c 100644
--- a/tests/dfx/perf/tests/test_higgs_audio_v3.json
+++ b/tests/dfx/perf/tests/test_higgs_audio_v3.json
@@ -1,6 +1,18 @@
 [
   {
     "test_name": "test_higgs_audio_v3",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "bosonai/higgs-audio-v3-tts-4b",
       "stage_overrides": {
diff --git a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
index 200251ba6..1b6902cba 100644
--- a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
+++ b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json
@@ -1,139 +1,152 @@
 [
-    {
-        "test_name": "test_hunyuan_image3_ti2i_inline",
-        "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
-        "server_type": "vllm-omni",
-        "benchmark_endpoint": "/v1/images/edits",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "deploy-config-inline": {
-                    "type": "yaml",
-                    "content": {
-                        "pipeline": "hunyuan_image_3_moe",
-                        "async_chunk": false,
-                        "trust_remote_code": true,
-                        "connectors": {
-                            "rdma_connector": {
-                                "name": "MooncakeTransferEngineConnector",
-                                "extra": {
-                                    "host": "auto",
-                                    "zmq_port": 50051,
-                                    "protocol": "rdma",
-                                    "device_name": "",
-                                    "memory_pool_size": 17179869184,
-                                    "memory_pool_device": "cpu"
-                                }
-                            }
-                        },
-                        "stages": [
-                            {
-                                "stage_id": 0,
-                                "is_comprehension": true,
-                                "final_output": true,
-                                "final_output_type": "text",
-                                "max_num_seqs": 16,
-                                "gpu_memory_utilization": 0.95,
-                                "enforce_eager": false,
-                                "enable_prefix_caching": false,
-                                "max_num_batched_tokens": 32768,
-                                "devices": "0,1",
-                                "tensor_parallel_size": 2,
-                                "hf_overrides": {
-                                    "rope_parameters": {
-                                        "mrope_section": [
-                                            0,
-                                            32,
-                                            32
-                                        ],
-                                        "rope_type": "default"
-                                    }
-                                },
-                                "omni_kv_config": {
-                                    "need_send_cache": true
-                                },
-                                "output_connectors": {
-                                    "to_stage_1": "rdma_connector"
-                                },
-                                "default_sampling_params": {
-                                    "temperature": 0.0,
-                                    "top_p": 1,
-                                    "top_k": -1,
-                                    "max_tokens": 8192,
-                                    "detokenize": true,
-                                    "skip_special_tokens": false,
-                                    "stop_token_ids": [
-                                        128025
-                                    ]
-                                }
-                            },
-                            {
-                                "stage_id": 1,
-                                "max_num_seqs": 1,
-                                "enforce_eager": true,
-                                "devices": "2,3",
-                                "distributed_executor_backend": "mp",
-                                "omni_kv_config": {
-                                    "need_recv_cache": true
-                                },
-                                "parallel_config": {
-                                    "tensor_parallel_size": 2,
-                                    "vae_patch_parallel_size": 2,
-                                    "enable_expert_parallel": true
-                                },
-                                "input_connectors": {
-                                    "from_stage_0": "rdma_connector"
-                                },
-                                "default_sampling_params": {
-                                    "num_inference_steps": 50,
-                                    "guidance_scale": 0
-                                }
-                            }
-                        ],
-                        "edges": [
-                            {
-                                "from": 0,
-                                "to": 1,
-                                "window_size": -1,
-                                "max_inflight": 1
-                            }
-                        ]
-                    }
-                },
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "inline_dataset_5_requests",
-                "dataset": "custom",
-                "dataset-path-inline": {
-                    "type": "jsonl",
-                    "content": [
-                        {
-                            "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
-                            "image_urls": [
-                                "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"
-                            ],
-                            "num_inference_steps": 8,
-                            "seed": 42,
-                            "guidance_scale": 1.0
-                        }
-                    ]
-                },
-                "task": "ti2i",
-                "bot-task": "think_recaption",
-                "max-concurrency": 16,
-                "num-prompts": 96,
-                "skip-performance-assertion": true,
-                "baseline": {
-                    "throughput_qps": 0.1874,
-                    "latency_p99": 101.852,
-                    "latency_mean": 79.0385,
-                    "peak_memory_mb_max": 70000,
-                    "peak_memory_mb_mean": 70000
+  {
+    "test_name": "test_hunyuan_image3_ti2i_inline",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "diffusion",
+      "local_model"
+    ],
+    "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
+    "server_type": "vllm-omni",
+    "benchmark_endpoint": "/v1/images/edits",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "trust-remote-code": true,
+        "deploy-config-inline": {
+          "type": "yaml",
+          "content": {
+            "pipeline": "hunyuan_image_3_moe",
+            "async_chunk": false,
+            "trust_remote_code": true,
+            "connectors": {
+              "rdma_connector": {
+                "name": "MooncakeTransferEngineConnector",
+                "extra": {
+                  "host": "auto",
+                  "zmq_port": 50051,
+                  "protocol": "rdma",
+                  "device_name": "",
+                  "memory_pool_size": 17179869184,
+                  "memory_pool_device": "cpu"
                 }
+              }
+            },
+            "stages": [
+              {
+                "stage_id": 0,
+                "is_comprehension": true,
+                "final_output": true,
+                "final_output_type": "text",
+                "max_num_seqs": 16,
+                "gpu_memory_utilization": 0.95,
+                "enforce_eager": false,
+                "enable_prefix_caching": false,
+                "max_num_batched_tokens": 32768,
+                "devices": "0,1",
+                "tensor_parallel_size": 2,
+                "hf_overrides": {
+                  "rope_parameters": {
+                    "mrope_section": [
+                      0,
+                      32,
+                      32
+                    ],
+                    "rope_type": "default"
+                  }
+                },
+                "omni_kv_config": {
+                  "need_send_cache": true
+                },
+                "output_connectors": {
+                  "to_stage_1": "rdma_connector"
+                },
+                "default_sampling_params": {
+                  "temperature": 0.0,
+                  "top_p": 1,
+                  "top_k": -1,
+                  "max_tokens": 8192,
+                  "detokenize": true,
+                  "skip_special_tokens": false,
+                  "stop_token_ids": [
+                    128025
+                  ]
+                }
+              },
+              {
+                "stage_id": 1,
+                "max_num_seqs": 1,
+                "enforce_eager": true,
+                "devices": "2,3",
+                "distributed_executor_backend": "mp",
+                "omni_kv_config": {
+                  "need_recv_cache": true
+                },
+                "parallel_config": {
+                  "tensor_parallel_size": 2,
+                  "vae_patch_parallel_size": 2,
+                  "enable_expert_parallel": true
+                },
+                "input_connectors": {
+                  "from_stage_0": "rdma_connector"
+                },
+                "default_sampling_params": {
+                  "num_inference_steps": 50,
+                  "guidance_scale": 0
+                }
+              }
+            ],
+            "edges": [
+              {
+                "from": 0,
+                "to": 1,
+                "window_size": -1,
+                "max_inflight": 1
+              }
+            ]
+          }
+        },
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "inline_dataset_5_requests",
+        "dataset": "custom",
+        "dataset-path-inline": {
+          "type": "jsonl",
+          "content": [
+            {
+              "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
+              "image_urls": [
+                "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"
+              ],
+              "num_inference_steps": 8,
+              "seed": 42,
+              "guidance_scale": 1.0
             }
-        ]
-    }
+          ]
+        },
+        "task": "ti2i",
+        "bot-task": "think_recaption",
+        "max-concurrency": 16,
+        "num-prompts": 96,
+        "skip-performance-assertion": true,
+        "baseline": {
+          "throughput_qps": 0.1874,
+          "latency_p99": 101.852,
+          "latency_mean": 79.0385,
+          "peak_memory_mb_max": 70000,
+          "peak_memory_mb_mean": 70000
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
index e6d7ce659..7a699f53a 100644
--- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
+++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
@@ -1,35 +1,47 @@
 [
-    {
-        "test_name": "test_hunyuan_image_tp2_cfgp2",
-        "description": "TP=2 CfgP=2 baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-                "tensor-parallel-size": 2,
-                "cfg-parallel-size": 2,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1024x1024_steps8",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1024,
-                "height": 1024,
-                "num-inference-steps": 8,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "skip-performance-assertion": true,
-                "baseline": {
-                    "throughput_qps": 0.21,
-                    "latency_p99": 4.7469,
-                    "peak_memory_mb_max": 101100,
-                    "peak_memory_mb_mean": 101100
-                }
-            }
-        ]
-    }
+  {
+    "test_name": "test_hunyuan_image_tp2_cfgp2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "diffusion",
+      "local_model"
+    ],
+    "description": "TP=2 CfgP=2 baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+        "tensor-parallel-size": 2,
+        "cfg-parallel-size": 2,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1024x1024_steps8",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 8,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "skip-performance-assertion": true,
+        "baseline": {
+          "throughput_qps": 0.21,
+          "latency_p99": 4.7469,
+          "peak_memory_mb_max": 101100,
+          "peak_memory_mb_mean": 101100
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
index fbbe1643e..521febd2d 100644
--- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
+++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
@@ -1,35 +1,47 @@
 [
-    {
-        "test_name": "test_hunyuan_image_tp2_sp2",
-        "description": "TP=2 SP=2 baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-                "tensor-parallel-size": 2,
-                "ulysses-degree": 2,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1024x1024_steps8",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1024,
-                "height": 1024,
-                "num-inference-steps": 8,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "skip-performance-assertion": true,
-                "baseline": {
-                    "throughput_qps": 0.20,
-                    "latency_p99": 5.1025,
-                    "peak_memory_mb_max": 97402,
-                    "peak_memory_mb_mean": 97402
-                }
-            }
-        ]
-    }
+  {
+    "test_name": "test_hunyuan_image_tp2_sp2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "diffusion",
+      "local_model"
+    ],
+    "description": "TP=2 SP=2 baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+        "tensor-parallel-size": 2,
+        "ulysses-degree": 2,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1024x1024_steps8",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 8,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "skip-performance-assertion": true,
+        "baseline": {
+          "throughput_qps": 0.20,
+          "latency_p99": 5.1025,
+          "peak_memory_mb_max": 97402,
+          "peak_memory_mb_mean": 97402
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp4.json b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json
index eb37f7a22..f2d05564b 100644
--- a/tests/dfx/perf/tests/test_hunyuan_image_tp4.json
+++ b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json
@@ -1,35 +1,47 @@
 [
-    {
-        "test_name": "test_hunyuan_image_tp4",
-        "description": "TP=4 baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "tencent/HunyuanImage-3.0-Instruct",
-            "serve_args": {
-                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-                "tensor-parallel-size": 4,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1024x1024_steps8",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1024,
-                "height": 1024,
-                "num-inference-steps": 8,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "skip-performance-assertion": true,
-                "baseline": {
-                    "throughput_qps": 0.213,
-                    "latency_mean": 4.6953,
-                    "latency_p99": 4.8601,
-                    "peak_memory_mb_max": 56912.0,
-                    "peak_memory_mb_mean": 56912.0
-                }
-            }
-        ]
-    }
+  {
+    "test_name": "test_hunyuan_image_tp4",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "diffusion",
+      "local_model"
+    ],
+    "description": "TP=4 baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+        "tensor-parallel-size": 4,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "1024x1024_steps8",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 8,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "skip-performance-assertion": true,
+        "baseline": {
+          "throughput_qps": 0.213,
+          "latency_mean": 4.6953,
+          "latency_p99": 4.8601,
+          "peak_memory_mb_max": 56912.0,
+          "peak_memory_mb_mean": 56912.0
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json
index 4a6f9e350..1a0f31e98 100644
--- a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json
@@ -1,217 +1,273 @@
 [
-    {
-        "test_name": "test_ltx2_baseline_eager",
-        "description": "Single-device baseline with enforce-eager (no torch.compile)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
+  {
+    "test_name": "test_ltx2_baseline_eager",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline with enforce-eager (no torch.compile)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "enforce-eager": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-
-    {
-        "test_name": "test_ltx2_torch_compile",
-        "description": "Single-device with torch.compile (default, no enforce-eager)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
+  {
+    "test_name": "test_ltx2_torch_compile",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device with torch.compile (default, no enforce-eager)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-
-    {
-        "test_name": "test_ltx2_cfg2_eager",
-        "description": "CFG-parallel=2 with enforce-eager",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cfg-parallel-size": 2,
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
+  {
+    "test_name": "test_ltx2_cfg2_eager",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "CFG-parallel=2 with enforce-eager",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "enforce-eager": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-
-    {
-        "test_name": "test_ltx2_cfg2_compile",
-        "description": "CFG-parallel=2 with torch.compile",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cfg-parallel-size": 2,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
+  {
+    "test_name": "test_ltx2_cfg2_compile",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "CFG-parallel=2 with torch.compile",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "cfg-parallel-size": 2,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-
-    {
-        "test_name": "test_ltx2_cache_dit_eager",
-        "description": "CacheDiT with enforce-eager",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cache-backend": "cache_dit",
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
-    }
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  },
+  {
+    "test_name": "test_ltx2_cache_dit_eager",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "CacheDiT with enforce-eager",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Lightricks/LTX-2",
+      "serve_args": {
+        "cache-backend": "cache_dit",
+        "enforce-eager": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "256x256_145f_steps6",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 256,
+        "height": 256,
+        "num-frames": 145,
+        "fps": 24,
+        "num-inference-steps": 6,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      },
+      {
+        "name": "480x768_41f_steps20",
+        "dataset": "random",
+        "task": "t2v",
+        "backend": "v1/videos",
+        "width": 768,
+        "height": 480,
+        "num-frames": 41,
+        "fps": 24,
+        "num-inference-steps": 20,
+        "num-prompts": 3,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
index 459554f63..66ffa231f 100644
--- a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
+++ b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
@@ -1,331 +1,343 @@
 [
-    {
-        "test_name": "test_qwen3_omni_chunk",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": [
-                "--async-chunk"
-            ]
+  {
+    "test_name": "test_qwen3_omni_chunk",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "omni"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "extra_cli_args": [
+        "--async-chunk"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          4,
+          16,
+          32,
+          64,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          4,
+          8,
+          16,
+          32
+        ],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            135.4379,
+            187.0896,
+            400,
+            435.675,
+            724.6836
+          ],
+          "mean_e2el_ms": [
+            17619.9851,
+            34811.8335,
+            43656.9501,
+            56657.1548,
+            85097.6016
+          ],
+          "mean_audio_rtf": [
+            0.1279,
+            0.2238,
+            0.2735,
+            0.3712,
+            0.5634
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          10
+        ],
+        "request_rate": [
+          0.1
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
         },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    4,
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        135.4379,
-                        187.0896,
-                        400,
-                        435.675,
-                        724.6836
-                    ],
-                    "mean_e2el_ms": [
-                        17619.9851,
-                        34811.8335,
-                        43656.9501,
-                        56657.1548,
-                        85097.6016
-                    ],
-                    "mean_audio_rtf": [
-                        0.1279,
-                        0.2238,
-                        0.2735,
-                        0.3712,
-                        0.5634
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    10
-                ],
-                "request_rate": [
-                    0.1
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        340.329
-                    ],
-                    "mean_e2el_ms": [
-                        4305.477
-                    ],
-                    "mean_audio_rtf": [
-                        0.1305
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    40
-                ],
-                "request_rate": [
-                    0.5
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        220.0909
-                    ],
-                    "mean_e2el_ms": [
-                        5475.9565
-                    ],
-                    "mean_audio_rtf": [
-                        0.1818
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    100
-                ],
-                "request_rate": [
-                    1.0
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        450.0155
-                    ],
-                    "mean_e2el_ms": [
-                        8516.8316
-                    ],
-                    "mean_audio_rtf": [
-                        0.2717
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    4,
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        78.4278,
-                        195.8037,
-                        333.7318,
-                        561.805,
-                        1038.043
-                    ],
-                    "mean_e2el_ms": [
-                        4288.3029,
-                        7429.9429,
-                        9451.6184,
-                        12607.5292,
-                        18153.2547
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    10
-                ],
-                "request_rate": [
-                    0.1
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        275.3875
-                    ],
-                    "mean_e2el_ms": [
-                        730.6995
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    40
-                ],
-                "request_rate": [
-                    0.5
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        166.2201
-                    ],
-                    "mean_e2el_ms": [
-                        673.5789
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    100
-                ],
-                "request_rate": [
-                    1.0
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        357.7781
-                    ],
-                    "mean_e2el_ms": [
-                        953.3493
-                    ]
-                }
-            }
-        ]
-    }
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            340.329
+          ],
+          "mean_e2el_ms": [
+            4305.477
+          ],
+          "mean_audio_rtf": [
+            0.1305
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          40
+        ],
+        "request_rate": [
+          0.5
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            220.0909
+          ],
+          "mean_e2el_ms": [
+            5475.9565
+          ],
+          "mean_audio_rtf": [
+            0.1818
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          1.0
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            450.0155
+          ],
+          "mean_e2el_ms": [
+            8516.8316
+          ],
+          "mean_audio_rtf": [
+            0.2717
+          ]
+        }
+      },
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          4,
+          16,
+          32,
+          64,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          4,
+          8,
+          16,
+          32
+        ],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            78.4278,
+            195.8037,
+            333.7318,
+            561.805,
+            1038.043
+          ],
+          "mean_e2el_ms": [
+            4288.3029,
+            7429.9429,
+            9451.6184,
+            12607.5292,
+            18153.2547
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          10
+        ],
+        "request_rate": [
+          0.1
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            275.3875
+          ],
+          "mean_e2el_ms": [
+            730.6995
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          40
+        ],
+        "request_rate": [
+          0.5
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            166.2201
+          ],
+          "mean_e2el_ms": [
+            673.5789
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          1.0
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            357.7781
+          ],
+          "mean_e2el_ms": [
+            953.3493
+          ]
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
index bee799a77..ac6f679c1 100644
--- a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
+++ b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
@@ -1,115 +1,127 @@
 [
-    {
-        "test_name": "test_qwen3_omni_3gpu_replica2",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "stage_overrides": {
-                "1": {"num_replicas": 2, "devices": "1,2"},
-                "2": {"num_replicas": 2, "devices": "1,2"}
-            },
-            "extra_cli_args": [
-                "--async-chunk"
-            ]
+  {
+    "test_name": "test_qwen3_omni_3gpu_replica2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 3
+        }
+      },
+      "full_model",
+      "omni"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "stage_overrides": {
+          "1": {"num_replicas": 2, "devices": "1,2"},
+          "2": {"num_replicas": 2, "devices": "1,2"}
+      },
+      "extra_cli_args": [
+        "--async-chunk"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [128, 128, 128, 128],
+        "max_concurrency": [8, 16, 24, 32],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+            "mean_ttft_ms": [203, 283, 470, 869],
+            "mean_tpot_ms": [28, 36, 41, 45],
+            "mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
+            "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [10],
+        "request_rate": [0.1],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
         },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [128, 128, 128, 128],
-                "max_concurrency": [8, 16, 24, 32],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [203, 283, 470, 869],
-                    "mean_tpot_ms": [28, 36, 41, 45],
-                    "mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
-                    "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [10],
-                "request_rate": [0.1],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [377],
-                    "mean_audio_ttfp_ms": [684],
-                    "mean_audio_rtf": [0.14]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [40],
-                "request_rate": [0.5],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [228],
-                    "mean_audio_ttfp_ms": [686],
-                    "mean_audio_rtf": [0.16]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [100],
-                "request_rate": [1.0],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [522],
-                    "mean_audio_ttfp_ms": [1119],
-                    "mean_audio_rtf": [0.21]
-                }
-            }
-        ]
-    }
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+            "mean_ttft_ms": [377],
+            "mean_audio_ttfp_ms": [684],
+            "mean_audio_rtf": [0.14]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [40],
+        "request_rate": [0.5],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+            "mean_ttft_ms": [228],
+            "mean_audio_ttfp_ms": [686],
+            "mean_audio_rtf": [0.16]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [100],
+        "request_rate": [1.0],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+            "mean_ttft_ms": [522],
+            "mean_audio_ttfp_ms": [1119],
+            "mean_audio_rtf": [0.21]
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
index d01ca06b4..78daf691b 100644
--- a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
+++ b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
@@ -1,170 +1,182 @@
 [
-    {
-        "test_name": "test_qwen3_omni",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": [
-                "--no-async-chunk"
-            ]
+  {
+    "test_name": "test_qwen3_omni",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "omni"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "extra_cli_args": [
+        "--no-async-chunk"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          4,
+          16,
+          32,
+          64,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          4,
+          8,
+          16,
+          32
+        ],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            109.1218,
+            247.4411,
+            354.9384,
+            447.7885,
+            694.7487
+          ],
+          "mean_e2el_ms": [
+            22413.667,
+            30808.9715,
+            35677.844,
+            46345.8978,
+            68350.4853
+          ],
+          "mean_audio_rtf": [
+            0.1695,
+            0.2033,
+            0.2467,
+            0.3234,
+            0.4438
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          10
+        ],
+        "request_rate": [
+          0.1
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
         },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    4,
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        109.1218,
-                        247.4411,
-                        354.9384,
-                        447.7885,
-                        694.7487
-                    ],
-                    "mean_e2el_ms": [
-                        22413.667,
-                        30808.9715,
-                        35677.844,
-                        46345.8978,
-                        68350.4853
-                    ],
-                    "mean_audio_rtf": [
-                        0.1695,
-                        0.2033,
-                        0.2467,
-                        0.3234,
-                        0.4438
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    10
-                ],
-                "request_rate": [
-                    0.1
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        323.3573
-                    ],
-                    "mean_e2el_ms": [
-                        4623.4173
-                    ],
-                    "mean_audio_rtf": [
-                        0.1406
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    40
-                ],
-                "request_rate": [
-                    0.5
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        215.4502
-                    ],
-                    "mean_e2el_ms": [
-                        4924.149
-                    ],
-                    "mean_audio_rtf": [
-                        0.1637
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "openai-chat-omni",
-                "endpoint": "/v1/chat/completions",
-                "num_prompts": [
-                    100
-                ],
-                "request_rate": [
-                    1.0
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        442.5815
-                    ],
-                    "mean_e2el_ms": [
-                        6441.0986
-                    ],
-                    "mean_audio_rtf": [
-                        0.2091
-                    ]
-                }
-            }
-        ]
-    }
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            323.3573
+          ],
+          "mean_e2el_ms": [
+            4623.4173
+          ],
+          "mean_audio_rtf": [
+            0.1406
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          40
+        ],
+        "request_rate": [
+          0.5
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            215.4502
+          ],
+          "mean_e2el_ms": [
+            4924.149
+          ],
+          "mean_audio_rtf": [
+            0.1637
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "openai-chat-omni",
+        "endpoint": "/v1/chat/completions",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          1.0
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+        "baseline": {
+          "mean_ttft_ms": [
+            442.5815
+          ],
+          "mean_e2el_ms": [
+            6441.0986
+          ],
+          "mean_audio_rtf": [
+            0.2091
+          ]
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
index f608eca0e..b9d553098 100644
--- a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
+++ b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
@@ -1,180 +1,192 @@
 [
-    {
-        "test_name": "test_qwen3_omni_vllm_text",
-        "server_params": {
-            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "use_omni": false,
-            "extra_cli_args": [
-                "--no-enable-prefix-caching",
-                "--skip-mm-profiling"
-            ]
+  {
+    "test_name": "test_qwen3_omni_vllm_text",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "omni"
+    ],
+    "server_params": {
+      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+      "use_omni": false,
+      "extra_cli_args": [
+        "--no-enable-prefix-caching",
+        "--skip-mm-profiling"
+      ]
+    },
+    "benchmark_params": [
+      {
+        "dataset_name": "random",
+        "backend": "vllm",
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "endpoint": "/v1/completions",
+        "num_prompts": [
+          4,
+          16,
+          32,
+          64,
+          128
+        ],
+        "max_concurrency": [
+          1,
+          4,
+          8,
+          16,
+          32
+        ],
+        "random_input_len": 2500,
+        "random_output_len": 900,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
         },
-        "benchmark_params": [
-            {
-                "dataset_name": "random",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [
-                    4,
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        69.1276,
-                        255.1882,
-                        310.639,
-                        422.0337,
-                        532.0582
-                    ],
-                    "mean_e2el_ms": [
-                        4382.7211,
-                        7510.3815,
-                        9440.4854,
-                        12553.0186,
-                        16932.6065
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [
-                    10
-                ],
-                "request_rate": [
-                    0.1
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 1,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(0, 60, 3)": 1.0
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        26.3381
-                    ],
-                    "mean_e2el_ms": [
-                        485.2436
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [
-                    40
-                ],
-                "request_rate": [
-                    0.5
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 2,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.5,
-                    "(720, 1280, 2)": 0.5
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        26.4517
-                    ],
-                    "mean_e2el_ms": [
-                        517.2017
-                    ]
-                }
-            },
-            {
-                "dataset_name": "random-mm",
-                "backend": "vllm",
-                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-                "endpoint": "/v1/completions",
-                "num_prompts": [
-                    100
-                ],
-                "request_rate": [
-                    1.0
-                ],
-                "random_input_len": 100,
-                "random_output_len": 100,
-                "random_range_ratio": 0.0,
-                "ignore_eos": true,
-                "extra_body": {
-                    "modalities": [
-                        "text"
-                    ]
-                },
-                "random_mm_base_items_per_request": 3,
-                "random_mm_num_mm_items_range_ratio": 0.5,
-                "random_mm_limit_mm_per_prompt": {
-                    "image": 1,
-                    "video": 1,
-                    "audio": 1
-                },
-                "random_mm_bucket_config": {
-                    "(256, 256, 1)": 0.34,
-                    "(720, 1280, 2)": 0.33,
-                    "(0, 60, 3)": 0.33
-                },
-                "percentile-metrics": "ttft,tpot,itl,e2el",
-                "baseline": {
-                    "mean_ttft_ms": [
-                        27.6972
-                    ],
-                    "mean_e2el_ms": [
-                        561.3684
-                    ]
-                }
-            }
-        ]
-    }
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            69.1276,
+            255.1882,
+            310.639,
+            422.0337,
+            532.0582
+          ],
+          "mean_e2el_ms": [
+            4382.7211,
+            7510.3815,
+            9440.4854,
+            12553.0186,
+            16932.6065
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "vllm",
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "endpoint": "/v1/completions",
+        "num_prompts": [
+          10
+        ],
+        "request_rate": [
+          0.1
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 1,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(0, 60, 3)": 1.0
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            26.3381
+          ],
+          "mean_e2el_ms": [
+            485.2436
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "vllm",
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "endpoint": "/v1/completions",
+        "num_prompts": [
+          40
+        ],
+        "request_rate": [
+          0.5
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 2,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.5,
+          "(720, 1280, 2)": 0.5
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            26.4517
+          ],
+          "mean_e2el_ms": [
+            517.2017
+          ]
+        }
+      },
+      {
+        "dataset_name": "random-mm",
+        "backend": "vllm",
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "endpoint": "/v1/completions",
+        "num_prompts": [
+          100
+        ],
+        "request_rate": [
+          1.0
+        ],
+        "random_input_len": 100,
+        "random_output_len": 100,
+        "random_range_ratio": 0.0,
+        "ignore_eos": true,
+        "extra_body": {
+          "modalities": [
+            "text"
+          ]
+        },
+        "random_mm_base_items_per_request": 3,
+        "random_mm_num_mm_items_range_ratio": 0.5,
+        "random_mm_limit_mm_per_prompt": {
+          "image": 1,
+          "video": 1,
+          "audio": 1
+        },
+        "random_mm_bucket_config": {
+          "(256, 256, 1)": 0.34,
+          "(720, 1280, 2)": 0.33,
+          "(0, 60, 3)": 0.33
+        },
+        "percentile-metrics": "ttft,tpot,itl,e2el",
+        "baseline": {
+          "mean_ttft_ms": [
+            27.6972
+          ],
+          "mean_e2el_ms": [
+            561.3684
+          ]
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
index 118fa6733..53e13daed 100644
--- a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
@@ -1,149 +1,185 @@
 [
-    {
-        "test_name": "test_qwen_image_edit_2511_single_device",
-        "description": "Single-device baseline (two input images)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2511",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0721,
-                    "latency_mean": 13.8398,
-                    "peak_memory_mb_max": 54731.4,
-                    "peak_memory_mb_mean": 54731.4
-                }
-            },
-            {
-                "name": "1536x1536_steps35_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0186,
-                    "latency_mean": 53.6942,
-                    "peak_memory_mb_max": 64402.4,
-                    "peak_memory_mb_mean": 64402.4
-                }
-            }
-        ]
+  {
+    "test_name": "test_qwen_image_edit_2511_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline (two input images)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit-2511",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
-        "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2511",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "vae-patch-parallel-size": 4,
-                "vae-use-tiling": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1536x1536_steps35_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0538,
-                    "latency_mean": 18.5591,
-                    "peak_memory_mb_max": 53485.0,
-                    "peak_memory_mb_mean": 53485.0
-                }
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0721,
+          "latency_mean": 13.8398,
+          "peak_memory_mb_max": 54731.4,
+          "peak_memory_mb_mean": 54731.4
+        }
+      },
+      {
+        "name": "1536x1536_steps35_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0186,
+          "latency_mean": 53.6942,
+          "peak_memory_mb_max": 64402.4,
+          "peak_memory_mb_mean": 64402.4
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit-2511",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
-        "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit-2511",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "cache-backend": "cache_dit",
-                "cache-config": {
-                    "Fn_compute_blocks": 1,
-                    "Bn_compute_blocks": 0,
-                    "max_warmup_steps": 4,
-                    "residual_diff_threshold": 0.24,
-                    "max_continuous_cached_steps": 3,
-                    "enable_taylorseer": false,
-                    "taylorseer_order": 1,
-                    "scm_steps_mask_policy": null,
-                    "scm_steps_policy": "dynamic"
-                },
-                "enable-diffusion-pipeline-profiler": true
-            }
+    "benchmark_params": [
+      {
+        "name": "1536x1536_steps35_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0538,
+          "latency_mean": 18.5591,
+          "peak_memory_mb_max": 53485.0,
+          "peak_memory_mb_mean": 53485.0
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit-2511",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "cache-backend": "cache_dit",
+        "cache-config": {
+          "Fn_compute_blocks": 1,
+          "Bn_compute_blocks": 0,
+          "max_warmup_steps": 4,
+          "residual_diff_threshold": 0.24,
+          "max_continuous_cached_steps": 3,
+          "enable_taylorseer": false,
+          "taylorseer_order": 1,
+          "scm_steps_mask_policy": null,
+          "scm_steps_policy": "dynamic"
         },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.2776,
-                    "latency_mean": 3.5938,
-                    "peak_memory_mb_max": 55214.0,
-                    "peak_memory_mb_mean": 55214.0
-                }
-            },
-            {
-                "name": "1536x1536_steps35_i2i_2img",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 2,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1051,
-                    "latency_mean": 9.4918,
-                    "peak_memory_mb_max": 64265.6,
-                    "peak_memory_mb_mean": 64265.6
-                }
-            }
-        ]
-    }
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.2776,
+          "latency_mean": 3.5938,
+          "peak_memory_mb_max": 55214.0,
+          "peak_memory_mb_mean": 55214.0
+        }
+      },
+      {
+        "name": "1536x1536_steps35_i2i_2img",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 2,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.1051,
+          "latency_mean": 9.4918,
+          "peak_memory_mb_max": 64265.6,
+          "peak_memory_mb_mean": 64265.6
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
index da3288182..bfda30731 100644
--- a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
@@ -1,144 +1,180 @@
 [
-    {
-        "test_name": "test_qwen_image_edit_single_device",
-        "description": "Single-device baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1129,
-                    "latency_mean": 8.8352,
-                    "peak_memory_mb_max": 55356.5,
-                    "peak_memory_mb_mean": 55356.5
-                }
-            },
-            {
-                "name": "1536x1536_steps35_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0240,
-                    "latency_mean": 41.5686,
-                    "peak_memory_mb_max": 62762.7,
-                    "peak_memory_mb_mean": 62762.7
-                }
-            }
-        ]
+  {
+    "test_name": "test_qwen_image_edit_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4",
-        "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "vae-patch-parallel-size": 4,
-                "vae-use-tiling": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1536x1536_steps35_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0672,
-                    "latency_mean": 14.8517,
-                    "peak_memory_mb_max": 53931.5,
-                    "peak_memory_mb_mean": 53931.5
-                }
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.1129,
+          "latency_mean": 8.8352,
+          "peak_memory_mb_max": 55356.5,
+          "peak_memory_mb_mean": 55356.5
+        }
+      },
+      {
+        "name": "1536x1536_steps35_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0240,
+          "latency_mean": 41.5686,
+          "peak_memory_mb_max": 62762.7,
+          "peak_memory_mb_mean": 62762.7
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit",
-        "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Edit",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "cache-backend": "cache_dit",
-                "cache-config": {
-                    "Fn_compute_blocks": 1,
-                    "Bn_compute_blocks": 0,
-                    "max_warmup_steps": 4,
-                    "residual_diff_threshold": 0.24,
-                    "max_continuous_cached_steps": 3,
-                    "enable_taylorseer": false,
-                    "taylorseer_order": 1,
-                    "scm_steps_mask_policy": null,
-                    "scm_steps_policy": "dynamic"
-                },
-                "enable-diffusion-pipeline-profiler": true
-            }
+    "benchmark_params": [
+      {
+        "name": "1536x1536_steps35_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0672,
+          "latency_mean": 14.8517,
+          "peak_memory_mb_max": 53931.5,
+          "peak_memory_mb_mean": 53931.5
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Edit",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "cache-backend": "cache_dit",
+        "cache-config": {
+          "Fn_compute_blocks": 1,
+          "Bn_compute_blocks": 0,
+          "max_warmup_steps": 4,
+          "residual_diff_threshold": 0.24,
+          "max_continuous_cached_steps": 3,
+          "enable_taylorseer": false,
+          "taylorseer_order": 1,
+          "scm_steps_mask_policy": null,
+          "scm_steps_policy": "dynamic"
         },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.2879,
-                    "latency_mean": 3.4701,
-                    "peak_memory_mb_max": 55525.6,
-                    "peak_memory_mb_mean": 55525.6
-                }
-            },
-            {
-                "name": "1536x1536_steps35_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1113,
-                    "latency_mean": 8.9656,
-                    "peak_memory_mb_max": 63758.3,
-                    "peak_memory_mb_mean": 63758.3
-                }
-            }
-        ]
-    }
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.2879,
+          "latency_mean": 3.4701,
+          "peak_memory_mb_max": 55525.6,
+          "peak_memory_mb_mean": 55525.6
+        }
+      },
+      {
+        "name": "1536x1536_steps35_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.1113,
+          "latency_mean": 8.9656,
+          "peak_memory_mb_max": 63758.3,
+          "peak_memory_mb_mean": 63758.3
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
index a94db56f5..3445dd94c 100644
--- a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
@@ -1,49 +1,61 @@
 [
-    {
-        "test_name": "test_qwen_image_layered_single_device",
-        "description": "Single-device baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image-Layered",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "640x640_steps20_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 640,
-                "height": 640,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0703,
-                    "latency_mean": 14.208,
-                    "peak_memory_mb_max": 60543.5,
-                    "peak_memory_mb_mean": 60543.5
-                }
-            },
-            {
-                "name": "1024x1024_steps35_i2i",
-                "dataset": "random",
-                "task": "i2i",
-                "width": 1024,
-                "height": 1024,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0417,
-                    "latency_mean": 23.914,
-                    "peak_memory_mb_max": 60543.5,
-                    "peak_memory_mb_mean": 60543.5
-                }
-            }
-        ]
-    }
+  {
+    "test_name": "test_qwen_image_layered_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image-Layered",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "640x640_steps20_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 640,
+        "height": 640,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0703,
+          "latency_mean": 14.208,
+          "peak_memory_mb_max": 60543.5,
+          "peak_memory_mb_mean": 60543.5
+        }
+      },
+      {
+        "name": "1024x1024_steps35_i2i",
+        "dataset": "random",
+        "task": "i2i",
+        "width": 1024,
+        "height": 1024,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0417,
+          "latency_mean": 23.914,
+          "peak_memory_mb_max": 60543.5,
+          "peak_memory_mb_mean": 60543.5
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
index 5ca6ed61a..b3daf7587 100644
--- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
@@ -1,219 +1,279 @@
 [
-    {
-        "test_name": "test_qwen_image_single_device",
-        "description": "Single-device baseline (no parallelism)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.4153,
-                    "latency_mean": 2.4065,
-                    "peak_memory_mb_mean": 52289.9
-                }
-            },
-            {
-                "name": "1536x1536_steps35",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0423,
-                    "latency_mean": 23.5991,
-                    "peak_memory_mb_mean": 61360.5
-                }
-            }
-        ]
+  {
+    "test_name": "test_qwen_image_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline (no parallelism)",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_single_device_step_execution",
-        "description": "Single-device baseline (no parallelism) with step execution",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true,
-                "step-execution": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.4175,
-                    "latency_mean": 2.3953,
-                    "peak_memory_mb_mean": 56000
-                }
-            },
-            {
-                "name": "1536x1536_steps35",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.0423,
-                    "latency_mean": 23.5768,
-                    "peak_memory_mb_mean": 65000
-                }
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.4153,
+          "latency_mean": 2.4065,
+          "peak_memory_mb_mean": 52289.9
+        }
+      },
+      {
+        "name": "1536x1536_steps35",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0423,
+          "latency_mean": 23.5991,
+          "peak_memory_mb_mean": 61360.5
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_single_device_step_execution",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline (no parallelism) with step execution",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "step-execution": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
-        "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "vae-patch-parallel-size": 4,
-                "vae-use-tiling": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "1536x1536_steps35",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1214,
-                    "latency_mean": 8.2211,
-                    "peak_memory_mb_mean": 51704.7
-                }
-            }
-        ]
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.4175,
+          "latency_mean": 2.3953,
+          "peak_memory_mb_mean": 56000
+        }
+      },
+      {
+        "name": "1536x1536_steps35",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.0423,
+          "latency_mean": 23.5768,
+          "peak_memory_mb_mean": 65000
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit",
-        "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "ulysses-degree": 2,
-                "cfg-parallel-size": 2,
-                "cache-backend": "cache_dit",
-                "cache-config": {
-                    "Fn_compute_blocks": 1,
-                    "Bn_compute_blocks": 0,
-                    "max_warmup_steps": 4,
-                    "residual_diff_threshold": 0.24,
-                    "max_continuous_cached_steps": 3,
-                    "enable_taylorseer": false,
-                    "taylorseer_order": 1,
-                    "scm_steps_mask_policy": null,
-                    "scm_steps_policy": "dynamic"
-                },
-                "enable-diffusion-pipeline-profiler": true
-            }
+    "benchmark_params": [
+      {
+        "name": "1536x1536_steps35",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.1214,
+          "latency_mean": 8.2211,
+          "peak_memory_mb_mean": 51704.7
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 4
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "ulysses-degree": 2,
+        "cfg-parallel-size": 2,
+        "cache-backend": "cache_dit",
+        "cache-config": {
+          "Fn_compute_blocks": 1,
+          "Bn_compute_blocks": 0,
+          "max_warmup_steps": 4,
+          "residual_diff_threshold": 0.24,
+          "max_continuous_cached_steps": 3,
+          "enable_taylorseer": false,
+          "taylorseer_order": 1,
+          "scm_steps_mask_policy": null,
+          "scm_steps_policy": "dynamic"
         },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.559,
-                    "latency_mean": 1.794,
-                    "peak_memory_mb_mean": 52397.6775
-                }
-            },
-            {
-                "name": "1536x1536_steps35",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 1536,
-                "height": 1536,
-                "num-inference-steps": 35,
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": 0.1918,
-                    "latency_mean": 5.2022,
-                    "peak_memory_mb_mean": 61468.8
-                }
-            }
-        ]
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
-        "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Qwen/Qwen-Image",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true,
-                "step-execution": true,
-                "max-num-seqs": 8
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "512x512_steps20_high_concurrency",
-                "dataset": "random",
-                "task": "t2i",
-                "width": 512,
-                "height": 512,
-                "num-inference-steps": 20,
-                "num-prompts": [20, 40, 80, 160],
-                "max-concurrency": [1, 2, 4, 8],
-                "warmup-requests": 8,
-                "warmup-concurrency": 8,
-                "warmup-num-inference-steps": 20,
-                "enable-negative-prompt": true,
-                "baseline": {
-                    "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
-                    "latency_mean": [2.4215, 2.992, 5.9293, 9.656],
-                    "peak_memory_mb_mean": [67000, 67000, 67000, 67000]
-                }
-            }
-        ]
-    }
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.559,
+          "latency_mean": 1.794,
+          "peak_memory_mb_mean": 52397.6775
+        }
+      },
+      {
+        "name": "1536x1536_steps35",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 1536,
+        "height": 1536,
+        "num-inference-steps": 35,
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "enable-negative-prompt": true,
+        "baseline": {
+          "throughput_qps": 0.1918,
+          "latency_mean": 5.2022,
+          "peak_memory_mb_mean": 61468.8
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Qwen/Qwen-Image",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true,
+        "step-execution": true,
+        "max-num-seqs": 8
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "512x512_steps20_high_concurrency",
+        "dataset": "random",
+        "task": "t2i",
+        "width": 512,
+        "height": 512,
+        "num-inference-steps": 20,
+        "num-prompts": [20, 40, 80, 160],
+        "max-concurrency": [1, 2, 4, 8],
+        "warmup-requests": 8,
+        "warmup-concurrency": 8,
+        "warmup-num-inference-steps": 20,
+        "enable-negative-prompt": true,
+        "baseline": {
+            "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
+            "latency_mean": [2.4215, 2.992, 5.9293, 9.656],
+            "peak_memory_mb_mean": [67000, 67000, 67000, 67000]
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py
index a195c4546..71d4aa95b 100644
--- a/tests/dfx/perf/tests/test_runner_metadata.py
+++ b/tests/dfx/perf/tests/test_runner_metadata.py
@@ -86,3 +86,202 @@ def test_enabled_false_entry_is_skipped():
     # Only the enabled=True entry should appear
     assert len(params) == 1
     assert params[0].get("task") == "voice_clone"
+
+
+def test_resolve_pytest_marks_hardware_dict_with_extra():
+    from tests.dfx.conftest import resolve_pytest_marks
+
+    marks = resolve_pytest_marks(
+        [
+            {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
+            "full_model",
+            "diffusion",
+            "local_model",
+        ]
+    )
+    names = {m.name for m in marks}
+    assert "H100" in names
+    assert "cuda" in names
+    assert "gpu" in names
+    assert "distributed_cuda" in names
+    assert "full_model" in names
+    assert "diffusion" in names
+    assert "local_model" in names
+
+
+def test_resolve_pytest_marks_rejects_legacy_object_format():
+    from tests.dfx.conftest import resolve_pytest_marks
+
+    with pytest.raises(ValueError, match="mark must be a list"):
+        resolve_pytest_marks(
+            {
+                "hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1},
+                "marks": ["full_model"],
+            }
+        )
+
+
+def test_extract_mark_resource_label():
+    from tests.dfx.conftest import extract_mark_resource_label
+
+    assert extract_mark_resource_label(None) == "na"
+    assert (
+        extract_mark_resource_label(
+            [
+                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                "full_model",
+            ]
+        )
+        == "H100"
+    )
+    assert (
+        extract_mark_resource_label(
+            [
+                {
+                    "hardware_marks": {
+                        "res": {"cuda": "H100", "rocm": "MI325"},
+                        "num_cards": 2,
+                    }
+                },
+                "diffusion",
+            ]
+        )
+        == "H100-MI325"
+    )
+
+
+def test_resource_label_for_filename():
+    from tests.dfx.conftest import resource_label_for_filename
+
+    assert resource_label_for_filename("H100") == ""
+    assert resource_label_for_filename("L4") == "L4"
+    assert resource_label_for_filename("910B") == "910B"
+    assert resource_label_for_filename("na") == "na"
+
+
+def test_hardware_json_value():
+    from tests.dfx.conftest import hardware_json_value
+
+    assert hardware_json_value("H100") == "H100"
+    assert hardware_json_value("na") == ""
+    assert hardware_json_value(None) == ""
+
+
+def test_extract_configs_resource_label(monkeypatch):
+    from tests.dfx.conftest import extract_configs_resource_label, get_runtime_resource_label
+
+    monkeypatch.setattr(
+        "tests.dfx.conftest._read_runtime_device_name",
+        lambda *, device_id=0: "NVIDIA H100 80GB HBM3",
+    )
+    get_runtime_resource_label(refresh=True)
+    assert extract_configs_resource_label([]) == "H100"
+    get_runtime_resource_label(refresh=True)
+    monkeypatch.setattr(
+        "tests.dfx.conftest._read_runtime_device_name",
+        lambda *, device_id=0: "Ascend910B2",
+    )
+    assert get_runtime_resource_label(refresh=True) == "910B"
+
+
+def test_load_benchmark_configs_from_dir(tmp_path):
+    from tests.dfx.conftest import load_benchmark_configs
+
+    (tmp_path / "a.json").write_text(
+        json.dumps([{"test_name": "test_a", "server_params": {"model": "m/a"}, "benchmark_params": []}]),
+        encoding="utf-8",
+    )
+    (tmp_path / "b.json").write_text(
+        json.dumps([{"test_name": "test_b", "server_params": {"model": "m/b"}, "benchmark_params": []}]),
+        encoding="utf-8",
+    )
+    configs = load_benchmark_configs(config_dir=tmp_path)
+    assert [c["test_name"] for c in configs] == ["test_a", "test_b"]
+
+
+def test_create_unique_server_pytest_params_applies_marks(tmp_path):
+    from tests.dfx.conftest import create_unique_server_pytest_params
+
+    configs = [
+        {
+            "test_name": "test_with_mark",
+            "mark": [
+                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                "full_model",
+            ],
+            "server_params": {"model": "some/model"},
+            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
+        },
+        {
+            "test_name": "test_without_mark",
+            "server_params": {"model": "other/model"},
+            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
+        },
+    ]
+    params = create_unique_server_pytest_params(configs, tmp_path)
+    by_id = {p.id: p for p in params}
+    assert len(by_id["test_with_mark"].values) == 1
+    assert isinstance(by_id["test_with_mark"].values[0], tuple)
+    assert any(m.name == "H100" for m in by_id["test_with_mark"].marks)
+    assert not any(m.name == "H100" for m in by_id["test_without_mark"].marks)
+
+
+def test_is_diffusion_perf_config():
+    from tests.dfx.conftest import is_diffusion_perf_config
+
+    assert not is_diffusion_perf_config(
+        {"test_name": "omni_a", "mark": [{"hardware_marks": {"res": {"cuda": "H100"}}}, "omni"]}
+    )
+    assert is_diffusion_perf_config(
+        {
+            "test_name": "diff_a",
+            "server_type": "vllm-omni",
+            "mark": [{"hardware_marks": {"res": {"cuda": "H100"}}}, "diffusion"],
+        }
+    )
+
+
+def test_benchmark_param_id_suffix_from_task_eval_phase():
+    from tests.dfx.conftest import _unique_benchmark_param_id_suffixes
+
+    params = [
+        {"task": "default_voice", "eval_phase": "latency"},
+        {"task": "default_voice", "eval_phase": "throughput"},
+    ]
+    assert _unique_benchmark_param_id_suffixes(params) == [
+        "default_voice_latency",
+        "default_voice_throughput",
+    ]
+
+
+def test_create_paired_omni_benchmark_pytest_params(tmp_path):
+    from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params
+
+    configs = [
+        {
+            "test_name": "test_omni",
+            "mark": [
+                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                "omni",
+            ],
+            "server_params": {"model": "m/omni"},
+            "benchmark_params": [{"name": "p0", "num_prompts": 1}, {"name": "p1", "num_prompts": 2}],
+        },
+        {
+            "test_name": "test_tts",
+            "mark": [
+                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
+                "tts",
+            ],
+            "server_params": {"model": "m/tts"},
+            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
+        },
+    ]
+    params = create_paired_omni_benchmark_pytest_params(configs, tmp_path)
+    by_id = {p.id: p for p in params}
+    assert set(by_id) == {"test_omni-p0", "test_omni-p1", "test_tts-p0"}
+    omni_row, bench_row = by_id["test_tts-p0"].values
+    assert omni_row[0] == "test_tts"
+    assert bench_row == ("test_tts", 0)
+    assert any(m.name == "tts" for m in by_id["test_tts-p0"].marks)
+    assert not any(m.name == "tts" for m in by_id["test_omni-p0"].marks)
diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json
index 1ba273a2c..85236eab7 100644
--- a/tests/dfx/perf/tests/test_tts.json
+++ b/tests/dfx/perf/tests/test_tts.json
@@ -1,6 +1,19 @@
 [
   {
     "test_name": "test_qwen3_tts_base",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100",
+            "npu": "A3"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
       "stage_overrides": {
@@ -98,6 +111,18 @@
   },
   {
     "test_name": "test_qwen3_tts_customvoice",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
       "stage_overrides": {
@@ -342,6 +367,18 @@
   },
   {
     "test_name": "test_voxtral_tts",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "mistralai/Voxtral-4B-TTS-2603"
     },
diff --git a/tests/dfx/perf/tests/test_voxcpm2.json b/tests/dfx/perf/tests/test_voxcpm2.json
index beb40d5d3..ef72b5d88 100644
--- a/tests/dfx/perf/tests/test_voxcpm2.json
+++ b/tests/dfx/perf/tests/test_voxcpm2.json
@@ -1,9 +1,24 @@
 [
   {
     "test_name": "test_voxcpm2",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "tts"
+    ],
     "server_params": {
       "model": "openbmb/VoxCPM2",
-      "trust_remote_code": true
+      "trust_remote_code": true,
+      "extra_cli_args": [
+        "--trust-remote-code"
+      ]
     },
     "benchmark_params": [
       {
diff --git a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
index 07873a64b..78f720a6a 100644
--- a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
+++ b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
@@ -1,107 +1,131 @@
 [
-    {
-        "test_name": "test_wan22_i2v_single_device",
-        "description": "Single-device baseline",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "832x480_frames81_steps4",
-                "dataset": "random",
-                "task": "i2v",
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 1,
-                "seed": 42,
-                "enable-negative-prompt": true,
-                "random-request-config": [
-                    {
-                        "width": 832,
-                        "height": 480,
-                        "num_inference_steps": 4,
-                        "num_frames": 81,
-                        "fps": 16,
-                        "weight": 1
-                    }
-                ],
-                "baseline": {
-                    "throughput_qps": 0.0465,
-                    "latency_mean": 21.4354,
-                    "peak_memory_mb_mean": 72557.58
-                }
-            }
-        ]
+  {
+    "test_name": "test_wan22_i2v_single_device",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 1
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "Single-device baseline",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+      "serve_args": {
+        "enable-diffusion-pipeline-profiler": true
+      }
     },
-    {
-        "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing",
-        "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-            "serve_args": {
-                "usp": 2,
-                "vae-patch-parallel-size": 2,
-                "use-hsdp": true,
-                "vae-use-slicing": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "832x480_frames81_steps4",
-                "dataset": "random",
-                "task": "i2v",
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 1,
-                "seed": 42,
-                "enable-negative-prompt": true,
-                "random-request-config": [
-                    {
-                        "width": 832,
-                        "height": 480,
-                        "num_inference_steps": 4,
-                        "num_frames": 81,
-                        "fps": 16,
-                        "weight": 1
-                    }
-                ],
-                "baseline": {
-                    "throughput_qps": 0.0642,
-                    "latency_mean": 15.5589,
-                    "peak_memory_mb_mean": 45742.88
-                }
-            },
-            {
-                "name": "1280x720_frames121_steps4",
-                "dataset": "random",
-                "task": "i2v",
-                "num-prompts": 10,
-                "max-concurrency": 1,
-                "num-input-images": 1,
-                "seed": 42,
-                "enable-negative-prompt": true,
-                "random-request-config": [
-                    {
-                        "width": 1280,
-                        "height": 720,
-                        "num_inference_steps": 4,
-                        "num_frames": 121,
-                        "fps": 16,
-                        "weight": 1
-                    }
-                ],
-                "baseline": {
-                    "throughput_qps": 0.0128,
-                    "latency_mean": 77.8565,
-                    "peak_memory_mb_mean": 53947.84
-                }
-            }
-        ]
-    }
+    "benchmark_params": [
+      {
+        "name": "832x480_frames81_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 832,
+            "height": 480,
+            "num_inference_steps": 4,
+            "num_frames": 81,
+            "fps": 16,
+            "weight": 1
+          }
+        ],
+        "baseline": {
+          "throughput_qps": 0.0465,
+          "latency_mean": 21.4354,
+          "peak_memory_mb_mean": 72557.58
+        }
+      }
+    ]
+  },
+  {
+    "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing",
+    "mark": [
+      {
+        "hardware_marks": {
+          "res": {
+            "cuda": "H100"
+          },
+          "num_cards": 2
+        }
+      },
+      "full_model",
+      "diffusion"
+    ],
+    "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing",
+    "server_type": "vllm-omni",
+    "server_params": {
+      "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+      "serve_args": {
+        "usp": 2,
+        "vae-patch-parallel-size": 2,
+        "use-hsdp": true,
+        "vae-use-slicing": true,
+        "enable-diffusion-pipeline-profiler": true
+      }
+    },
+    "benchmark_params": [
+      {
+        "name": "832x480_frames81_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 832,
+            "height": 480,
+            "num_inference_steps": 4,
+            "num_frames": 81,
+            "fps": 16,
+            "weight": 1
+          }
+        ],
+        "baseline": {
+          "throughput_qps": 0.0642,
+          "latency_mean": 15.5589,
+          "peak_memory_mb_mean": 45742.88
+        }
+      },
+      {
+        "name": "1280x720_frames121_steps4",
+        "dataset": "random",
+        "task": "i2v",
+        "num-prompts": 10,
+        "max-concurrency": 1,
+        "num-input-images": 1,
+        "seed": 42,
+        "enable-negative-prompt": true,
+        "random-request-config": [
+          {
+            "width": 1280,
+            "height": 720,
+            "num_inference_steps": 4,
+            "num_frames": 121,
+            "fps": 16,
+            "weight": 1
+          }
+        ],
+        "baseline": {
+          "throughput_qps": 0.0128,
+          "latency_mean": 77.8565,
+          "peak_memory_mb_mean": 53947.84
+        }
+      }
+    ]
+  }
 ]
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py
index fa7f0f905..5ffd25632 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py
+++ b/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",
         ),
diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py
index 6f7cf984d..5e9424c88 100644
--- a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py
+++ b/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},
diff --git a/tests/diffusion/attention/test_kernels_hub.py b/tests/diffusion/attention/test_kernels_hub.py
new file mode 100644
index 000000000..0f22a1244
--- /dev/null
+++ b/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}"
diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py
index b9b4a72b7..babc23b4a 100644
--- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py
+++ b/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."""
diff --git a/tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py b/tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py
new file mode 100644
index 000000000..5faa42eaf
--- /dev/null
+++ b/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,
+    )
diff --git a/tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py b/tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py
new file mode 100644
index 000000000..6def7ebad
--- /dev/null
+++ b/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"
diff --git a/tests/diffusion/distributed/test_expert_parallel_layout.py b/tests/diffusion/distributed/test_expert_parallel_layout.py
index 8fc1db605..c4d3c6ae7 100644
--- a/tests/diffusion/distributed/test_expert_parallel_layout.py
+++ b/tests/diffusion/distributed/test_expert_parallel_layout.py
@@ -19,6 +19,10 @@ class _FakeGroup:
         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)
@@ -55,6 +59,23 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         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,
@@ -68,6 +89,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     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"):
@@ -87,7 +109,7 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         backend="gloo",
     )
 
-    assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP
+    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
@@ -96,6 +118,13 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
     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 == [
@@ -155,7 +184,9 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch):
         ],
     ]
 
-    ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"]
+    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]
 
 
diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py
index e6d782c0a..91c6784fc 100644
--- a/tests/diffusion/distributed/test_hsdp.py
+++ b/tests/diffusion/distributed/test_hsdp.py
@@ -1,19 +1,97 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.
+"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
 
-These tests verify HSDP configuration logic without requiring a distributed environment.
-"""
+import gc
+import os
+import socket
 
 import pytest
+import torch
+import torch.distributed as dist
 import torch.nn as nn
+from torch.distributed.tensor import DeviceMesh, DTensor
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    _unshardable_parameters,
+    shard_model,
+)
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
 
 
+class _PackedBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(2, 2))
+        self.packed_weight = nn.Parameter(torch.ones(2, 2, dtype=torch.uint8), requires_grad=False)
+        self.input_global_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False)
+
+
+class _PackedModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.block = _PackedBlock()
+        self.root_weight = nn.Parameter(torch.ones(2))
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+@pytest.fixture(scope="module")
+def cpu_process_group():
+    if dist.is_initialized():
+        yield
+        return
+
+    master_port = _find_free_port()
+    os.environ.update(
+        {
+            "RANK": "0",
+            "LOCAL_RANK": "0",
+            "WORLD_SIZE": "1",
+            "MASTER_ADDR": "127.0.0.1",
+            "MASTER_PORT": str(master_port),
+        }
+    )
+    dist.init_process_group("gloo", rank=0, world_size=1)
+    try:
+        yield
+    finally:
+        dist.destroy_process_group()
+        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
+            os.environ.pop(key, None)
+        gc.collect()
+
+
+def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
+    model = _PackedModel()
+    ignored_params = _unshardable_parameters(model)
+    expected_ignored = {model.block.packed_weight, model.block.input_global_scale}
+    assert ignored_params == expected_ignored
+
+    packed_weight = model.block.packed_weight
+    input_global_scale = model.block.input_global_scale
+    shard_model(
+        model,
+        mesh=DeviceMesh("cpu", [0]),
+        hsdp_shard_conditions=[lambda name, _module: name == "block"],
+        ignored_params=ignored_params,
+    )
+
+    assert isinstance(model.block.weight, DTensor)
+    assert isinstance(model.root_weight, DTensor)
+    assert model.block.packed_weight is packed_weight
+    assert model.block.input_global_scale is input_global_scale
+    assert not isinstance(model.block.packed_weight, DTensor)
+    assert not isinstance(model.block.input_global_scale, DTensor)
+
+
 class TestHSDPInferenceConfig:
     """Tests for HSDPInferenceConfig dataclass."""
 
diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py
new file mode 100644
index 000000000..29fcf1804
--- /dev/null
+++ b/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()
diff --git a/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py
deleted file mode 100644
index 0c81dc801..000000000
--- a/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py
+++ /dev/null
@@ -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]))
diff --git a/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py
deleted file mode 100644
index 6b7392c80..000000000
--- a/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py
+++ /dev/null
@@ -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
diff --git a/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py
deleted file mode 100644
index cccea56fb..000000000
--- a/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py
+++ /dev/null
@@ -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
diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py
index 376dbe685..f370d5b49 100644
--- a/tests/diffusion/model_loader/test_diffusers_loader.py
+++ b/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"))
@@ -202,6 +179,59 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert get_current_diffusion_config_or_none() is None
 
 
+def test_hsdp_processes_quantized_weights_before_sharding(mocker):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.diffusion.offloader.module_collector import PipelineModules
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=True,
+            hsdp_replicate_size=1,
+            hsdp_shard_size=2,
+        ),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    loader.quant_config = object()
+
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+
+    loader._init_from_load_format = mocker.Mock(return_value=model)  # type: ignore[method-assign]
+    loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load"))  # type: ignore[method-assign]
+    loader._process_weights_after_loading = mocker.Mock(  # type: ignore[method-assign]
+        side_effect=lambda _model, _device: events.append("process")
+    )
+    mocker.patch.object(
+        loader_mod.ModuleDiscovery,
+        "discover",
+        return_value=PipelineModules(
+            dits=[model.transformer],
+            dit_names=["transformer"],
+            vaes=[],
+            encoders=[],
+            encoder_names=[],
+            resident_modules=[],
+            resident_names=[],
+        ),
+    )
+    mocker.patch(
+        "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp",
+        side_effect=lambda _model: events.append("prepare"),
+    )
+    mocker.patch.object(
+        loader_mod,
+        "apply_hsdp_to_model",
+        side_effect=lambda *_args, **_kwargs: events.append("shard"),
+    )
+
+    loader._load_model_with_hsdp(torch.device("cpu"))
+
+    assert events == ["load", "process", "prepare", "shard"]
+
+
 def test_get_all_weights(prefetch_helios_model, mock_tp_group):
     """Ensure that get all weights on a tiny model resolves to nonempty weights."""
     od_config = OmniDiffusionConfig(
diff --git a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py b/tests/diffusion/model_loader/test_diffusers_loader_gguf.py
deleted file mode 100644
index f747dca81..000000000
--- a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py
+++ /dev/null
@@ -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"]
diff --git a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py b/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py
index 513e80d28..8c429a00f 100644
--- a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py
+++ b/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py
@@ -9,6 +9,7 @@ import torch.nn as nn
 
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     ModelOptFp8CheckpointAdapter,
+    ModelOptNvFp4CheckpointAdapter,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -42,6 +43,70 @@ class _QuantizedPackedModelOptModel(nn.Module):
         )
 
 
+class _RemappedModelOptModel(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Linear(2, 2, bias=False)
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        return {
+            "transformer.orig.proj.weight": "runtime.proj.weight",
+        }.get(name, name)
+
+
+class _QuantizedRemappedModelOptModel(nn.Module):
+    """FP8 target whose params live under a namespace the raw checkpoint keys
+    never use, reachable only via ``remap_checkpoint_key``.
+
+    Mirrors the Cosmos3 case that regressed: the generic weights-mapper cannot
+    resolve ``transformer.orig.*`` to ``runtime.*``, so without the remap hook
+    the adapter treats every tensor as unresolved and silently drops the scales.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Module()
+        self.runtime.proj.register_parameter(
+            "weight",
+            nn.Parameter(torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "weight_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+        self.runtime.proj.register_parameter(
+            "input_scale",
+            nn.Parameter(torch.empty(1), requires_grad=False),
+        )
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        prefix = "transformer.orig.proj."
+        if name.startswith(prefix):
+            return "runtime.proj." + name[len(prefix) :]
+        return name
+
+
+class _RemappedNvFp4Model(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.runtime = nn.Module()
+        self.runtime.proj = nn.Module()
+        self.runtime.proj.register_parameter(
+            "weight",
+            nn.Parameter(torch.empty(1, 1, dtype=torch.uint8), requires_grad=False),
+        )
+        for name in ("input_scale", "weight_scale", "weight_scale_2"):
+            self.runtime.proj.register_parameter(name, nn.Parameter(torch.empty(1), requires_grad=False))
+
+    @staticmethod
+    def remap_checkpoint_key(name: str) -> str:
+        return name.replace("transformer.orig.proj.", "runtime.proj.")
+
+
 def _make_source() -> SimpleNamespace:
     return SimpleNamespace(
         subfolder="transformer",
@@ -92,3 +157,85 @@ def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target():
         "transformer.block.to_q.weight_scale",
         "transformer.block.to_q.input_scale",
     ]
+
+
+def test_modelopt_adapter_uses_checkpoint_key_remap_for_target_dtype():
+    model = _RemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    scale = torch.tensor([0.5], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight", fp8_weight),
+                    ("transformer.orig.proj.weight_scale", scale),
+                ]
+            )
+        )
+    )
+
+    assert [name for name, _ in adapted] == ["runtime.proj.weight"]
+    assert adapted[0][1].dtype == model.runtime.proj.weight.dtype
+    assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale)
+
+
+def test_modelopt_adapter_remaps_and_keeps_scales_for_quantized_target():
+    """Regression: scales were dropped for a quantized (FP8) target that only the
+    model's ``remap_checkpoint_key`` hook can resolve. Assert both scales are
+    emitted and the FP8 weight passes through unchanged (no dequantization)."""
+    model = _QuantizedRemappedModelOptModel()
+    adapter = ModelOptFp8CheckpointAdapter(model, _make_source())
+    fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn)
+    weight_scale = torch.tensor([0.5], dtype=torch.float32)
+    input_scale = torch.tensor([1.0], dtype=torch.float32)
+
+    adapted = list(
+        adapter.adapt(
+            iter(
+                [
+                    ("transformer.orig.proj.weight_scale", weight_scale),
+                    ("transformer.orig.proj.input_scale", input_scale),
+                    ("transformer.orig.proj.weight", fp8_weight),
+                ]
+            )
+        )
+    )
+
+    emitted = dict(adapted)
+    assert len(adapted) == len(emitted) == 3
+
+    # Both scales are retained and emitted under the remapped runtime names.
+    assert torch.equal(emitted["runtime.proj.weight_scale"], weight_scale)
+    assert torch.equal(emitted["runtime.proj.input_scale"], input_scale)
+
+    # The FP8 weight is passed through unchanged under the remapped target name
+    # (the quantized param is FP8, so it must NOT be dequantized). Compare via
+    # float32 since torch.equal has no FP8 CPU kernel.
+    assert emitted["runtime.proj.weight"].dtype == torch.float8_e4m3fn
+    assert torch.equal(emitted["runtime.proj.weight"].to(torch.float32), fp8_weight.to(torch.float32))
+
+
+def test_modelopt_nvfp4_adapter_remaps_quantized_weights_and_scales():
+    adapter = ModelOptNvFp4CheckpointAdapter(_RemappedNvFp4Model(), _make_source())
+    prefix = "transformer.orig.proj"
+    checkpoint_tensors = [
+        (f"{prefix}.input_scale", torch.tensor([1.0])),
+        (f"{prefix}.weight_scale", torch.tensor([0.5])),
+        (f"{prefix}.weight_scale_2", torch.tensor([0.25])),
+        (f"{prefix}.weight", torch.tensor([[1]], dtype=torch.uint8)),
+    ]
+
+    adapted = list(adapter.adapt(iter(checkpoint_tensors)))
+
+    assert [name for name, _ in adapted] == [
+        name.replace(f"{prefix}.", "runtime.proj.") for name, _ in checkpoint_tensors
+    ]
+
+
+def test_modelopt_nvfp4_adapter_rejects_unconsumed_pre_quant_scale():
+    adapter = ModelOptNvFp4CheckpointAdapter(nn.Module(), _make_source())
+
+    with pytest.raises(ValueError, match="does not consume pre_quant_scale"):
+        list(adapter.adapt(iter([("transformer.proj.pre_quant_scale", torch.tensor([1.0]))])))
diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
index da4ae32e3..1085e3878 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py
@@ -20,26 +20,52 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+def test_pipeline_declares_layerwise_offload_components() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    assert Cosmos3OmniDiffusersPipeline._dit_modules == ["transformer.language_model", "transformer"]
+    assert Cosmos3OmniDiffusersPipeline._encoder_modules == []
+    assert Cosmos3OmniDiffusersPipeline._vae_modules == ["vae"]
+    assert Cosmos3OmniDiffusersPipeline._resident_modules == []
+    assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload")
+
+
 class StubScheduler:
     def __init__(
         self,
         timesteps: list[int] | None = None,
         *,
         flow_shift: float = 1.0,
-        use_karras_sigmas: bool = True,
     ) -> None:
         self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
         self.config = SimpleNamespace(
             num_train_timesteps=1000,
             flow_shift=flow_shift,
-            use_karras_sigmas=use_karras_sigmas,
         )
-        self.set_timesteps_calls: list[tuple[int, torch.device]] = []
+        self.set_timesteps_calls: list[dict[str, Any]] = []
         self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
 
-    def set_timesteps(self, num_steps: int, device: torch.device) -> None:
-        self.set_timesteps_calls.append((num_steps, device))
-        self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64, device=device)
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device | None = None,
+        *,
+        shift: float | None = None,
+        sigmas: list[float] | None = None,
+    ) -> None:
+        self.set_timesteps_calls.append(
+            {
+                "num_inference_steps": num_inference_steps,
+                "device": device,
+                "shift": shift,
+                "sigmas": sigmas,
+            }
+        )
+        if sigmas is not None:
+            self.timesteps = torch.tensor(sigmas, device=device)
+        else:
+            assert num_inference_steps is not None
+            self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device)
 
     def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs):
         del kwargs
@@ -190,6 +216,7 @@ def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
 def make_cosmos3_pipeline():
     def _make():
         from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+            COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT,
             Cosmos3OmniDiffusersPipeline,
         )
 
@@ -203,11 +230,10 @@ def make_cosmos3_pipeline():
         pipeline.vae_scale_factor_temporal = 4
         pipeline.vae_scale_factor_spatial = 8
         pipeline.scheduler = StubScheduler([9, 3], flow_shift=1.0)
-        pipeline._base_scheduler_config = pipeline.scheduler.config
-        pipeline._engine_init_flow_shift = 1.0
-        pipeline._current_flow_shift = 1.0
-        pipeline._base_scheduler_use_karras_sigmas = True
-        pipeline._current_scheduler_use_karras_sigmas = True
+        pipeline._engine_init_flow_shift = COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT
+        pipeline._current_flow_shift = COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT
+        pipeline.is_distilled_model = False
+        pipeline.is_edge_model = False
         pipeline._guidance_scale = None
         pipeline._num_timesteps = None
         pipeline._cosmos3_branch_caches = None
@@ -233,6 +259,7 @@ def make_sampling_params(**overrides: Any) -> SimpleNamespace:
         "num_frames": None,
         "num_inference_steps": None,
         "guidance_scale": None,
+        "guidance_scale_provided": False,
         "generator": None,
         "seed": 123,
         "num_outputs_per_prompt": 1,
@@ -278,13 +305,171 @@ def _mask() -> torch.Tensor:
     return torch.ones(1, 1, dtype=torch.long)
 
 
+def _capture_tokenize_calls(pipeline: Any) -> list[dict[str, Any]]:
+    calls: list[dict[str, Any]] = []
+
+    def _tokenize(
+        text: str,
+        max_sequence_length: int,
+        use_system_prompt: bool = False,
+        system_prompt: str | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        index = len(calls) + 1
+        calls.append(
+            {
+                "text": text,
+                "max_sequence_length": max_sequence_length,
+                "use_system_prompt": use_system_prompt,
+                "system_prompt": system_prompt,
+            }
+        )
+        return _ids(index), torch.full((1, 1), index, dtype=torch.long)
+
+    pipeline._tokenize_prompt = _tokenize
+    return calls
+
+
+@pytest.mark.parametrize(
+    ("provided", "value", "default", "is_distilled", "expected"),
+    [
+        (False, 1.0, 7.0, False, 7.0),
+        (True, 1.0, 7.0, False, 1.0),
+        (True, 4.5, 7.0, False, 4.5),
+        (False, 1.0, 7.0, True, 1.0),
+        (True, 4.5, 7.0, True, 1.0),
+    ],
+)
+def test_resolve_guidance_scale(
+    make_cosmos3_pipeline,
+    provided: bool,
+    value: float,
+    default: float,
+    is_distilled: bool,
+    expected: float,
+) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.is_distilled_model = is_distilled
+    sp = make_sampling_params(
+        guidance_scale=value,
+        guidance_scale_provided=provided,
+    )
+
+    assert pipeline._resolve_guidance_scale(sp, default) == expected
+
+
+def test_distilled_generation_accepts_t2i_and_i2v(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.is_distilled_model = True
+    common = {
+        "action_enabled": False,
+        "transfer_config": None,
+        "is_v2v": False,
+        "sound_enabled": False,
+    }
+
+    pipeline._validate_distilled_generation_mode(
+        is_t2i=True,
+        image_tensor=None,
+        **common,
+    )
+    pipeline._validate_distilled_generation_mode(
+        is_t2i=False,
+        image_tensor=torch.zeros(1),
+        **common,
+    )
+
+
+@pytest.mark.parametrize(
+    ("overrides", "message"),
+    [
+        ({"action_enabled": True}, "action requests are unsupported"),
+        ({"transfer_config": SimpleNamespace()}, "transfer requests are unsupported"),
+        ({"is_v2v": True}, "video-to-video requests are unsupported"),
+        ({"sound_enabled": True}, "sound generation is unsupported"),
+        (
+            {"is_t2i": False, "image_tensor": None},
+            "text-to-video requests are unsupported",
+        ),
+        (
+            {"is_t2i": True, "image_tensor": torch.zeros(1)},
+            "image-conditioned image generation is unsupported",
+        ),
+    ],
+)
+def test_distilled_generation_rejects_unsupported_modes(
+    make_cosmos3_pipeline,
+    overrides: dict[str, Any],
+    message: str,
+) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.is_distilled_model = True
+    kwargs = {
+        "is_t2i": True,
+        "image_tensor": None,
+        "action_enabled": False,
+        "transfer_config": None,
+        "is_v2v": False,
+        "sound_enabled": False,
+    }
+    kwargs.update(overrides)
+
+    with pytest.raises(ValueError, match=message):
+        pipeline._validate_distilled_generation_mode(**kwargs)
+
+
+def test_distilled_generation_rejects_robolab_policy(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.is_distilled_model = True
+
+    with pytest.raises(ValueError, match="RoboLab/action policy requests are unsupported"):
+        pipeline._forward_robolab_policy(make_sampling_params(), None, 0.0)
+
+
+@pytest.mark.parametrize(
+    ("prompt", "sampling_params", "message"),
+    [
+        (
+            {"prompt": "x", "modalities": ["video"], "generate_sound": True},
+            make_sampling_params(),
+            "do not support sound generation",
+        ),
+        (
+            {"prompt": "x", "modalities": ["video"]},
+            make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}),
+            "do not support transfer inference",
+        ),
+        (
+            {
+                "prompt": "x",
+                "modalities": ["video"],
+                "additional_information": {
+                    "preprocessed_video": torch.zeros(1, 3, 5, 16, 16),
+                },
+            },
+            make_sampling_params(height=16, width=16, num_frames=5),
+            "do not support video-to-video generation",
+        ),
+    ],
+)
+def test_edge_forward_rejects_unsupported_modes(
+    make_cosmos3_pipeline,
+    prompt: dict[str, Any],
+    sampling_params: SimpleNamespace,
+    message: str,
+) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.is_edge_model = True
+
+    with pytest.raises(ValueError, match=message):
+        pipeline.forward(SimpleNamespace(prompts=[prompt], sampling_params=sampling_params))
+
+
 def test_pipeline_registered_and_exported() -> None:
     from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
     from vllm_omni.diffusion.models import cosmos3
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
     from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
     from vllm_omni.diffusion.registry import (
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
         _DIFFUSION_IR_OP_PRIORITY_FUNCS,
         _DIFFUSION_MODELS,
         _DIFFUSION_POST_PROCESS_FUNCS,
@@ -294,20 +479,45 @@ def test_pipeline_registered_and_exported() -> None:
     assert issubclass(Cosmos3OmniDiffusersPipeline, nn.Module)
     assert issubclass(Cosmos3OmniDiffusersPipeline, ProgressBarMixin)
     assert Cosmos3OmniDiffusersPipeline.support_image_input is True
-    assert _DIFFUSION_MODELS["Cosmos3OmniDiffusersPipeline"] == (
-        "cosmos3",
-        "pipeline_cosmos3",
-        "Cosmos3OmniDiffusersPipeline",
-    )
-    assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func"
-    assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func"
-    assert (
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_action_post_process_func"
-    )
-    assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_ir_op_priority_func"
-    assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS
     assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__
 
+    for pipeline_name in ("Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"):
+        assert _DIFFUSION_MODELS[pipeline_name] == (
+            "cosmos3",
+            "pipeline_cosmos3",
+            "Cosmos3OmniDiffusersPipeline",
+        )
+        assert _DIFFUSION_PRE_PROCESS_FUNCS[pipeline_name] == "get_cosmos3_pre_process_func"
+        assert _DIFFUSION_POST_PROCESS_FUNCS[pipeline_name] == "get_cosmos3_post_process_func"
+        assert _DIFFUSION_IR_OP_PRIORITY_FUNCS[pipeline_name] == "get_cosmos3_ir_op_priority_func"
+        assert pipeline_name in CUSTOM_DIT_ENABLERS
+
+
+@pytest.mark.parametrize(
+    "pipeline_name",
+    ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"],
+)
+def test_cosmos3_model_index_resolves_pipeline(
+    tmp_path,
+    pipeline_name: str,
+) -> None:
+    import json
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        Cosmos3OmniDiffusersPipeline,
+    )
+    from vllm_omni.diffusion.registry import DiffusionModelRegistry
+
+    (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": pipeline_name}))
+
+    config = OmniDiffusionConfig(model=str(tmp_path))
+    config.enrich_config()
+
+    assert config.model_class_name == pipeline_name
+    resolved_pipeline_cls = DiffusionModelRegistry._try_load_model_cls(config.model_class_name)
+    assert resolved_pipeline_cls is Cosmos3OmniDiffusersPipeline
+
 
 @pytest.fixture
 def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
@@ -329,25 +539,24 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
             return self
 
     class _StubDiffusersScheduler:
+        load_config_calls: list[dict[str, Any]] = []
         from_config_calls: list[dict[str, Any]] = []
 
-        def __init__(self, *, flow_shift: float = 1.0, use_karras_sigmas: bool = True) -> None:
-            self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas)
+        def __init__(self, *, flow_shift: float = 1.0) -> None:
+            self.config = SimpleNamespace(flow_shift=flow_shift)
 
         @classmethod
-        def from_pretrained(cls, *args, **kwargs):
-            return cls()
+        def load_config(cls, *args, **kwargs):
+            cls.load_config_calls.append({"args": args, "kwargs": dict(kwargs)})
+            return {
+                "_class_name": "UniPCMultistepScheduler",
+                "flow_shift": 1.0,
+            }
 
         @classmethod
         def from_config(cls, config, **kwargs):
             cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
-            return cls(
-                flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)),
-                use_karras_sigmas=kwargs.get(
-                    "use_karras_sigmas",
-                    getattr(config, "use_karras_sigmas", True),
-                ),
-            )
+            return cls(flow_shift=kwargs.get("shift", config.get("flow_shift", 1.0)))
 
     class _StubVideoProcessor:
         def __init__(self, *args, **kwargs) -> None:
@@ -355,13 +564,18 @@ def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
 
     monkeypatch.setattr(pipeline_cosmos3, "AutoTokenizer", _StubAutoTokenizer)
     monkeypatch.setattr(pipeline_cosmos3, "DistributedAutoencoderKLWan", _StubDiffusersVAE)
-    monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", _StubDiffusersScheduler)
+    monkeypatch.setattr(pipeline_cosmos3, "FlowUniPCMultistepScheduler", _StubDiffusersScheduler)
     monkeypatch.setattr(pipeline_cosmos3, "VideoProcessor", _StubVideoProcessor)
     monkeypatch.setattr(pipeline_cosmos3, "get_local_device", lambda: torch.device("cpu"))
     return _StubDiffusersScheduler
 
 
-def _make_od_config(*, sound_gen: bool) -> SimpleNamespace:
+def _make_od_config(
+    *,
+    sound_gen: bool,
+    tf_model_config_overrides: dict[str, Any] | None = None,
+    model_config: dict[str, Any] | None = None,
+) -> SimpleNamespace:
     tf_model_config = {
         "hidden_size": 8,
         "num_hidden_layers": 0,
@@ -376,6 +590,8 @@ def _make_od_config(*, sound_gen: bool) -> SimpleNamespace:
     }
     if sound_gen:
         tf_model_config["sound_gen"] = True
+    if tf_model_config_overrides:
+        tf_model_config.update(tf_model_config_overrides)
     return SimpleNamespace(
         enable_cpu_offload=False,
         enable_diffusion_pipeline_profiler=False,
@@ -384,7 +600,7 @@ def _make_od_config(*, sound_gen: bool) -> SimpleNamespace:
         flow_shift=None,
         quantization_config=None,
         custom_pipeline_args={},
-        model_config={},
+        model_config=model_config or {},
         tf_model_config=tf_model_config,
     )
 
@@ -394,16 +610,13 @@ def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_in
 
     pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False))
 
-    assert stub_real_pipeline_init.from_config_calls == []
     assert pipeline._sound_tokenizer is None
     assert pipeline.transformer.sound_gen is False
     assert not hasattr(pipeline.transformer, "audio_proj_in")
     assert not hasattr(pipeline.transformer, "audio_proj_out")
-    assert pipeline._base_scheduler_use_karras_sigmas is True
-    assert pipeline._current_scheduler_use_karras_sigmas is True
 
 
-def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeline_init) -> None:
+def test_pipeline_init_uses_flow_unipc_with_cosmos3_defaults(stub_real_pipeline_init) -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
 
     od_config = _make_od_config(sound_gen=False)
@@ -411,55 +624,302 @@ def test_pipeline_init_rebuilds_scheduler_with_cosmos3_defaults(stub_real_pipeli
 
     pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config)
 
-    assert stub_real_pipeline_init.from_config_calls
-    call = stub_real_pipeline_init.from_config_calls[-1]
-    assert getattr(call["config"], "flow_shift") == 1.0
-    assert call["kwargs"] == {"flow_shift": 2.5}
+    assert len(stub_real_pipeline_init.load_config_calls) == 1
+    assert len(stub_real_pipeline_init.from_config_calls) == 1
+    call = stub_real_pipeline_init.from_config_calls[0]
+    assert call["config"]["_class_name"] == "UniPCMultistepScheduler"
+    assert call["kwargs"] == {
+        "shift": 1.0,
+        "use_dynamic_shifting": False,
+        "prediction_type": "flow_prediction",
+    }
+    assert pipeline.is_distilled_model is False
     assert pipeline._engine_init_flow_shift == 2.5
-    assert pipeline._base_scheduler_use_karras_sigmas is True
 
 
-def test_set_flow_shift_restores_checkpoint_scheduler_mode(
-    make_cosmos3_pipeline,
+@pytest.mark.parametrize(
+    ("scheduler_class_name", "expected_distilled"),
+    [
+        ("FlowMatchEulerDiscreteScheduler", True),
+        ("UniPCMultistepScheduler", False),
+    ],
+)
+def test_pipeline_resolves_scheduler_class_from_checkpoint_file(
+    tmp_path,
+    stub_real_pipeline_init,
     monkeypatch: pytest.MonkeyPatch,
+    scheduler_class_name: str,
+    expected_distilled: bool,
 ) -> None:
+    import json
+
     from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+    from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
+        FlowUniPCMultistepScheduler as RealFlowUniPCMultistepScheduler,
+    )
 
-    class RecordingScheduler:
-        from_config_calls: list[dict[str, Any]] = []
+    t_list = [1.0, 0.75, 0.5, 0.25]
+    scheduler_dir = tmp_path / "scheduler"
+    scheduler_dir.mkdir()
+    scheduler_config = {"_class_name": scheduler_class_name}
+    if expected_distilled:
+        scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list}
+    (scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config))
 
-        def __init__(self, *, flow_shift: float, use_karras_sigmas: bool) -> None:
-            self.config = SimpleNamespace(flow_shift=flow_shift, use_karras_sigmas=use_karras_sigmas)
+    class StubFlowUniPCScheduler(StubScheduler):
+        from_config_calls: list[tuple[Any, dict[str, Any]]] = []
+
+        @classmethod
+        def load_config(cls, *args, **kwargs):
+            return RealFlowUniPCMultistepScheduler.load_config(*args, **kwargs)
 
         @classmethod
         def from_config(cls, config, **kwargs):
-            cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)})
-            return cls(
-                flow_shift=kwargs.get("flow_shift", getattr(config, "flow_shift", 1.0)),
-                use_karras_sigmas=kwargs.get(
-                    "use_karras_sigmas",
-                    getattr(config, "use_karras_sigmas", True),
-                ),
-            )
+            cls.from_config_calls.append((config, dict(kwargs)))
+            return cls()
 
-    monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", RecordingScheduler)
-    pipeline = make_cosmos3_pipeline()
+    class StubFlowMatchScheduler(StubScheduler):
+        from_config_calls: list[tuple[Any, dict[str, Any]]] = []
 
-    pipeline._set_flow_shift(10.0, use_karras_sigmas=False)
-    assert RecordingScheduler.from_config_calls[-1]["kwargs"] == {
-        "flow_shift": 10.0,
-        "use_karras_sigmas": False,
+        @classmethod
+        def from_config(cls, config, **kwargs):
+            cls.from_config_calls.append((config, dict(kwargs)))
+            scheduler = cls()
+            scheduler.config.fixed_step_sampler_config = config["fixed_step_sampler_config"]
+            return scheduler
+
+    monkeypatch.setattr(pipeline_cosmos3, "FlowUniPCMultistepScheduler", StubFlowUniPCScheduler)
+    monkeypatch.setattr(
+        pipeline_cosmos3,
+        "FlowMatchEulerDiscreteScheduler",
+        StubFlowMatchScheduler,
+    )
+
+    od_config = _make_od_config(sound_gen=False)
+    od_config.model = str(tmp_path)
+    pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config)
+
+    assert pipeline.is_distilled_model is expected_distilled
+    if expected_distilled:
+        assert len(StubFlowMatchScheduler.from_config_calls) == 1
+        assert StubFlowMatchScheduler.from_config_calls[0][1] == {"stochastic_sampling": True}
+        assert StubFlowUniPCScheduler.from_config_calls == []
+        assert pipeline._scheduler_init_t_list == t_list
+    else:
+        assert len(StubFlowUniPCScheduler.from_config_calls) == 1
+        assert StubFlowMatchScheduler.from_config_calls == []
+        assert not hasattr(pipeline, "_scheduler_init_t_list")
+
+
+def test_distilled_pipeline_initializes_sde_scheduler(
+    stub_real_pipeline_init,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from diffusers import FlowMatchEulerDiscreteScheduler
+
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        Cosmos3OmniDiffusersPipeline,
+    )
+
+    t_list = [1.0, 0.9375, 0.8333333333333334, 0.625]
+    scheduler_config = {
+        "_class_name": "FlowMatchEulerDiscreteScheduler",
+        "shift": 1.0,
+        "stochastic_sampling": False,
+        "fixed_step_sampler_config": {
+            "sample_type": "sde",
+            "t_list": t_list,
+        },
     }
-    assert pipeline._current_flow_shift == 10.0
-    assert pipeline._current_scheduler_use_karras_sigmas is False
+    monkeypatch.setattr(
+        stub_real_pipeline_init,
+        "load_config",
+        classmethod(lambda cls, *args, **kwargs: scheduler_config),
+    )
 
-    pipeline._set_flow_shift(10.0)
-    assert RecordingScheduler.from_config_calls[-1]["kwargs"] == {"flow_shift": 10.0}
-    assert pipeline._current_flow_shift == 10.0
-    assert pipeline._current_scheduler_use_karras_sigmas is True
+    pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False))
+    assert pipeline.is_distilled_model is True
+    assert isinstance(pipeline.scheduler, FlowMatchEulerDiscreteScheduler)
+    assert pipeline.scheduler.config.stochastic_sampling is True
+    assert pipeline._scheduler_init_t_list == t_list
 
-    pipeline._set_flow_shift(10.0)
-    assert len(RecordingScheduler.from_config_calls) == 2
+
+@pytest.mark.parametrize(
+    ("fixed_step_config", "error_pattern"),
+    [
+        pytest.param(
+            {"t_list": [1.0, 0.5]},
+            r"fixed_step_sampler_config\.sample_type=sde",
+            id="missing-sample-type",
+        ),
+        pytest.param(
+            {"sample_type": "ode", "t_list": [1.0, 0.5]},
+            r"fixed_step_sampler_config\.sample_type=sde",
+            id="unsupported-sample-type",
+        ),
+        pytest.param(
+            {"sample_type": "sde"},
+            r"non-empty fixed_step_sampler_config\.t_list",
+            id="missing-t-list",
+        ),
+        pytest.param(
+            {"sample_type": "sde", "t_list": []},
+            r"non-empty fixed_step_sampler_config\.t_list",
+            id="empty-t-list",
+        ),
+    ],
+)
+def test_distilled_scheduler_validates_fixed_step_config(
+    stub_real_pipeline_init,
+    monkeypatch: pytest.MonkeyPatch,
+    fixed_step_config: dict[str, Any],
+    error_pattern: str,
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        Cosmos3OmniDiffusersPipeline,
+    )
+
+    scheduler_config = {
+        "_class_name": "FlowMatchEulerDiscreteScheduler",
+        "fixed_step_sampler_config": fixed_step_config,
+    }
+    monkeypatch.setattr(
+        stub_real_pipeline_init,
+        "load_config",
+        classmethod(lambda cls, *args, **kwargs: scheduler_config),
+    )
+
+    with pytest.raises(ValueError, match=error_pattern):
+        Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False))
+
+
+def test_pipeline_init_selects_edge_transformer_from_backbone_type(stub_real_pipeline_init) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
+        COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT,
+        Cosmos3OmniDiffusersPipeline,
+    )
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        COSMOS3_EDGE_BACKBONE_TYPE,
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    pipeline = Cosmos3OmniDiffusersPipeline(
+        od_config=_make_od_config(
+            sound_gen=False,
+            tf_model_config_overrides={
+                "backbone_type": COSMOS3_EDGE_BACKBONE_TYPE,
+                "qk_norm_for_text": False,
+                "latent_channel": 48,
+                "latent_patch_size": 2,
+                "temporal_compression_factor": 4,
+                "layer_norm_epsilon": 1e-5,
+            },
+        )
+    )
+
+    assert isinstance(pipeline.transformer, Cosmos3EdgeVFMTransformer)
+    assert pipeline.is_edge_model is True
+    assert pipeline._engine_init_flow_shift == COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT
+
+
+def test_pipeline_init_does_not_select_edge_from_model_index_class_name(stub_real_pipeline_init) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import Cosmos3EdgeVFMTransformer
+
+    pipeline = Cosmos3OmniDiffusersPipeline(
+        od_config=_make_od_config(
+            sound_gen=False,
+            model_config={"_class_name": "Cosmos3EdgeOmniDiffusersPipeline"},
+        )
+    )
+
+    assert isinstance(pipeline.transformer, Cosmos3VFMTransformer)
+    assert not isinstance(pipeline.transformer, Cosmos3EdgeVFMTransformer)
+
+
+def test_flow_unipc_reuses_scheduler_and_forwards_each_request_shift(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    scheduler = pipeline.scheduler
+
+    assert pipeline._engine_init_flow_shift == 10.0
+    assert pipeline._current_flow_shift == 10.0
+    assert scheduler.config.flow_shift == 1.0
+
+    for shift in (3.0, 10.0):
+        pipeline._set_flow_shift(shift)
+        pipeline._set_timesteps(4, torch.device("cpu"), shift=pipeline._current_flow_shift)
+
+    assert pipeline.scheduler is scheduler
+    assert pipeline._current_flow_shift == 10.0
+    assert [call["shift"] for call in scheduler.set_timesteps_calls] == [3.0, 10.0]
+    assert all(call["num_inference_steps"] == 4 for call in scheduler.set_timesteps_calls)
+    assert all(call["sigmas"] is None for call in scheduler.set_timesteps_calls)
+
+
+def test_flow_unipc_reproducible_with_same_seed(make_cosmos3_pipeline) -> None:
+    from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
+        FlowUniPCMultistepScheduler,
+    )
+
+    pipeline = make_cosmos3_pipeline()
+    pipeline.scheduler = FlowUniPCMultistepScheduler(
+        shift=1.0,
+        use_dynamic_shifting=False,
+        prediction_type="flow_prediction",
+    )
+    pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (
+        _ids(1),
+        _mask(),
+        _ids(0),
+        _mask(),
+    )
+
+    def run(seed: int) -> torch.Tensor:
+        request = make_request_batch(
+            {"prompt": "A test video.", "modalities": ["video"]},
+            make_sampling_params(
+                seed=seed,
+                guidance_scale=1.0,
+                guidance_scale_provided=True,
+                num_inference_steps=4,
+                num_frames=5,
+                height=16,
+                width=16,
+                extra_args={"flow_shift": 3.0},
+            ),
+        )
+        output = pipeline.forward(request)
+        return output.output["video"]
+
+    first = run(123)
+    second = run(123)
+    different_seed = run(456)
+
+    torch.testing.assert_close(first, second)
+    assert not torch.equal(first, different_seed)
+
+
+def test_distilled_set_timesteps_uses_fixed_sigma_schedule(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    pipeline.is_distilled_model = True
+    pipeline._scheduler_init_t_list = [1.0, 0.75, 0.5, 0.25]
+
+    pipeline._set_timesteps(
+        num_inference_steps=99,
+        device=torch.device("cpu"),
+        shift=7.0,
+    )
+
+    assert pipeline.scheduler.set_timesteps_calls == [
+        {
+            "num_inference_steps": None,
+            "device": torch.device("cpu"),
+            "shift": None,
+            "sigmas": pipeline._scheduler_init_t_list,
+        }
+    ]
 
 
 def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
@@ -708,11 +1168,11 @@ def test_postprocess_handles_image_video_audio_and_validation() -> None:
 def test_action_postprocess_handles_robolab_policy_outputs() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
         RoboLabPolicyInputs,
-        get_cosmos3_action_post_process_func,
+        get_cosmos3_post_process_func,
         make_robolab_action_postprocess_inputs,
     )
 
-    func = get_cosmos3_action_post_process_func(SimpleNamespace())
+    func = get_cosmos3_post_process_func(SimpleNamespace())
     inputs = RoboLabPolicyInputs(
         prompt="Pick the cube.",
         video_tensor=torch.zeros(1, 3, 3, 16, 16),
@@ -736,13 +1196,41 @@ def test_action_postprocess_handles_robolab_policy_outputs() -> None:
     )
 
     action = torch.tensor([[[0.0, 0.25], [1.0, 0.75]]])
-    custom_output = {"robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs)}
-    processed = func(action, custom_output=custom_output)
+    processed = func(
+        {
+            "payload": {
+                "actions": action,
+            },
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": 2,
+                    "action_mode": "policy",
+                    "domain_id": 7,
+                },
+                "common": {
+                    "action_only_output": True,
+                },
+                "internal": {
+                    "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
+                },
+            },
+        }
+    )
 
-    assert processed.shape == (1, 2)
-    assert processed.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype
-    torch.testing.assert_close(torch.from_numpy(processed), torch.tensor([[1.0, 0.25]]))
-    assert "robolab_action_postprocess" not in custom_output
+    processed_action = processed["payload"]["actions"]
+    assert processed_action.shape == (1, 2)
+    assert processed_action.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype
+    torch.testing.assert_close(torch.from_numpy(processed_action), torch.tensor([[1.0, 0.25]]))
+    assert processed["metadata"] == {
+        "actions": {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        },
+        "common": {
+            "action_only_output": True,
+        },
+    }
 
 
 def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -772,12 +1260,35 @@ def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.Monke
     assert merged.custom_op == ["platform_kernel", "native"]
 
 
-def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None:
-    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+def test_format_and_tokenize_prompts_leaves_plain_prompts_unchanged_by_default(make_cosmos3_pipeline) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_SYSTEM_PROMPT
 
     pipeline = make_cosmos3_pipeline()
-    captured: list[str] = []
-    pipeline._tokenize_prompt = lambda text, *args, **kwargs: (captured.append(text) or _ids(len(captured)), _mask())
+    calls = _capture_tokenize_calls(pipeline)
+
+    result = pipeline._format_and_tokenize_prompts(
+        "  A robot.  ",
+        "  bad.  ",
+        num_frames=48,
+        frame_rate=24,
+        height=720,
+        width=1280,
+        max_sequence_length=32,
+        sp=SimpleNamespace(extra_args={}),
+        use_system_prompt=False,
+        is_t2i=False,
+    )
+
+    assert [call["text"] for call in calls] == ["A robot.", "bad."]
+    assert all(call["max_sequence_length"] == 32 for call in calls)
+    assert all(call["use_system_prompt"] is False for call in calls)
+    assert all(call["system_prompt"] == COSMOS3_SYSTEM_PROMPT for call in calls)
+    assert [tensor.item() for tensor in result] == [1, 1, 2, 2]
+
+
+def test_format_and_tokenize_prompts_applies_video_templates_and_system_override(make_cosmos3_pipeline) -> None:
+    pipeline = make_cosmos3_pipeline()
+    calls = _capture_tokenize_calls(pipeline)
 
     pipeline._format_and_tokenize_prompts(
         "A robot",
@@ -787,21 +1298,167 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No
         height=720,
         width=1280,
         max_sequence_length=32,
+        sp=SimpleNamespace(
+            system_prompt="direct system prompt",
+            extra_args={
+                "use_duration_template": True,
+                "use_resolution_template": True,
+                "system_prompt": "API system prompt",
+            },
+        ),
+        use_system_prompt=True,
+        is_t2i=False,
+    )
+
+    assert calls == [
+        {
+            "text": ("A robot. The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution."),
+            "max_sequence_length": 32,
+            "use_system_prompt": True,
+            "system_prompt": "API system prompt",
+        },
+        {
+            "text": (
+                "bad. The video is not 2.0 seconds long and is not of 24 FPS. This video is not of 720x1280 resolution."
+            ),
+            "max_sequence_length": 32,
+            "use_system_prompt": True,
+            "system_prompt": "API system prompt",
+        },
+    ]
+
+
+def test_format_and_tokenize_prompts_uses_image_templates_for_t2i(make_cosmos3_pipeline) -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_T2I_SYSTEM_PROMPT
+
+    pipeline = make_cosmos3_pipeline()
+    calls = _capture_tokenize_calls(pipeline)
+
+    pipeline._format_and_tokenize_prompts(
+        "A robot",
+        "bad",
+        num_frames=1,
+        frame_rate=24,
+        height=1024,
+        width=768,
+        max_sequence_length=64,
         sp=SimpleNamespace(
             extra_args={
-                "negative_metadata_mode": "inverse",
-                # Duration/resolution metadata templates are off by default
-                # (commit "Change resolution and duration templates to off by
-                # default"); enable them explicitly to exercise the formatting.
                 "use_duration_template": True,
                 "use_resolution_template": True,
             }
         ),
         use_system_prompt=True,
+        is_t2i=True,
+    )
+
+    assert [call["text"] for call in calls] == [
+        "A robot. This image is of 1024x768 resolution.",
+        "bad. This image is not of 1024x768 resolution.",
+    ]
+    assert all("seconds" not in call["text"] and "FPS" not in call["text"] for call in calls)
+    assert all(call["system_prompt"] == COSMOS3_T2I_SYSTEM_PROMPT for call in calls)
+
+
+def test_format_and_tokenize_prompts_rewrites_json_object_metadata(make_cosmos3_pipeline) -> None:
+    import json
+
+    pipeline = make_cosmos3_pipeline()
+    calls = _capture_tokenize_calls(pipeline)
+
+    pipeline._format_and_tokenize_prompts(
+        ('{"caption": "A robot", "duration": "old", "fps": 1, "resolution": {"H": 1, "W": 2}, "aspect_ratio": "1:1"}'),
+        "bad",
+        num_frames=48,
+        frame_rate=24,
+        height=720,
+        width=1280,
+        max_sequence_length=32,
+        sp=SimpleNamespace(
+            extra_args={
+                "aspect_ratio": "16:9",
+                "use_duration_template": True,
+                "use_resolution_template": True,
+            }
+        ),
         is_t2i=False,
     )
-    assert "The video is 2.0 seconds long" in captured[0]
-    assert "The video is not 2.0 seconds long" in captured[1]
+
+    assert json.loads(calls[0]["text"]) == {
+        "caption": "A robot",
+        "duration": "2s",
+        "fps": 24.0,
+        "resolution": {"H": 720, "W": 1280},
+        "aspect_ratio": "16:9",
+    }
+    assert "The video is 2.0 seconds long" not in calls[0]["text"]
+    assert calls[1]["text"] == (
+        "bad. The video is not 2.0 seconds long and is not of 24 FPS. This video is not of 720x1280 resolution."
+    )
+
+
+def test_format_and_tokenize_prompts_removes_video_metadata_from_t2i_json(make_cosmos3_pipeline) -> None:
+    import json
+
+    pipeline = make_cosmos3_pipeline()
+    calls = _capture_tokenize_calls(pipeline)
+
+    pipeline._format_and_tokenize_prompts(
+        '{"caption": "A robot", "duration": "2s", "fps": 24}',
+        "",
+        num_frames=1,
+        frame_rate=24,
+        height=1024,
+        width=768,
+        max_sequence_length=32,
+        sp=SimpleNamespace(extra_args={}),
+        is_t2i=True,
+    )
+
+    assert json.loads(calls[0]["text"]) == {
+        "caption": "A robot",
+        "resolution": {"H": 1024, "W": 768},
+    }
+
+
+@pytest.mark.parametrize(
+    "prompt",
+    [
+        "{malformed",
+        '["A robot"]',
+    ],
+)
+def test_format_and_tokenize_prompts_falls_back_for_non_object_json(
+    make_cosmos3_pipeline,
+    prompt: str,
+) -> None:
+    pipeline = make_cosmos3_pipeline()
+    calls = _capture_tokenize_calls(pipeline)
+
+    pipeline._format_and_tokenize_prompts(
+        prompt,
+        "",
+        num_frames=48,
+        frame_rate=24,
+        height=720,
+        width=1280,
+        max_sequence_length=32,
+        sp=SimpleNamespace(
+            extra_args={
+                "use_duration_template": True,
+                "use_resolution_template": True,
+            }
+        ),
+        is_t2i=False,
+    )
+
+    assert calls[0]["text"] == (
+        f"{prompt}. The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution."
+    )
+
+
+def test_checkpoint_key_remap() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
 
     remaps = {
         "embed_tokens.weight": "transformer.language_model.embed_tokens.weight",
@@ -813,9 +1470,14 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No
         "layers.3.self_attn.to_q.weight": "transformer.language_model.layers.3.self_attn.to_q.weight",
         "layers.3.self_attn.to_out.weight": "transformer.language_model.layers.3.self_attn.to_out.weight",
         "layers.3.self_attn.norm_q.weight": "transformer.language_model.layers.3.self_attn.norm_q.weight",
+        "layers.3.self_attn.k_norm_und_for_gen.weight": (
+            "transformer.language_model.layers.3.self_attn.k_norm_und_for_gen.weight"
+        ),
         "layers.3.self_attn.add_q_proj.weight": "transformer.gen_layers.3.cross_attention.to_q.weight",
         "layers.3.self_attn.to_add_out.weight": "transformer.gen_layers.3.cross_attention.to_out.weight",
         "layers.3.self_attn.norm_added_q.weight": "transformer.gen_layers.3.cross_attention.norm_q.weight",
+        "layers.3.mlp_moe_gen.up_proj.weight": "transformer.gen_layers.3.mlp.up_proj.weight",
+        "layers.3.mlp_moe_gen.down_proj.weight": "transformer.gen_layers.3.mlp.down_proj.weight",
         "transformer.model.layers.3.self_attn.add_k_proj.weight": (
             "transformer.gen_layers.3.cross_attention.to_k.weight"
         ),
@@ -916,6 +1578,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)
@@ -1116,9 +1821,9 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
     pipeline._prepare_transfer_latents = fake_prepare
     pipeline.diffuse_transfer = fake_diffuse_transfer
 
-    def fake_set_flow_shift(target, *, use_karras_sigmas=None):
+    def fake_set_flow_shift(target):
         captured.setdefault("flow_shifts", []).append(target)
-        captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas)
+        pipeline._current_flow_shift = float(target)
 
     pipeline._set_flow_shift = fake_set_flow_shift
     pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta")
@@ -1153,10 +1858,10 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint
     assert captured["format_frame_rate"] == expected_fps
     assert captured["shared_kwargs"]["fps"] == expected_fps
     assert captured["flow_shifts"] == [10.0]
-    # Transfer uses the V2V flow-sigma schedule (karras off) so flow_shift applies.
-    assert captured["scheduler_use_karras_sigmas"] == [False]
-    assert output.custom_output["fps"] == expected_fps
-    assert output.output["video"].device.type == "meta"
+    # Transfer applies the V2V flow shift when building its timestep schedule.
+    assert [call["shift"] for call in pipeline.scheduler.set_timesteps_calls] == [10.0]
+    assert output.output["metadata"]["video"]["fps"] == expected_fps
+    assert output.output["payload"]["video"].device.type == "meta"
 
 
 def test_forward_transfer_runs_multichunk_overlap_path(
@@ -1225,12 +1930,12 @@ def test_forward_transfer_runs_multichunk_overlap_path(
 
     assert captured["conditional_frames"] == [2, 1]
     assert len(captured["decode_calls"]) == 2
-    assert output.output["video"].shape == (1, 3, 8, 16, 16)
+    assert output.output["payload"]["video"].shape == (1, 3, 8, 16, 16)
     torch.testing.assert_close(
-        output.output["video"][0, 0, :, 0, 0],
+        output.output["payload"]["video"][0, 0, :, 0, 0],
         torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2, 0.2, 0.3, 0.4]),
     )
-    assert output.custom_output["transfer_controls"]["edge"].shape == (1, 3, 8, 16, 16)
+    assert output.output["metadata"]["transfer"]["controls"]["edge"].shape == (1, 3, 8, 16, 16)
     torch.testing.assert_close(captured["targets"][0][:, :, 0], torch.full((1, 3, 16, 16), -1.0))
     torch.testing.assert_close(captured["targets"][0][:, :, 1], torch.full((1, 3, 16, 16), 1.0))
     torch.testing.assert_close(captured["targets"][0][:, :, 2:], torch.full((1, 3, 3, 16, 16), 1.0))
@@ -1314,9 +2019,9 @@ class TestForwardRouting:
         pipeline._format_and_tokenize_prompts = fake_format
         pipeline._prepare_latents = fake_prepare
 
-        def fake_set_flow_shift(target, *, use_karras_sigmas=None):
+        def fake_set_flow_shift(target):
             captured.setdefault("flow_shifts", []).append(target)
-            captured.setdefault("scheduler_use_karras_sigmas", []).append(use_karras_sigmas)
+            pipeline._current_flow_shift = float(target)
 
         pipeline._set_flow_shift = fake_set_flow_shift
         pipeline.diffuse = fake_diffuse
@@ -1343,7 +2048,7 @@ class TestForwardRouting:
                 {
                     "key": "video",
                     "is_t2i": False,
-                    "flow": [1.0],
+                    "flow": [10.0],
                     "steps": [35],
                     "frames": 189,
                 },
@@ -1366,8 +2071,8 @@ class TestForwardRouting:
         assert captured["format"]["is_t2i"] is expected["is_t2i"]
         assert captured["format"]["num_frames"] == expected["frames"]
         assert captured["flow_shifts"] == expected["flow"]
-        assert captured["scheduler_use_karras_sigmas"] == [None]
-        assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
+        assert [call["num_inference_steps"] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
+        assert all(call["shift"] == expected["flow"][0] for call in pipeline.scheduler.set_timesteps_calls)
 
     def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None:
         pipeline = make_cosmos3_pipeline()
@@ -1414,7 +2119,6 @@ class TestForwardRouting:
             )
         )
         assert captured["flow_shifts"][-1] == 10.0
-        assert captured["scheduler_use_karras_sigmas"][-1] is False
         assert captured["format"]["negative_prompt"] == ""
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
         assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
@@ -1454,8 +2158,13 @@ class TestForwardRouting:
             )
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
-        assert output.custom_output["action"].shape == (1, 2, 2)
-        assert "action_only_output" not in output.custom_output
+        assert output.output["payload"]["actions"].shape == (1, 2, 2)
+        assert output.output["metadata"]["actions"] == {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        }
+        assert "common" not in output.output["metadata"]
 
     def test_forward_dispatches_robolab_policy_flow(
         self,
@@ -1539,12 +2248,15 @@ class TestForwardRouting:
         assert captured["prepare_action_video"]["kwargs"] == {"image_size": None}
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
         assert captured["diffuse_calls"][-1]["timesteps"].tolist() == [4, 3, 2, 1]
-        assert output.output == {}
-        assert output.custom_output["action_only_output"] is True
-        assert output.custom_output["action"].shape == (1, 2, 2)
-        assert "actions" not in output.custom_output
-        assert "robolab_action_postprocess" in output.custom_output
-        assert "robolab_policy_inputs" not in output.custom_output
+        assert output.output["payload"]["actions"].shape == (1, 2, 2)
+        assert output.output["metadata"]["actions"] == {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        }
+        assert output.output["metadata"]["common"]["action_only_output"] is True
+        assert "robolab_action_postprocess" in output.output["metadata"]["internal"]
+        assert "robolab_policy_inputs" not in output.output["metadata"]
 
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),
diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py
index 70598c952..cb59dffe5 100644
--- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py
+++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py
@@ -10,9 +10,38 @@ import pytest
 import torch
 from torch import nn
 
+from vllm_omni.platforms import current_omni_platform
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
+@pytest.fixture(autouse=True)
+def _single_rank_tensor_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Provide the TP metadata required by vLLM parallel linear layers."""
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+
+    from vllm_omni.diffusion.models.cosmos3 import (
+        transformer_cosmos3,
+        transformer_cosmos3_edge,
+    )
+
+    monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(transformer_cosmos3, "get_tensor_model_parallel_world_size", lambda: 1)
+    monkeypatch.setattr(transformer_cosmos3_edge, "get_tensor_model_parallel_world_size", lambda: 1)
+
+
+@pytest.fixture
+def accelerator_device() -> torch.device:
+    """Provide an accelerator device, skipping when none is available."""
+    if current_omni_platform.get_device_count() == 0:
+        pytest.skip("Accelerator required for this test")
+    return current_omni_platform.get_torch_device(0)
+
+
 def _tiny_cosmos3_config(**overrides):
     config = {
         "hidden_size": 8,
@@ -30,6 +59,22 @@ def _tiny_cosmos3_config(**overrides):
     return config
 
 
+def _tiny_cosmos3_edge_config(**overrides):
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        COSMOS3_EDGE_BACKBONE_TYPE,
+    )
+
+    config = _tiny_cosmos3_config(
+        latent_patch_size=2,
+        latent_channel=48,
+        temporal_compression_factor=4,
+        backbone_type=COSMOS3_EDGE_BACKBONE_TYPE,
+        qk_norm_for_text=False,
+    )
+    config.update(overrides)
+    return config
+
+
 def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
         compute_mrope_position_ids_action,
@@ -68,6 +113,16 @@ def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     assert action_offset == 14
 
 
+def test_timestep_embedder_stores_frequencies_in_fp32() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        TimestepEmbedder,
+    )
+
+    embedder = TimestepEmbedder(hidden_size=8, frequency_embedding_size=16)
+
+    assert embedder.freqs.dtype == torch.float32
+
+
 @pytest.mark.parametrize(
     ("key", "value"),
     [
@@ -87,29 +142,245 @@ def test_validate_supported_config_rejects_unsupported_flags(key: str, value) ->
     Cosmos3VFMTransformer._validate_supported_config(None)
 
 
-def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
+def test_edge_config_resolves_nemotron_defaults() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
 
-    model = object.__new__(Cosmos3VFMTransformer)
+    model = Cosmos3EdgeVFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_edge_config(
+                layer_norm_epsilon=1e-5,
+                rope_scaling={},
+                rope_parameters={"mrope_section": [24, 20, 20]},
+            ),
+            dtype=torch.float32,
+        )
+    )
+
+    assert model.rms_norm_eps == 1e-5
+    assert model.rope_theta == 100_000_000
+    assert model.mrope_section == [24, 20, 20]
+    assert model.latent_channel_size == 48
+    assert model.latent_patch_size == 2
+    assert model.temporal_compression_factor == 4
+
+
+def test_edge_config_requires_backbone_type() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    config = _tiny_cosmos3_edge_config()
+    config.pop("backbone_type")
+
+    with pytest.raises(ValueError, match="must declare backbone_type"):
+        Cosmos3EdgeVFMTransformer(
+            SimpleNamespace(
+                tf_model_config=config,
+                dtype=torch.float32,
+            )
+        )
+
+
+def test_edge_und_and_gen_use_relu2_and_no_und_qk_norm() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+        Cosmos3Relu2MLP,
+    )
+
+    model = Cosmos3EdgeVFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_edge_config(
+                num_hidden_layers=1,
+                use_k_norm_und_for_gen=False,
+            ),
+            dtype=torch.float32,
+        )
+    )
+
+    layer = model.language_model.layers[0]
+    gen_layer = model.gen_layers[0]
+    assert isinstance(layer.mlp, Cosmos3Relu2MLP)
+    assert not hasattr(layer.mlp, "gate_proj")
+    assert isinstance(gen_layer.mlp, Cosmos3Relu2MLP)
+    assert not hasattr(gen_layer.mlp, "gate_proj")
+    assert not hasattr(layer.self_attn, "norm_q")
+    assert not hasattr(layer.self_attn, "norm_k")
+    assert layer.self_attn.k_norm_und_for_gen is None
+
+
+def test_edge_creates_gen_k_norm_when_enabled() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    model = Cosmos3EdgeVFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_edge_config(
+                num_hidden_layers=1,
+                qk_norm_for_diffusion=True,
+                use_k_norm_und_for_gen=True,
+            ),
+            dtype=torch.float32,
+        )
+    )
+
+    k_norm = model.language_model.layers[0].self_attn.k_norm_und_for_gen
+    assert model.use_k_norm_und_for_gen is True
+    assert isinstance(k_norm, RMSNorm)
+    assert isinstance(model.gen_layers[0].cross_attention.norm_q, RMSNorm)
+    assert isinstance(model.gen_layers[0].cross_attention.norm_k, RMSNorm)
+
+
+def test_edge_supports_no_gen_qk_norm_variant() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    model = Cosmos3EdgeVFMTransformer(
+        SimpleNamespace(
+            tf_model_config=_tiny_cosmos3_edge_config(
+                num_hidden_layers=1,
+                qk_norm_for_diffusion=False,
+                use_k_norm_und_for_gen=False,
+            ),
+            dtype=torch.float32,
+        )
+    )
+
+    assert model.qk_norm_for_diffusion is False
+    assert model.language_model.layers[0].self_attn.k_norm_und_for_gen is None
+    assert model.gen_layers[0].cross_attention.qk_norm is False
+    assert not hasattr(model.gen_layers[0].cross_attention, "norm_q")
+    assert not hasattr(model.gen_layers[0].cross_attention, "norm_k")
+
+
+def test_edge_validates_required_relu2_weights() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    model = object.__new__(Cosmos3EdgeVFMTransformer)
     nn.Module.__init__(model)
-    model.language_model = nn.Module()
-    model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)])
-    model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
-    model.norm_moe_gen = nn.LayerNorm(2)
+    model.num_hidden_layers = 1
+    model.qk_norm_for_diffusion = True
+    model.use_k_norm_und_for_gen = True
 
-    matched = [
-        name
-        for name, module in model.named_modules()
-        if any(condition(name, module) for condition in model._hsdp_shard_conditions)
-    ]
-    assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
-    assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
-    assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"]
+    complete = {
+        "transformer.language_model.layers.0.mlp.up_proj.weight",
+        "transformer.language_model.layers.0.mlp.down_proj.weight",
+        "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight",
+        "transformer.gen_layers.0.mlp.up_proj.weight",
+        "transformer.gen_layers.0.mlp.down_proj.weight",
+    }
+    model.validate_loaded_weights(complete)
 
-    model.latent_patch_size = 2
-    model.latent_channel_size = 3
-    latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
-    torch.testing.assert_close(model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), latents)
+    mlp_weights = complete - {
+        "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight",
+    }
+    for missing_weight in mlp_weights:
+        with pytest.raises(ValueError, match="missing required weights"):
+            model.validate_loaded_weights(complete - {missing_weight})
+
+    missing_k_norm = complete - {"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight"}
+    with pytest.raises(ValueError, match=r"self_attn\.k_norm_und_for_gen"):
+        model.validate_loaded_weights(missing_k_norm)
+
+    model.use_k_norm_und_for_gen = False
+    model.validate_loaded_weights(missing_k_norm)
+
+
+def test_edge_gen_cached_k_is_normalized_but_reasoner_uses_raw_k() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeCausalAttention,
+    )
+
+    class ScaleProjection(nn.Module):
+        def __init__(self, scale: float) -> None:
+            super().__init__()
+            self.scale = scale
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return x * self.scale
+
+    class IdentityOutput(nn.Module):
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return x
+
+    class RawKAttention(nn.Module):
+        def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+            del v
+            return q + k
+
+    attn = object.__new__(Cosmos3EdgeCausalAttention)
+    nn.Module.__init__(attn)
+    attn.hidden_size = 4
+    attn.num_heads = 2
+    attn.num_kv_heads = 2
+    attn.head_dim = 2
+    attn.num_heads_local = 2
+    attn.num_kv_heads_local = 2
+    attn.to_q = ScaleProjection(1.0)
+    attn.to_k = ScaleProjection(1.0)
+    attn.to_v = ScaleProjection(1.0)
+    attn.to_out = IdentityOutput()
+    attn.k_norm_und_for_gen = RMSNorm(2, eps=1e-5)
+    attn.attn = RawKAttention()
+
+    hidden_states = torch.tensor(
+        [[[1.0, 2.0, 3.0, 4.0], [2.0, 1.0, 4.0, 3.0]]],
+        dtype=torch.float32,
+    )
+    freqs_cos = torch.ones(1, 2, 1, 2)
+    freqs_sin = torch.zeros(1, 2, 1, 2)
+
+    reasoner_out, gen_k, gen_v = attn(hidden_states, freqs_cos, freqs_sin)
+    attn.to_k.scale = 3.0
+    scaled_reasoner_out, scaled_gen_k, scaled_gen_v = attn(hidden_states, freqs_cos, freqs_sin)
+
+    torch.testing.assert_close(gen_k, scaled_gen_k)
+    torch.testing.assert_close(gen_v, scaled_gen_v)
+    assert not torch.allclose(reasoner_out, scaled_reasoner_out)
+
+
+def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
+        Cosmos3EdgeLanguageModel,
+        Cosmos3EdgeVFMTransformer,
+    )
+
+    for transformer_cls, language_model_cls in (
+        (Cosmos3VFMTransformer, Cosmos3LanguageModel),
+        (Cosmos3EdgeVFMTransformer, Cosmos3EdgeLanguageModel),
+    ):
+        model = object.__new__(transformer_cls)
+        nn.Module.__init__(model)
+        model.language_model = nn.Module()
+        model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)])
+        model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
+        model.norm_moe_gen = nn.LayerNorm(2)
+
+        matched = [
+            name
+            for name, module in model.named_modules()
+            if any(condition(name, module) for condition in model._hsdp_shard_conditions)
+        ]
+        assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
+        assert transformer_cls._layerwise_offload_blocks_attrs == ["gen_layers"]
+        assert language_model_cls._layerwise_offload_blocks_attrs == ["layers"]
+        assert transformer_cls._repeated_blocks == ["Cosmos3GenDecoderLayer"]
+
+        model.latent_patch_size = 2
+        model.latent_channel_size = 3
+        latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
+        torch.testing.assert_close(
+            model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5),
+            latents,
+        )
 
 
 def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -131,6 +402,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
+def test_qwen_gen_mlp_remains_gated() -> None:
+    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
+
+    model = Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1), dtype=torch.float32)
+    )
+
+    assert isinstance(model.gen_layers[0].mlp, Cosmos3GatedMLP)
+    assert hasattr(model.gen_layers[0].mlp, "gate_proj")
+
+
+def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    class ToyReasonerLayer(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, hidden: torch.Tensor, freqs):
+            del freqs
+            kv = hidden.new_zeros(hidden.shape[0], hidden.shape[1], 1, 1)
+            return hidden + self.weight.to(hidden.dtype), kv, kv
+
+    class ToyGeneratorLayer(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.ones(1))
+
+        def forward(self, hidden: torch.Tensor, **kwargs):
+            del kwargs
+            return hidden + self.weight.to(hidden.dtype)
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.language_model.layers = nn.ModuleList([ToyReasonerLayer()])
+    model.gen_layers = nn.ModuleList([ToyGeneratorLayer()])
+    model.enable_model_cpu_offload(device=torch.device("cpu"), pin_memory=False)
+
+    assert model._model_cpu_offload_enabled
+    assert set(model._model_cpu_offload_components()) == {"reasoner", "generator"}
+    assert model.device == torch.device("cpu")
+
+    output = model(
+        hidden_states=torch.zeros(1, 2, 1, 2, 2),
+        timestep=torch.tensor([1.0]),
+        text_ids=torch.tensor([[1, 2]], dtype=torch.long),
+        text_mask=torch.ones(1, 2, dtype=torch.long),
+        video_shape=(1, 2, 2),
+        fps=24.0,
+    )
+
+    assert tuple(output.shape) == (1, 2, 1, 2, 2)
+    # The reasoner runs once then the generator component stays resident for GEN.
+    assert model._active_model_cpu_offload_component == "generator"
+
+    model.disable_model_cpu_offload()
+    assert not model._model_cpu_offload_enabled
+
+
+def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device(
+    monkeypatch: pytest.MonkeyPatch, accelerator_device: torch.device
+) -> None:
+    from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
+
+    monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
+
+    model = transformer_cosmos3.Cosmos3VFMTransformer(
+        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
+    )
+    model.language_model.layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.to(accelerator_device)
+
+    reasoner_param = model.language_model.layers[0].weight
+    generator_param = model.gen_layers[0].weight
+
+    model.enable_model_cpu_offload(device=accelerator_device, pin_memory=False)
+
+    assert model._model_cpu_offload_enabled
+    # On enable, every group is parked on CPU until a phase activates it.
+    assert reasoner_param.device.type == "cpu"
+    assert generator_param.device.type == "cpu"
+
+    model._activate_model_cpu_offload_component("reasoner")
+    assert reasoner_param.device == accelerator_device
+    assert generator_param.device.type == "cpu"
+
+    model._activate_model_cpu_offload_component("generator")
+    assert reasoner_param.device.type == "cpu"
+    assert generator_param.device == accelerator_device
+
+    model.disable_model_cpu_offload()
+    assert not model._model_cpu_offload_enabled
+    # Disable restores both components onto the device.
+    assert reasoner_param.device == accelerator_device
+    assert generator_param.device == accelerator_device
+
+
 def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 
diff --git a/tests/diffusion/models/flux/test_flux_pipeline.py b/tests/diffusion/models/flux/test_flux_pipeline.py
index fd5d7a2ac..5a06133f4 100644
--- a/tests/diffusion/models/flux/test_flux_pipeline.py
+++ b/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
diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py
index 626f78eed..e8ade2e51 100644
--- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py
+++ b/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."""
 
diff --git a/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py b/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py
index dc4d49b11..9609a39ba 100644
--- a/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py
+++ b/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py
@@ -524,11 +524,11 @@ class TestAr2DiffusionIntegration:
         self,
         token_ids: list[int],
         prompt_kwargs: dict[str, Any] | None = None,
-    ) -> tuple[Any, Any]:
+    ) -> dict[str, Any] | None:
         """Invoke ``ar2diffusion`` with a single fake AR output.
 
-        Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can
-        assert on the resolved height/width.
+        Returns the diffusion request payload so callers can assert on the
+        resolved height/width.
         """
         from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
             ar2diffusion,
@@ -601,8 +601,8 @@ class TestAr2DiffusionIntegration:
             token_ids=[150000, 150001, token_id],
             prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
         )
-        assert len(result) == 1
-        diffusion_input = result[0]
+        assert result is not None
+        diffusion_input = result
         assert diffusion_input["height"] == expected_h, (
             f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}"
         )
@@ -617,8 +617,8 @@ class TestAr2DiffusionIntegration:
             token_ids=[150000, 150001],  # no ratio tokens
             prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024},
         )
-        assert len(result) == 1
-        diffusion_input = result[0]
+        assert result is not None
+        diffusion_input = result
         assert diffusion_input["height"] == 768
         assert diffusion_input["width"] == 1024
 
@@ -633,9 +633,9 @@ class TestAr2DiffusionIntegration:
                 token_ids=[150000],
                 prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024},
             )
-        assert len(result) == 1
-        assert result[0]["height"] == 512
-        assert result[0]["width"] == 512
+        assert result is not None
+        assert result["height"] == 512
+        assert result["width"] == 512
 
     # -- cache behaviour -----------------------------------------------------
 
diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py
new file mode 100644
index 000000000..3e592f14b
--- /dev/null
+++ b/tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py
@@ -0,0 +1,436 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LTX-2.3 forward-parallel behavior."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestCFGParallelIntegration:
+    def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
+        """LTX23Pipeline must use the shared CFG parallel implementation."""
+        from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        assert issubclass(LTX23Pipeline, CFGParallelMixin)
+
+    def test_ltx2_transformer_has_dit_cache_config(self):
+        """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
+        from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+        from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
+
+        adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
+        assert isinstance(adapter_config, CacheDiTAdapterConfig)
+        assert adapter_config.has_separate_cfg
+
+
+class TestCFGParallelHelpers:
+    """Test LTX-2.3 CFG helper math without loading model weights."""
+
+    def test_combine_cfg_noise_matches_x0_space_formula(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        pipe = object.__new__(LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        video_combined, audio_combined = pipe.combine_cfg_noise(
+            (video_pos, audio_pos),
+            (video_neg, audio_neg),
+            scale,
+            video_latents=video_sample,
+            audio_latents=audio_sample,
+            video_sigma=video_sigma,
+            audio_sigma=audio_sigma,
+        )
+
+        x0_video_cond = video_sample - video_pos * video_sigma
+        x0_video_uncond = video_sample - video_neg * video_sigma
+        x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond)
+        expected_video = (video_sample - x0_video_guided) / video_sigma
+
+        x0_audio_cond = audio_sample - audio_pos * audio_sigma
+        x0_audio_uncond = audio_sample - audio_neg * audio_sigma
+        x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond)
+        expected_audio = (audio_sample - x0_audio_guided) / audio_sigma
+        assert torch.allclose(video_combined, expected_video)
+        assert torch.allclose(audio_combined, expected_audio)
+
+    def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        video_sample = torch.tensor([[[1.0, -2.0]]])
+        audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
+        audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
+        video_sigma = torch.tensor(0.25)
+        audio_sigma = torch.tensor(0.5)
+        scale = 4.0
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if tensor.shape == video_pos.shape:
+                    return [video_pos, video_neg]
+                return [audio_pos, audio_neg]
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_sample,
+            video_pos,
+            video_neg,
+            video_sigma,
+            scale,
+        )
+        expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_sample,
+            audio_pos,
+            audio_neg,
+            audio_sigma,
+            scale,
+        )
+
+        for rank, expected_branch in ((0, "positive"), (1, "negative")):
+            calls = []
+            monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
+
+            def fake_predict_noise(**kwargs):
+                calls.append(kwargs["branch"])
+                if kwargs["branch"] == "positive":
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+            object.__setattr__(pipe, "predict_noise", fake_predict_noise)
+            video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
+                true_cfg_scale=scale,
+                positive_kwargs={"branch": "positive"},
+                negative_kwargs={"branch": "negative"},
+                cfg_normalize=False,
+                video_latents=video_sample,
+                audio_latents=audio_sample,
+                video_sigma=video_sigma,
+                audio_sigma=audio_sigma,
+            )
+
+            assert calls == [expected_branch]
+            torch.testing.assert_close(video_combined, expected_video)
+            torch.testing.assert_close(audio_combined, expected_audio)
+
+        assert "_cfg_video_latents" not in pipe.__dict__
+        assert "_cfg_audio_latents" not in pipe.__dict__
+
+
+class TestCFGParallelForwardPath:
+    """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
+
+    def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.device = torch.device("cpu")
+        pipe.tokenizer_max_length = 4
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 1)
+
+        class StopAtEncodePromptError(Exception):
+            pass
+
+        captured = {}
+
+        def fake_encode_prompt(**kwargs):
+            captured.update(kwargs)
+            raise StopAtEncodePromptError
+
+        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
+
+        prompt_embeds_a = torch.zeros(2, 3)
+        prompt_embeds_b = torch.ones(2, 3)
+        negative_prompt_embeds_a = torch.full((2, 3), 2.0)
+        negative_prompt_embeds_b = torch.full((2, 3), 3.0)
+        prompt_attention_mask_a = torch.tensor([True, True])
+        prompt_attention_mask_b = torch.tensor([True, False])
+        negative_attention_mask_a = torch.tensor([False, True])
+        negative_attention_mask_b = torch.tensor([False, False])
+
+        requests = [
+            OmniDiffusionRequest(
+                prompt={
+                    "prompt": "prompt-a",
+                    "negative_prompt": "negative-a",
+                    "prompt_embeds": prompt_embeds_a,
+                    "negative_prompt_embeds": negative_prompt_embeds_a,
+                    "prompt_attention_mask": prompt_attention_mask_a,
+                    "negative_prompt_attention_mask": negative_attention_mask_a,
+                },
+                sampling_params=OmniDiffusionSamplingParams(
+                    height=32,
+                    width=32,
+                    num_frames=1,
+                    frame_rate=1.0,
+                    num_inference_steps=2,
+                ),
+                request_id="ltx23-prompt-local-a",
+            ),
+            OmniDiffusionRequest(
+                prompt={
+                    "prompt": "prompt-b",
+                    "negative_prompt": "negative-b",
+                    "prompt_embeds": prompt_embeds_b,
+                    "negative_prompt_embeds": negative_prompt_embeds_b,
+                    "attention_mask": prompt_attention_mask_b,
+                    "negative_attention_mask": negative_attention_mask_b,
+                },
+                sampling_params=OmniDiffusionSamplingParams(
+                    height=32,
+                    width=32,
+                    num_frames=1,
+                    frame_rate=1.0,
+                    num_inference_steps=2,
+                ),
+                request_id="ltx23-prompt-local-b",
+            ),
+        ]
+
+        with pytest.raises(StopAtEncodePromptError):
+            pipe.forward(DiffusionRequestBatch(requests=requests))
+
+        assert captured["prompt"] is None
+        assert captured["negative_prompt"] is None
+        torch.testing.assert_close(
+            captured["prompt_embeds"],
+            torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["negative_prompt_embeds"],
+            torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["prompt_attention_mask"],
+            torch.stack([prompt_attention_mask_a, prompt_attention_mask_b], dim=0),
+        )
+        torch.testing.assert_close(
+            captured["negative_prompt_attention_mask"],
+            torch.stack([negative_attention_mask_a, negative_attention_mask_b], dim=0),
+        )
+
+    @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)])
+    @pytest.mark.parametrize(
+        ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"),
+        [(1.0, 1, 1.0), (None, 24, 24.0)],
+    )
+    def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
+        self,
+        monkeypatch,
+        cfg_rank,
+        expected_prompt_value,
+        frame_rate_input,
+        audio_sampling_rate,
+        expected_frame_rate,
+    ):
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = object.__new__(ltx23.LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.device = torch.device("cpu")
+        pipe.tokenizer_max_length = 1
+        pipe.vae_spatial_compression_ratio = 32
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.audio_sampling_rate = audio_sampling_rate
+        pipe.audio_hop_length = 1
+        pipe.audio_vae_temporal_compression_ratio = 1
+        pipe.audio_vae_mel_compression_ratio = 1
+        pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+        pipe.tokenizer = SimpleNamespace(padding_side="left")
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.audio_vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(mel_bins=2, latent_channels=1),
+        )
+
+        video_pos = torch.tensor([[[0.2, -0.3]]])
+        video_neg = torch.tensor([[[-0.4, 0.1]]])
+        audio_pos = torch.tensor([[[0.7, -0.2]]])
+        audio_neg = torch.tensor([[[0.1, 0.4]]])
+
+        class FakeCfgGroup:
+            def all_gather(self, tensor, separate_tensors=True):
+                assert separate_tensors
+                if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg):
+                    return [video_pos, video_neg]
+                if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg):
+                    return [audio_pos, audio_neg]
+                raise AssertionError(f"Unexpected gathered tensor: {tensor}")
+
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+
+        def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
+            scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
+            return torch.tensor([1.0, 0.5], device=device), 2
+
+        monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
+
+        class FakeScheduler:
+            def __init__(self, name="video", calls=None):
+                self.name = name
+                self.calls = [] if calls is None else calls
+                self.config = {
+                    "max_image_seq_len": 4096,
+                    "base_image_seq_len": 1024,
+                    "base_shift": 0.95,
+                    "max_shift": 2.05,
+                }
+                self.sigmas = torch.tensor([0.25, 0.25])
+
+            def __deepcopy__(self, memo):
+                return FakeScheduler("audio", self.calls)
+
+            def step(self, noise_pred, t, latents, return_dict=False, generator=None):
+                self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone()))
+                return (latents - noise_pred,)
+
+        class FakeConnectors:
+            def to(self, device):
+                return self
+
+            def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
+                assert padding_side == "left"
+                assert prompt_embeds.shape[0] == 2
+                return prompt_embeds, prompt_embeds, prompt_attention_mask
+
+        rope_video_fps: list[float] = []
+
+        class FakeRope:
+            def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps):
+                rope_video_fps.append(fps)
+                return torch.zeros(batch_size, num_frames * height * width, 3, device=device)
+
+            def prepare_audio_coords(self, batch_size, num_frames, device):
+                return torch.zeros(batch_size, num_frames, 1, device=device)
+
+        class FakeTransformer:
+            def __init__(self):
+                self.config = SimpleNamespace(in_channels=2)
+                self.rope = FakeRope()
+                self.audio_rope = FakeRope()
+                self.calls = []
+
+            def __call__(self, **kwargs):
+                self.calls.append(kwargs)
+                expected_prompt = torch.full((1, 1, 1), expected_prompt_value)
+                torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt)
+                torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt)
+                assert kwargs["hidden_states"].shape == (1, 1, 2)
+                assert kwargs["audio_hidden_states"].shape == (1, 1, 2)
+                if cfg_rank == 0:
+                    return video_pos, audio_pos
+                return video_neg, audio_neg
+
+        class DummyProgress:
+            def __enter__(self):
+                return self
+
+            def __exit__(self, exc_type, exc, tb):
+                return False
+
+            def update(self):
+                pass
+
+        pipe.scheduler = FakeScheduler()
+        pipe.connectors = FakeConnectors()
+        pipe.transformer = FakeTransformer()
+        object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress())
+
+        def fake_encode_prompt(**kwargs):
+            return (
+                torch.ones(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+                torch.zeros(1, 1, 1),
+                torch.ones(1, 1, dtype=torch.bool),
+            )
+
+        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
+
+        video_latents = torch.tensor([[[1.0, -2.0]]])
+        audio_latents = torch.tensor([[[0.5, 3.0]]])
+        req = OmniDiffusionRequest(
+            prompt={"prompt": "prompt", "negative_prompt": "negative"},
+            sampling_params=OmniDiffusionSamplingParams(
+                height=32,
+                width=32,
+                num_frames=1,
+                frame_rate=frame_rate_input,
+                num_inference_steps=2,
+                guidance_scale=4.0,
+                latents=video_latents,
+                audio_latents=audio_latents,
+                output_type="latent",
+            ),
+            request_id="ltx23-cfg-parallel-forward-test",
+        )
+
+        output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
+
+        expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            video_latents,
+            video_pos,
+            video_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+            audio_latents,
+            audio_pos,
+            audio_neg,
+            pipe.scheduler.sigmas[0],
+            4.0,
+        )
+        scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
+        assert scheduler_call_names == ["video", "audio", "video", "audio"]
+        assert len(pipe.transformer.calls) == 2
+        torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
+        torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
+
+        video_out, audio_out = output.output
+        torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
+        torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
+
+        # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve
+        # to the model's own 24.0 default, not crash on None; a provided rate is passed through.
+        assert rope_video_fps
+        assert all(fps == expected_frame_rate for fps in rope_video_fps)
diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py
new file mode 100644
index 000000000..9b883ff34
--- /dev/null
+++ b/tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py
@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the LTX-2.3 image-to-video pipeline."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_ltx23_request_pipe(cls):
+    pipe = object.__new__(cls)
+    torch.nn.Module.__init__(pipe)
+    pipe.device = torch.device("cpu")
+    pipe.tokenizer_max_length = 99
+    return pipe
+
+
+class TestLTX23ImageToVideoForwardStages:
+    def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline)
+        image = torch.zeros(3, 8, 8)
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt={
+                        "prompt": "make the image move",
+                        "negative_prompt": "jitter",
+                        "multi_modal_data": {"image": image},
+                    },
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=384,
+                        width=512,
+                        num_frames=25,
+                        num_inference_steps=2,
+                    ),
+                    request_id="ltx23-i2v-forward-stage-delegation",
+                )
+            ]
+        )
+        seen = {}
+
+        def fake_forward_impl(req_arg, request_inputs, **kwargs):
+            seen["req"] = req_arg
+            seen["request_inputs"] = request_inputs
+            seen["kwargs"] = kwargs
+            return ["i2v-delegated"]
+
+        object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
+
+        output = pipe.forward(req, noise_scale=0.5)
+
+        assert output == ["i2v-delegated"]
+        assert seen["req"] is req
+        assert seen["request_inputs"].prompt == ["make the image move"]
+        assert seen["request_inputs"].negative_prompt == ["jitter"]
+        assert seen["kwargs"]["image"] is image
+        assert seen["kwargs"]["noise_scale"] == 0.5
+
+    def test_denoise_timestep_kwargs_masks_video_only(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        ts = torch.tensor([2.0, 4.0])
+        denoise_ctx = SimpleNamespace(
+            conditioning_mask=torch.tensor([[1.0, 0.0]]),
+            conditioning_mask_for_model=torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+        )
+
+        kwargs = pipe._denoise_timestep_kwargs(
+            ts,
+            SimpleNamespace(cfg_parallel_ready=False),
+            denoise_ctx,
+        )
+
+        torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]]))
+        torch.testing.assert_close(kwargs["audio_timestep"], ts)
+        torch.testing.assert_close(kwargs["sigma"], ts)
+
+
+class TestLTX23ImageToVideoPipeline:
+    def test_ltx23_i2v_pipeline_reuses_ltx23_semantics(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        assert issubclass(LTX23ImageToVideoPipeline, LTX23Pipeline)
+        assert LTX23ImageToVideoPipeline.support_image_input is True
+
+    def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        image = object()
+
+        assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image
+        with pytest.raises(ValueError, match="exactly one image per prompt"):
+            LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
+
+    def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        image = torch.zeros(1, 3, 4, 4)
+        additional = {
+            "preprocessed_image": None,
+            "pixel_values": image,
+            "image": torch.ones_like(image),
+        }
+
+        assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
+
+    def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            raise AssertionError("packed I2V latents should not be noised")
+
+        monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
+
+        latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=None,
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, latents)
+
+    def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx2
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            return torch.ones(shape, device=device, dtype=dtype)
+
+        monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2, "randn_tensor", fake_randn_tensor)
+
+        latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=None,
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
+
+    def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+
+        class FakeScheduler:
+            def step(self, noise_pred, t, latents, return_dict=False):
+                return (latents + noise_pred + t,)
+
+        pipe.scheduler = FakeScheduler()
+        latents = torch.tensor([[[1.0], [2.0], [3.0]]])
+        noise_pred = torch.full_like(latents, 10.0)
+
+        out = pipe._step_video_latents_i2v(
+            noise_pred,
+            latents,
+            torch.tensor(0.5),
+            latent_num_frames=3,
+            latent_height=1,
+            latent_width=1,
+        )
+
+        torch.testing.assert_close(out[:, :1], latents[:, :1])
+        torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)
+
+
+class TestLTX23ImageToVideoModule:
+    """Test that pipeline_ltx2_3_image2video.py exposes I2V entry points."""
+
+    def test_i2v_classes_importable(self):
+        """I2V classes must be importable from the implementation module."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+
+        assert LTX23ImageToVideoPipeline is not None
+
+    def test_post_process_func_importable(self):
+        """get_ltx2_post_process_func must be importable for registry lookup."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func
+
+        assert callable(get_ltx2_post_process_func)
+
+    def test_i2v_class_matches_package_export(self):
+        """Package export must point at the I2V implementation module."""
+        from vllm_omni.diffusion.models.ltx2 import LTX23ImageToVideoPipeline as PackageExported
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import (
+            LTX23ImageToVideoPipeline as ModuleExported,
+        )
+
+        assert PackageExported is ModuleExported
diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
index 96051416b..25aae4510 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py
@@ -38,6 +38,279 @@ def _make_ltx23_pipeline(sequence_parallel_size: int = 1):
     return pipeline
 
 
+def _make_ltx23_request_pipe(cls):
+    pipe = object.__new__(cls)
+    torch.nn.Module.__init__(pipe)
+    pipe.device = torch.device("cpu")
+    pipe.tokenizer_max_length = 99
+    return pipe
+
+
+def _resolve_request_inputs_for_test(pipe, req):
+    return pipe._resolve_request_inputs(
+        req,
+        prompt=None,
+        negative_prompt=None,
+        height=None,
+        width=None,
+        num_frames=None,
+        frame_rate=None,
+        num_inference_steps=None,
+        timesteps=None,
+        guidance_scale=4.0,
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.0,
+        decode_noise_scale=None,
+        output_type="np",
+        max_sequence_length=None,
+    )
+
+
+class TestLTX23RequestParsing:
+    def test_t2v_and_i2v_share_request_input_resolution(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        prompt_embeds = torch.tensor([[1.0, 2.0]])
+        negative_prompt_embeds = torch.tensor([[3.0, 4.0]])
+        prompt_attention_mask = torch.tensor([True, False])
+        negative_attention_mask = torch.tensor([False, True])
+        video_latents = torch.ones(1, 2, 3)
+        audio_latents = torch.zeros(1, 2, 3)
+        generator = torch.Generator().manual_seed(123)
+
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt={
+                        "prompt": "shared prompt",
+                        "negative_prompt": "shared negative",
+                        "additional_information": {
+                            "prompt_embeds": prompt_embeds,
+                            "negative_prompt_embeds": negative_prompt_embeds,
+                            "attention_mask": prompt_attention_mask,
+                            "negative_attention_mask": negative_attention_mask,
+                        },
+                    },
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=384,
+                        width=512,
+                        num_frames=25,
+                        frame_rate=12.5,
+                        num_inference_steps=1,
+                        num_outputs_per_prompt=2,
+                        guidance_scale=6.0,
+                        generator=generator,
+                        latents=video_latents,
+                        extra_args={"audio_latents": audio_latents},
+                        decode_timestep=[0.1],
+                        decode_noise_scale=[0.2],
+                        output_type="latent",
+                        max_sequence_length=17,
+                    ),
+                    request_id="ltx23-shared-request-inputs",
+                )
+            ]
+        )
+
+        resolved_t2v = _resolve_request_inputs_for_test(
+            _make_ltx23_request_pipe(LTX23Pipeline),
+            req,
+        )
+        resolved_i2v = _resolve_request_inputs_for_test(
+            _make_ltx23_request_pipe(LTX23ImageToVideoPipeline),
+            req,
+        )
+
+        assert resolved_i2v.prompt == resolved_t2v.prompt is None
+        assert resolved_i2v.negative_prompt == resolved_t2v.negative_prompt is None
+        assert resolved_i2v.height == resolved_t2v.height == 384
+        assert resolved_i2v.width == resolved_t2v.width == 512
+        assert resolved_i2v.num_frames == resolved_t2v.num_frames == 25
+        assert resolved_i2v.frame_rate == resolved_t2v.frame_rate == 12.5
+        assert resolved_i2v.num_inference_steps == resolved_t2v.num_inference_steps == 2
+        assert resolved_i2v.guidance_scale == resolved_t2v.guidance_scale == 6.0
+        assert resolved_i2v.num_videos_per_prompt == resolved_t2v.num_videos_per_prompt == 2
+        assert isinstance(resolved_i2v.generator, list)
+        assert isinstance(resolved_t2v.generator, list)
+        assert [gen.initial_seed() for gen in resolved_i2v.generator] == [123, 123]
+        assert [gen.initial_seed() for gen in resolved_t2v.generator] == [123, 123]
+        assert resolved_i2v.decode_timestep == resolved_t2v.decode_timestep == [0.1]
+        assert resolved_i2v.decode_noise_scale == resolved_t2v.decode_noise_scale == [0.2]
+        assert resolved_i2v.output_type == resolved_t2v.output_type == "latent"
+        assert resolved_i2v.max_sequence_length == resolved_t2v.max_sequence_length == 17
+        torch.testing.assert_close(resolved_i2v.latents, video_latents)
+        torch.testing.assert_close(resolved_t2v.latents, video_latents)
+        torch.testing.assert_close(resolved_i2v.audio_latents, audio_latents)
+        torch.testing.assert_close(resolved_t2v.audio_latents, audio_latents)
+        torch.testing.assert_close(resolved_i2v.prompt_embeds, torch.stack([prompt_embeds]))
+        torch.testing.assert_close(resolved_t2v.prompt_embeds, torch.stack([prompt_embeds]))
+        torch.testing.assert_close(resolved_i2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds]))
+        torch.testing.assert_close(resolved_t2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds]))
+        torch.testing.assert_close(resolved_i2v.prompt_attention_mask, torch.stack([prompt_attention_mask]))
+        torch.testing.assert_close(resolved_t2v.prompt_attention_mask, torch.stack([prompt_attention_mask]))
+        torch.testing.assert_close(
+            resolved_i2v.negative_prompt_attention_mask,
+            torch.stack([negative_attention_mask]),
+        )
+        torch.testing.assert_close(
+            resolved_t2v.negative_prompt_attention_mask,
+            torch.stack([negative_attention_mask]),
+        )
+
+    def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = _make_ltx23_request_pipe(LTX23Pipeline)
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt={
+                        "prompt": "with embeds",
+                        "additional_information": {"prompt_embeds": torch.tensor([[1.0]])},
+                    },
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=384,
+                        width=512,
+                        num_frames=25,
+                        num_inference_steps=2,
+                    ),
+                    request_id="ltx23-mixed-precomputed-fields-0",
+                ),
+                OmniDiffusionRequest(
+                    prompt={"prompt": "without embeds"},
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=384,
+                        width=512,
+                        num_frames=25,
+                        num_inference_steps=2,
+                    ),
+                    request_id="ltx23-mixed-precomputed-fields-1",
+                ),
+            ]
+        )
+
+        with pytest.raises(ValueError, match="mix of provided and missing prompt_embeds"):
+            _resolve_request_inputs_for_test(pipe, req)
+
+
+class TestLTX23ForwardStages:
+    def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+
+        pipe = _make_ltx23_request_pipe(LTX23Pipeline)
+        prompt_embeds = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]])
+        negative_prompt_embeds = torch.tensor([[[-1.0], [-2.0]], [[-3.0], [-4.0]]])
+        prompt_attention_mask = torch.tensor([[True, False], [False, True]])
+        negative_prompt_attention_mask = torch.tensor([[False, False], [True, True]])
+
+        (
+            repeated_prompt_embeds,
+            repeated_prompt_attention_mask,
+            repeated_negative_prompt_embeds,
+            repeated_negative_prompt_attention_mask,
+        ) = pipe.encode_prompt(
+            prompt=None,
+            do_classifier_free_guidance=True,
+            num_videos_per_prompt=2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            device=torch.device("cpu"),
+        )
+
+        torch.testing.assert_close(
+            repeated_prompt_embeds,
+            torch.tensor([[[1.0], [2.0]], [[1.0], [2.0]], [[3.0], [4.0]], [[3.0], [4.0]]]),
+        )
+        torch.testing.assert_close(
+            repeated_negative_prompt_embeds,
+            torch.tensor([[[-1.0], [-2.0]], [[-1.0], [-2.0]], [[-3.0], [-4.0]], [[-3.0], [-4.0]]]),
+        )
+        torch.testing.assert_close(
+            repeated_prompt_attention_mask,
+            torch.tensor([[True, False], [True, False], [False, True], [False, True]]),
+        )
+        torch.testing.assert_close(
+            repeated_negative_prompt_attention_mask,
+            torch.tensor([[False, False], [False, False], [True, True], [True, True]]),
+        )
+
+    def test_t2v_forward_delegates_to_shared_forward_impl(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        pipe = _make_ltx23_request_pipe(LTX23Pipeline)
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="stage refactor prompt",
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=384,
+                        width=512,
+                        num_frames=25,
+                        frame_rate=16,
+                        num_inference_steps=3,
+                        guidance_scale=4.5,
+                        output_type="latent",
+                    ),
+                    request_id="ltx23-t2v-forward-stage-delegation",
+                )
+            ]
+        )
+        sigmas = [1.0, 0.5]
+        timesteps = [1000, 500]
+        attention_kwargs = {"scale": 1.0}
+        seen = {}
+
+        def fake_forward_impl(req_arg, request_inputs, **kwargs):
+            seen["req"] = req_arg
+            seen["request_inputs"] = request_inputs
+            seen["kwargs"] = kwargs
+            return ["delegated"]
+
+        object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
+
+        output = pipe.forward(
+            req,
+            sigmas=sigmas,
+            timesteps=timesteps,
+            noise_scale=0.25,
+            attention_kwargs=attention_kwargs,
+        )
+
+        assert output == ["delegated"]
+        assert seen["req"] is req
+        assert seen["request_inputs"].height == 384
+        assert seen["request_inputs"].width == 512
+        assert seen["request_inputs"].num_frames == 25
+        assert seen["request_inputs"].frame_rate == 16.0
+        assert seen["request_inputs"].guidance_scale == 4.5
+        assert seen["request_inputs"].output_type == "latent"
+        assert seen["kwargs"] == {
+            "noise_scale": 0.25,
+            "sigmas": sigmas,
+            "timesteps": timesteps,
+            "attention_kwargs": attention_kwargs,
+        }
+
+
 class TestPipelineIndependence:
     """Verify LTX23Pipeline is fully independent from LTX2Pipeline."""
 
@@ -65,13 +338,6 @@ class TestPipelineIndependence:
 
         assert issubclass(LTX23Pipeline, ProgressBarMixin)
 
-    def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
-        """LTX23Pipeline must use the shared CFG parallel implementation."""
-        from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-
-        assert issubclass(LTX23Pipeline, CFGParallelMixin)
-
     def test_ltx23_pipeline_declares_offload_components(self):
         """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
@@ -105,661 +371,40 @@ class TestPipelineIndependence:
         assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin)
 
 
-class TestLTX23VaeDecodeParallel:
-    """Test LTX-2.3 video VAE tiled parallel helpers without loading weights."""
+class TestLTX23DecodeConditioning:
+    def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value
 
-    def test_ltx23_video_vae_is_distributed_tile_only_class(self):
-        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
-            DistributedAutoencoderKLLTX2Video,
-        )
-        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
+        assert _expand_per_prompt_decode_value(
+            [0.1, 0.2],
+            prompt_batch_size=2,
+            effective_batch_size=4,
+            field_name="decode_timestep",
+        ) == [0.1, 0.1, 0.2, 0.2]
+        assert _expand_per_prompt_decode_value(
+            [0.3],
+            prompt_batch_size=2,
+            effective_batch_size=4,
+            field_name="decode_timestep",
+        ) == [0.3, 0.3, 0.3, 0.3]
+        assert _expand_per_prompt_decode_value(
+            [0.1, 0.2, 0.3, 0.4],
+            prompt_batch_size=2,
+            effective_batch_size=4,
+            field_name="decode_timestep",
+        ) == [0.1, 0.2, 0.3, 0.4]
 
-        assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin)
-        assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split")
+    def test_decode_conditioning_rejects_ambiguous_lengths(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value
 
-    def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self):
-        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
-            DistributedAutoencoderKLLTX2Video,
-        )
-
-        vae = SimpleNamespace(
-            spatial_compression_ratio=32,
-            tile_sample_min_height=512,
-            tile_sample_min_width=512,
-            tile_sample_stride_height=448,
-            tile_sample_stride_width=448,
-            temporal_compression_ratio=8,
-            dtype=torch.float32,
-        )
-
-        z = torch.zeros(1, 2, 5, 16, 24)
-        tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z)
-
-        assert grid_spec.grid_shape == (2, 2)
-        assert grid_spec.split_dims == (3, 4)
-        assert grid_spec.tile_spec["sample_height"] == 512
-        assert grid_spec.tile_spec["sample_width"] == 768
-        assert grid_spec.tile_spec["blend_height"] == 64
-        assert grid_spec.tile_spec["blend_width"] == 64
-        assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512)
-        assert grid_spec.tile_spec["tile_output_shapes"] == {
-            0: (1, 3, 33, 512, 512),
-            1: (1, 3, 33, 512, 320),
-            2: (1, 3, 33, 64, 512),
-            3: (1, 3, 33, 64, 320),
-        }
-        assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)]
-        assert [tuple(task.tensor.shape) for task in tasks] == [
-            (1, 2, 5, 16, 16),
-            (1, 2, 5, 16, 10),
-            (1, 2, 5, 2, 16),
-            (1, 2, 5, 2, 10),
-        ]
-        assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10]
-
-    def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self):
-        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
-            DistributedAutoencoderKLLTX2Video,
-        )
-        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec
-
-        class FakeVae:
-            def __init__(self):
-                self.blend_calls = []
-
-            def clear_cache(self):
-                pass
-
-            def blend_v(self, _previous, current, blend_height):
-                self.blend_calls.append(("v", blend_height))
-                return current
-
-            def blend_h(self, _previous, current, blend_width):
-                self.blend_calls.append(("h", blend_width))
-                return current
-
-        fake_vae = FakeVae()
-        grid_spec = GridSpec(
-            split_dims=(3, 4),
-            grid_shape=(2, 2),
-            tile_spec={
-                "sample_height": 10,
-                "sample_width": 10,
-                "blend_height": 1,
-                "blend_width": 2,
-                "tile_sample_stride_height": 5,
-                "tile_sample_stride_width": 5,
-            },
-        )
-        tiles = {
-            (0, 0): torch.full((1, 3, 2, 6, 6), 1.0),
-            (0, 1): torch.full((1, 3, 2, 6, 6), 2.0),
-            (1, 0): torch.full((1, 3, 2, 6, 6), 3.0),
-            (1, 1): torch.full((1, 3, 2, 6, 6), 4.0),
-        }
-
-        merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec)
-
-        assert merged.shape == (1, 3, 2, 10, 10)
-        assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)]
-        torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5))
-        torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0))
-        torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0))
-        torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0))
-
-    def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self):
-        from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2
-
-        z = torch.zeros(1, 2, 1, 16, 24)
-        expected = torch.ones(1, 3, 1, 512, 768)
-        seen = {}
-
-        class FakeExecutor:
-            def execute(self, tensor, operator, broadcast_result=True):
-                seen["tensor"] = tensor
-                seen["operator"] = operator
-                seen["broadcast_result"] = broadcast_result
-                return expected
-
-        vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True)
-        vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae)
-        vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae)
-        vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae)
-
-        output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode(
-            vae,
-            z,
-            temb=torch.tensor(0.5),
-            return_dict=False,
-        )
-
-        assert len(output) == 1
-        assert output[0] is expected
-        assert seen["tensor"] is z
-        assert seen["broadcast_result"] is False
-        assert seen["operator"].split.__name__ == "tile_split"
-        assert seen["operator"].merge.__name__ == "tile_merge"
-
-    def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self):
-        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor
-        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
-            DistributedOperator,
-            GridSpec,
-            TileTask,
-        )
-
-        z = torch.zeros(1, 1, 1, 1, 1)
-        tile_output_shapes = {
-            0: (1, 1, 1, 2, 2),
-            1: (1, 1, 1, 2, 1),
-            2: (1, 1, 1, 1, 2),
-            3: (1, 1, 1, 1, 1),
-        }
-        tasks = [
-            TileTask(0, (0, 0), z, workload=4),
-            TileTask(1, (0, 1), z, workload=2),
-            TileTask(2, (1, 0), z, workload=2),
-            TileTask(3, (1, 1), z, workload=1),
-        ]
-        grid_spec = GridSpec(
-            split_dims=(3, 4),
-            grid_shape=(2, 2),
-            tile_spec={
-                "max_tile_output_shape": (1, 1, 1, 2, 2),
-                "tile_output_shapes": tile_output_shapes,
-            },
-            output_dtype=torch.float32,
-        )
-        seen = {}
-
-        def exec_tile(task):
-            return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
-
-        def merge_tiles(coord_tensor_map, passed_grid_spec):
-            seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()}
-            assert passed_grid_spec is grid_spec
-            return torch.stack(
-                [
-                    coord_tensor_map[(0, 0)].flatten()[0],
-                    coord_tensor_map[(0, 1)].flatten()[0],
-                    coord_tensor_map[(1, 0)].flatten()[0],
-                    coord_tensor_map[(1, 1)].flatten()[0],
-                ]
+        with pytest.raises(ValueError, match="decode_timestep"):
+            _expand_per_prompt_decode_value(
+                [0.1, 0.2, 0.3],
+                prompt_batch_size=2,
+                effective_batch_size=4,
+                field_name="decode_timestep",
             )
 
-        operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles)
-
-        rank0_executor = object.__new__(LTX2VaeExecutor)
-        rank0_executor.parallel_size = 2
-        rank0_executor.world_size = 2
-        rank0_executor.rank = 0
-
-        def gather_rank0(local_tile_tensor):
-            assigned = rank0_executor._balance_tasks(tasks, 2)
-            rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]]
-            rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta(
-                rank1_results,
-                list(local_tile_tensor.shape),
-                z.device,
-                torch.float32,
-            )
-            seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape)
-            return [local_tile_tensor, rank1_tile_tensor]
-
-        def fail_final_sync(*_args, **_kwargs):
-            raise AssertionError("broadcast_result=False should not sync the final result")
-
-        rank0_executor.gather_tensors = gather_rank0
-        rank0_executor._sync_final_result = fail_final_sync
-
-        rank0_result = rank0_executor.execute(z, operator, broadcast_result=False)
-
-        torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0]))
-        assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2)
-        assert seen["merged_shapes"] == {
-            (0, 0): (1, 1, 1, 2, 2),
-            (0, 1): (1, 1, 1, 2, 1),
-            (1, 0): (1, 1, 1, 1, 2),
-            (1, 1): (1, 1, 1, 1, 1),
-        }
-
-        non_rank0_executor = object.__new__(LTX2VaeExecutor)
-        non_rank0_executor.parallel_size = 2
-        non_rank0_executor.world_size = 2
-        non_rank0_executor.rank = 1
-
-        def gather_rank1(local_tile_tensor):
-            seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape)
-            return None
-
-        def fail_non_rank0_merge(*_args, **_kwargs):
-            raise AssertionError("non-rank0 should not merge gathered tiles")
-
-        non_rank0_executor.gather_tensors = gather_rank1
-        non_rank0_executor._sync_final_result = fail_final_sync
-
-        empty_result = non_rank0_executor.execute(
-            z,
-            DistributedOperator(
-                split=lambda _z: (tasks, grid_spec),
-                exec=exec_tile,
-                merge=fail_non_rank0_merge,
-            ),
-            broadcast_result=False,
-        )
-
-        assert tuple(empty_result.shape) == (0,)
-        assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2)
-
-
-class TestCFGParallelHelpers:
-    """Test LTX-2.3 CFG helper math without loading model weights."""
-
-    def test_combine_cfg_noise_matches_x0_space_formula(self):
-        import torch
-
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-
-        pipe = object.__new__(LTX23Pipeline)
-        video_sample = torch.tensor([[[1.0, -2.0]]])
-        audio_sample = torch.tensor([[[0.5, 3.0]]])
-        video_pos = torch.tensor([[[0.2, -0.3]]])
-        video_neg = torch.tensor([[[-0.4, 0.1]]])
-        audio_pos = torch.tensor([[[0.7, -0.2]]])
-        audio_neg = torch.tensor([[[0.1, 0.4]]])
-        video_sigma = torch.tensor(0.25)
-        audio_sigma = torch.tensor(0.5)
-        scale = 4.0
-
-        video_combined, audio_combined = pipe.combine_cfg_noise(
-            (video_pos, audio_pos),
-            (video_neg, audio_neg),
-            scale,
-            video_latents=video_sample,
-            audio_latents=audio_sample,
-            video_sigma=video_sigma,
-            audio_sigma=audio_sigma,
-        )
-
-        x0_video_cond = video_sample - video_pos * video_sigma
-        x0_video_uncond = video_sample - video_neg * video_sigma
-        x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond)
-        expected_video = (video_sample - x0_video_guided) / video_sigma
-
-        x0_audio_cond = audio_sample - audio_pos * audio_sigma
-        x0_audio_uncond = audio_sample - audio_neg * audio_sigma
-        x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond)
-        expected_audio = (audio_sample - x0_audio_guided) / audio_sigma
-        assert torch.allclose(video_combined, expected_video)
-        assert torch.allclose(audio_combined, expected_audio)
-
-    def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
-
-        pipe = object.__new__(ltx23.LTX23Pipeline)
-        video_sample = torch.tensor([[[1.0, -2.0]]])
-        audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
-        video_pos = torch.tensor([[[0.2, -0.3]]])
-        video_neg = torch.tensor([[[-0.4, 0.1]]])
-        audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
-        audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
-        video_sigma = torch.tensor(0.25)
-        audio_sigma = torch.tensor(0.5)
-        scale = 4.0
-
-        class FakeCfgGroup:
-            def all_gather(self, tensor, separate_tensors=True):
-                assert separate_tensors
-                if tensor.shape == video_pos.shape:
-                    return [video_pos, video_neg]
-                return [audio_pos, audio_neg]
-
-        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
-        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
-
-        expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
-            video_sample,
-            video_pos,
-            video_neg,
-            video_sigma,
-            scale,
-        )
-        expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
-            audio_sample,
-            audio_pos,
-            audio_neg,
-            audio_sigma,
-            scale,
-        )
-
-        for rank, expected_branch in ((0, "positive"), (1, "negative")):
-            calls = []
-            monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
-
-            def fake_predict_noise(**kwargs):
-                calls.append(kwargs["branch"])
-                if kwargs["branch"] == "positive":
-                    return video_pos, audio_pos
-                return video_neg, audio_neg
-
-            object.__setattr__(pipe, "predict_noise", fake_predict_noise)
-            video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
-                true_cfg_scale=scale,
-                positive_kwargs={"branch": "positive"},
-                negative_kwargs={"branch": "negative"},
-                cfg_normalize=False,
-                video_latents=video_sample,
-                audio_latents=audio_sample,
-                video_sigma=video_sigma,
-                audio_sigma=audio_sigma,
-            )
-
-            assert calls == [expected_branch]
-            torch.testing.assert_close(video_combined, expected_video)
-            torch.testing.assert_close(audio_combined, expected_audio)
-
-        assert "_cfg_video_latents" not in pipe.__dict__
-        assert "_cfg_audio_latents" not in pipe.__dict__
-
-
-class TestCFGParallelForwardPath:
-    """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
-
-    def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
-        from vllm_omni.diffusion.request import OmniDiffusionRequest
-        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-        pipe = object.__new__(ltx23.LTX23Pipeline)
-        torch.nn.Module.__init__(pipe)
-        pipe.device = torch.device("cpu")
-        pipe.tokenizer_max_length = 4
-        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 1)
-
-        class StopAtEncodePromptError(Exception):
-            pass
-
-        captured = {}
-
-        def fake_encode_prompt(**kwargs):
-            captured.update(kwargs)
-            raise StopAtEncodePromptError
-
-        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
-
-        prompt_embeds_a = torch.zeros(2, 3)
-        prompt_embeds_b = torch.ones(2, 3)
-        negative_prompt_embeds_a = torch.full((2, 3), 2.0)
-        negative_prompt_embeds_b = torch.full((2, 3), 3.0)
-        prompt_attention_mask_a = torch.tensor([True, True])
-        prompt_attention_mask_b = torch.tensor([True, False])
-        negative_attention_mask_a = torch.tensor([False, True])
-        negative_attention_mask_b = torch.tensor([False, False])
-
-        requests = [
-            OmniDiffusionRequest(
-                prompt={
-                    "prompt": "prompt-a",
-                    "negative_prompt": "negative-a",
-                    "prompt_embeds": prompt_embeds_a,
-                    "negative_prompt_embeds": negative_prompt_embeds_a,
-                    "prompt_attention_mask": prompt_attention_mask_a,
-                    "negative_prompt_attention_mask": negative_attention_mask_a,
-                },
-                sampling_params=OmniDiffusionSamplingParams(
-                    height=32,
-                    width=32,
-                    num_frames=1,
-                    frame_rate=1.0,
-                    num_inference_steps=2,
-                ),
-                request_id="ltx23-prompt-local-a",
-            ),
-            OmniDiffusionRequest(
-                prompt={
-                    "prompt": "prompt-b",
-                    "negative_prompt": "negative-b",
-                    "prompt_embeds": prompt_embeds_b,
-                    "negative_prompt_embeds": negative_prompt_embeds_b,
-                    "attention_mask": prompt_attention_mask_b,
-                    "negative_attention_mask": negative_attention_mask_b,
-                },
-                sampling_params=OmniDiffusionSamplingParams(
-                    height=32,
-                    width=32,
-                    num_frames=1,
-                    frame_rate=1.0,
-                    num_inference_steps=2,
-                ),
-                request_id="ltx23-prompt-local-b",
-            ),
-        ]
-
-        with pytest.raises(StopAtEncodePromptError):
-            pipe.forward(DiffusionRequestBatch(requests=requests))
-
-        assert captured["prompt"] is None
-        assert captured["negative_prompt"] is None
-        torch.testing.assert_close(
-            captured["prompt_embeds"],
-            torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0),
-        )
-        torch.testing.assert_close(
-            captured["negative_prompt_embeds"],
-            torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0),
-        )
-        torch.testing.assert_close(
-            captured["prompt_attention_mask"],
-            torch.stack([prompt_attention_mask_a, prompt_attention_mask_b], dim=0),
-        )
-        torch.testing.assert_close(
-            captured["negative_prompt_attention_mask"],
-            torch.stack([negative_attention_mask_a, negative_attention_mask_b], dim=0),
-        )
-
-    @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)])
-    @pytest.mark.parametrize(
-        ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"),
-        [(1.0, 1, 1.0), (None, 24, 24.0)],
-    )
-    def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
-        self,
-        monkeypatch,
-        cfg_rank,
-        expected_prompt_value,
-        frame_rate_input,
-        audio_sampling_rate,
-        expected_frame_rate,
-    ):
-        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
-        from vllm_omni.diffusion.request import OmniDiffusionRequest
-        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-        pipe = object.__new__(ltx23.LTX23Pipeline)
-        torch.nn.Module.__init__(pipe)
-        pipe.device = torch.device("cpu")
-        pipe.tokenizer_max_length = 1
-        pipe.vae_spatial_compression_ratio = 32
-        pipe.vae_temporal_compression_ratio = 1
-        pipe.transformer_spatial_patch_size = 1
-        pipe.transformer_temporal_patch_size = 1
-        pipe.audio_sampling_rate = audio_sampling_rate
-        pipe.audio_hop_length = 1
-        pipe.audio_vae_temporal_compression_ratio = 1
-        pipe.audio_vae_mel_compression_ratio = 1
-        pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
-        pipe.tokenizer = SimpleNamespace(padding_side="left")
-        pipe.vae = SimpleNamespace(
-            latents_mean=torch.zeros(2),
-            latents_std=torch.ones(2),
-            config=SimpleNamespace(scaling_factor=1.0),
-        )
-        pipe.audio_vae = SimpleNamespace(
-            latents_mean=torch.zeros(2),
-            latents_std=torch.ones(2),
-            config=SimpleNamespace(mel_bins=2, latent_channels=1),
-        )
-
-        video_pos = torch.tensor([[[0.2, -0.3]]])
-        video_neg = torch.tensor([[[-0.4, 0.1]]])
-        audio_pos = torch.tensor([[[0.7, -0.2]]])
-        audio_neg = torch.tensor([[[0.1, 0.4]]])
-
-        class FakeCfgGroup:
-            def all_gather(self, tensor, separate_tensors=True):
-                assert separate_tensors
-                if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg):
-                    return [video_pos, video_neg]
-                if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg):
-                    return [audio_pos, audio_neg]
-                raise AssertionError(f"Unexpected gathered tensor: {tensor}")
-
-        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
-        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
-        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
-
-        def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
-            scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
-            return torch.tensor([1.0, 0.5], device=device), 2
-
-        monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
-
-        class FakeScheduler:
-            def __init__(self, name="video", calls=None):
-                self.name = name
-                self.calls = [] if calls is None else calls
-                self.config = {
-                    "max_image_seq_len": 4096,
-                    "base_image_seq_len": 1024,
-                    "base_shift": 0.95,
-                    "max_shift": 2.05,
-                }
-                self.sigmas = torch.tensor([0.25, 0.25])
-
-            def __deepcopy__(self, memo):
-                return FakeScheduler("audio", self.calls)
-
-            def step(self, noise_pred, t, latents, return_dict=False, generator=None):
-                self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone()))
-                return (latents - noise_pred,)
-
-        class FakeConnectors:
-            def to(self, device):
-                return self
-
-            def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
-                assert padding_side == "left"
-                assert prompt_embeds.shape[0] == 2
-                return prompt_embeds, prompt_embeds, prompt_attention_mask
-
-        rope_video_fps: list[float] = []
-
-        class FakeRope:
-            def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps):
-                rope_video_fps.append(fps)
-                return torch.zeros(batch_size, num_frames * height * width, 3, device=device)
-
-            def prepare_audio_coords(self, batch_size, num_frames, device):
-                return torch.zeros(batch_size, num_frames, 1, device=device)
-
-        class FakeTransformer:
-            def __init__(self):
-                self.config = SimpleNamespace(in_channels=2)
-                self.rope = FakeRope()
-                self.audio_rope = FakeRope()
-                self.calls = []
-
-            def __call__(self, **kwargs):
-                self.calls.append(kwargs)
-                expected_prompt = torch.full((1, 1, 1), expected_prompt_value)
-                torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt)
-                torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt)
-                assert kwargs["hidden_states"].shape == (1, 1, 2)
-                assert kwargs["audio_hidden_states"].shape == (1, 1, 2)
-                if cfg_rank == 0:
-                    return video_pos, audio_pos
-                return video_neg, audio_neg
-
-        class DummyProgress:
-            def __enter__(self):
-                return self
-
-            def __exit__(self, exc_type, exc, tb):
-                return False
-
-            def update(self):
-                pass
-
-        pipe.scheduler = FakeScheduler()
-        pipe.connectors = FakeConnectors()
-        pipe.transformer = FakeTransformer()
-        object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress())
-
-        def fake_encode_prompt(**kwargs):
-            return (
-                torch.ones(1, 1, 1),
-                torch.ones(1, 1, dtype=torch.bool),
-                torch.zeros(1, 1, 1),
-                torch.ones(1, 1, dtype=torch.bool),
-            )
-
-        object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
-
-        video_latents = torch.tensor([[[1.0, -2.0]]])
-        audio_latents = torch.tensor([[[0.5, 3.0]]])
-        req = OmniDiffusionRequest(
-            prompt={"prompt": "prompt", "negative_prompt": "negative"},
-            sampling_params=OmniDiffusionSamplingParams(
-                height=32,
-                width=32,
-                num_frames=1,
-                frame_rate=frame_rate_input,
-                num_inference_steps=2,
-                guidance_scale=4.0,
-                latents=video_latents,
-                audio_latents=audio_latents,
-                output_type="latent",
-            ),
-            request_id="ltx23-cfg-parallel-forward-test",
-        )
-
-        output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
-
-        expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
-            video_latents,
-            video_pos,
-            video_neg,
-            pipe.scheduler.sigmas[0],
-            4.0,
-        )
-        expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
-            audio_latents,
-            audio_pos,
-            audio_neg,
-            pipe.scheduler.sigmas[0],
-            4.0,
-        )
-        scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
-        assert scheduler_call_names == ["video", "audio", "video", "audio"]
-        assert len(pipe.transformer.calls) == 2
-        torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
-
-        video_out, audio_out = output.output
-        torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
-        torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
-
-        # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve
-        # to the model's own 24.0 default, not crash on None; a provided rate is passed through.
-        assert rope_video_fps
-        assert all(fps == expected_frame_rate for fps in rope_video_fps)
-
 
 class TestRegistryIntegration:
     """Verify all LTX-2.3 pipeline variants are registered."""
@@ -814,15 +459,6 @@ class TestRegistryIntegration:
         for name in expected:
             assert name not in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS"
 
-    def test_ltx2_transformer_has_dit_cache_config(self):
-        """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
-        from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
-        from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
-
-        adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
-        assert isinstance(adapter_config, CacheDiTAdapterConfig)
-        assert adapter_config.has_separate_cfg
-
 
 class TestVocoderSampleRateDetection:
     """Test _detect_vocoder_output_sample_rate logic."""
@@ -912,31 +548,6 @@ class TestPostProcessFunction:
         assert "audio_sample_rate" not in result
 
 
-class TestReExportModule:
-    """Test that pipeline_ltx2_3_image2video.py correctly re-exports."""
-
-    def test_i2v_classes_importable(self):
-        """I2V classes must be importable from the re-export module."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
-
-        assert LTX23ImageToVideoPipeline is not None
-
-    def test_post_process_func_importable(self):
-        """get_ltx2_post_process_func must be importable from re-export module."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func
-
-        assert callable(get_ltx2_post_process_func)
-
-    def test_i2v_classes_are_same_as_direct_import(self):
-        """Re-exported classes must be the same objects as direct imports."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23ImageToVideoPipeline as Direct
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import (
-            LTX23ImageToVideoPipeline as ReExported,
-        )
-
-        assert Direct is ReExported
-
-
 class TestInitExports:
     """Test that __init__.py exports all LTX-2.3 classes."""
 
diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py b/tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py
new file mode 100644
index 000000000..9b53bb4f8
--- /dev/null
+++ b/tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py
@@ -0,0 +1,282 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestLTX23OutputRank:
+    def test_single_process_rank_is_output_rank(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
+
+        monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: False)
+
+        assert ltx23._is_output_rank() is True
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is True
+
+    def test_non_output_rank_skips_decode_unless_vae_decode_is_distributed(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
+
+        monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx23.torch.distributed, "get_rank", lambda: 1)
+
+        assert ltx23._is_output_rank() is False
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is False
+        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: True)) is True
+
+
+class TestLTX23VaeDistributedDecode:
+    """Test LTX-2.3 distributed VAE helpers without loading weights."""
+
+    def test_ltx23_video_vae_is_distributed_tile_only_class(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
+
+        assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin)
+        assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split")
+
+    def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+            DistributedOperator,
+            GridSpec,
+            TileTask,
+        )
+
+        z = torch.zeros(1, 1, 1, 1, 1)
+        tile_output_shapes = {
+            0: (1, 1, 1, 2, 2),
+            1: (1, 1, 1, 2, 1),
+            2: (1, 1, 1, 1, 2),
+            3: (1, 1, 1, 1, 1),
+        }
+        tasks = [
+            TileTask(0, (0, 0), z, workload=4),
+            TileTask(1, (0, 1), z, workload=2),
+            TileTask(2, (1, 0), z, workload=2),
+            TileTask(3, (1, 1), z, workload=1),
+        ]
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "max_tile_output_shape": (1, 1, 1, 2, 2),
+                "tile_output_shapes": tile_output_shapes,
+            },
+            output_dtype=torch.float32,
+        )
+        seen = {}
+
+        def exec_tile(task):
+            return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
+
+        def merge_tiles(coord_tensor_map, passed_grid_spec):
+            seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()}
+            assert passed_grid_spec is grid_spec
+            return torch.stack(
+                [
+                    coord_tensor_map[(0, 0)].flatten()[0],
+                    coord_tensor_map[(0, 1)].flatten()[0],
+                    coord_tensor_map[(1, 0)].flatten()[0],
+                    coord_tensor_map[(1, 1)].flatten()[0],
+                ]
+            )
+
+        operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles)
+
+        rank0_executor = object.__new__(LTX2VaeExecutor)
+        rank0_executor.parallel_size = 2
+        rank0_executor.world_size = 2
+        rank0_executor.rank = 0
+
+        def gather_rank0(local_tile_tensor):
+            assigned = rank0_executor._balance_tasks(tasks, 2)
+            rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]]
+            rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta(
+                rank1_results,
+                list(local_tile_tensor.shape),
+                z.device,
+                torch.float32,
+            )
+            seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape)
+            return [local_tile_tensor, rank1_tile_tensor]
+
+        def fail_final_sync(*_args, **_kwargs):
+            raise AssertionError("broadcast_result=False should not sync the final result")
+
+        rank0_executor.gather_tensors = gather_rank0
+        rank0_executor._sync_final_result = fail_final_sync
+
+        rank0_result = rank0_executor.execute(z, operator, broadcast_result=False)
+
+        torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0]))
+        assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2)
+        assert seen["merged_shapes"] == {
+            (0, 0): (1, 1, 1, 2, 2),
+            (0, 1): (1, 1, 1, 2, 1),
+            (1, 0): (1, 1, 1, 1, 2),
+            (1, 1): (1, 1, 1, 1, 1),
+        }
+
+        non_rank0_executor = object.__new__(LTX2VaeExecutor)
+        non_rank0_executor.parallel_size = 2
+        non_rank0_executor.world_size = 2
+        non_rank0_executor.rank = 1
+
+        def gather_rank1(local_tile_tensor):
+            seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape)
+            return None
+
+        def fail_non_rank0_merge(*_args, **_kwargs):
+            raise AssertionError("non-rank0 should not merge gathered tiles")
+
+        non_rank0_executor.gather_tensors = gather_rank1
+        non_rank0_executor._sync_final_result = fail_final_sync
+
+        empty_result = non_rank0_executor.execute(
+            z,
+            DistributedOperator(
+                split=lambda _z: (tasks, grid_spec),
+                exec=exec_tile,
+                merge=fail_non_rank0_merge,
+            ),
+            broadcast_result=False,
+        )
+
+        assert tuple(empty_result.shape) == (0,)
+        assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2)
+
+
+class TestLTX23VaeTiling:
+    """Test LTX-2.3 video VAE tile helpers without loading weights."""
+
+    def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+
+        vae = SimpleNamespace(
+            spatial_compression_ratio=32,
+            tile_sample_min_height=512,
+            tile_sample_min_width=512,
+            tile_sample_stride_height=448,
+            tile_sample_stride_width=448,
+            temporal_compression_ratio=8,
+            dtype=torch.float32,
+        )
+
+        z = torch.zeros(1, 2, 5, 16, 24)
+        tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z)
+
+        assert grid_spec.grid_shape == (2, 2)
+        assert grid_spec.split_dims == (3, 4)
+        assert grid_spec.tile_spec["sample_height"] == 512
+        assert grid_spec.tile_spec["sample_width"] == 768
+        assert grid_spec.tile_spec["blend_height"] == 64
+        assert grid_spec.tile_spec["blend_width"] == 64
+        assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512)
+        assert grid_spec.tile_spec["tile_output_shapes"] == {
+            0: (1, 3, 33, 512, 512),
+            1: (1, 3, 33, 512, 320),
+            2: (1, 3, 33, 64, 512),
+            3: (1, 3, 33, 64, 320),
+        }
+        assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)]
+        assert [tuple(task.tensor.shape) for task in tasks] == [
+            (1, 2, 5, 16, 16),
+            (1, 2, 5, 16, 10),
+            (1, 2, 5, 2, 16),
+            (1, 2, 5, 2, 10),
+        ]
+        assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10]
+
+    def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self):
+        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+            DistributedAutoencoderKLLTX2Video,
+        )
+        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec
+
+        class FakeVae:
+            def __init__(self):
+                self.blend_calls = []
+
+            def clear_cache(self):
+                pass
+
+            def blend_v(self, _previous, current, blend_height):
+                self.blend_calls.append(("v", blend_height))
+                return current
+
+            def blend_h(self, _previous, current, blend_width):
+                self.blend_calls.append(("h", blend_width))
+                return current
+
+        fake_vae = FakeVae()
+        grid_spec = GridSpec(
+            split_dims=(3, 4),
+            grid_shape=(2, 2),
+            tile_spec={
+                "sample_height": 10,
+                "sample_width": 10,
+                "blend_height": 1,
+                "blend_width": 2,
+                "tile_sample_stride_height": 5,
+                "tile_sample_stride_width": 5,
+            },
+        )
+        tiles = {
+            (0, 0): torch.full((1, 3, 2, 6, 6), 1.0),
+            (0, 1): torch.full((1, 3, 2, 6, 6), 2.0),
+            (1, 0): torch.full((1, 3, 2, 6, 6), 3.0),
+            (1, 1): torch.full((1, 3, 2, 6, 6), 4.0),
+        }
+
+        merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec)
+
+        assert merged.shape == (1, 3, 2, 10, 10)
+        assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)]
+        torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5))
+        torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0))
+        torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0))
+
+    def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self):
+        from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2
+
+        z = torch.zeros(1, 2, 1, 16, 24)
+        expected = torch.ones(1, 3, 1, 512, 768)
+        seen = {}
+
+        class FakeExecutor:
+            def execute(self, tensor, operator, broadcast_result=True):
+                seen["tensor"] = tensor
+                seen["operator"] = operator
+                seen["broadcast_result"] = broadcast_result
+                return expected
+
+        vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True)
+        vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae)
+        vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae)
+        vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae)
+
+        output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode(
+            vae,
+            z,
+            temb=torch.tensor(0.5),
+            return_dict=False,
+        )
+
+        assert len(output) == 1
+        assert output[0] is expected
+        assert seen["tensor"] is z
+        assert seen["broadcast_result"] is False
+        assert seen["operator"].split.__name__ == "tile_split"
+        assert seen["operator"].merge.__name__ == "tile_merge"
diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py
index 3f85a6b17..7f23b4733 100644
--- a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py
+++ b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py
@@ -1,4 +1,3 @@
-import inspect
 from types import SimpleNamespace
 
 import pytest
@@ -193,6 +192,7 @@ def _make_request_batch_prompt_sampling(**overrides):
         "true_cfg_scale": None,
         "guidance_scale_provided": False,
         "guidance_scale": 1.0,
+        "output_type": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
@@ -344,16 +344,3 @@ def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_cl
 
     with pytest.raises(AssertionError, match="text encoder should not run"):
         pipeline.encode_prompt(prompt="short prompt")
-
-
-@pytest.mark.parametrize(
-    "pipeline_class",
-    [
-        QwenImagePipeline,
-        QwenImageLayeredPipeline,
-        QwenImageEditPipeline,
-        QwenImageEditPlusPipeline,
-    ],
-)
-def test_forward_max_sequence_length_default_is_1024(pipeline_class: type):
-    assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024
diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py
new file mode 100644
index 000000000..62d17f705
--- /dev/null
+++ b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py
@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image as qwen_image
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_qwen_image_postprocess_preserves_output_envelope(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path,
+) -> None:
+    model_path = tmp_path / "model"
+    vae_path = model_path / "vae"
+    vae_path.mkdir(parents=True)
+    (vae_path / "config.json").write_text(json.dumps({}), encoding="utf-8")
+
+    processed = object()
+    captured: dict[str, object] = {}
+
+    class FakeImageProcessor:
+        def __init__(self, *, vae_scale_factor: int) -> None:
+            captured["vae_scale_factor"] = vae_scale_factor
+
+        def postprocess(self, images: torch.Tensor) -> object:
+            captured["images"] = images
+            return processed
+
+    monkeypatch.setattr(qwen_image, "VaeImageProcessor", FakeImageProcessor)
+    postprocess = qwen_image.get_qwen_image_post_process_func(SimpleNamespace(model=str(model_path)))
+
+    image = torch.zeros(1, 3, 4, 4)
+    output = postprocess(
+        {
+            "payload": {"image": image},
+            "metadata": {"prompt_embeddings": {"present": True}},
+        }
+    )
+
+    assert captured["vae_scale_factor"] == 16
+    assert captured["images"] is image
+    assert output == {
+        "payload": {"image": processed},
+        "metadata": {"prompt_embeddings": {"present": True}},
+    }
diff --git a/tests/diffusion/models/sd3/test_sd3_pipeline.py b/tests/diffusion/models/sd3/test_sd3_pipeline.py
index 3164e5dda..4b0bf843f 100644
--- a/tests/diffusion/models/sd3/test_sd3_pipeline.py
+++ b/tests/diffusion/models/sd3/test_sd3_pipeline.py
@@ -21,6 +21,7 @@ def _make_sd3_sampling(**overrides):
         "generator": None,
         "latents": None,
         "guidance_scale": 4.0,
+        "output_type": None,
     }
     values.update(overrides)
     return SimpleNamespace(**values)
diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
index d695766f6..198c31908 100644
--- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
+++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py
@@ -20,6 +20,12 @@ class _StubTransformer(nn.Module):
         return torch.float32
 
 
+class _StubTextEncoder(nn.Module):
+    @property
+    def dtype(self) -> torch.dtype:
+        return torch.float32
+
+
 class _StubScheduler:
     def __init__(self, timesteps: list[int]) -> None:
         self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
@@ -41,12 +47,30 @@ def _noop_progress_bar(*args, **kwargs):
     yield _Bar()
 
 
+def _stub_encode_prompt(
+    prompt,
+    negative_prompt=None,
+    do_classifier_free_guidance=True,
+    num_videos_per_prompt=1,
+    max_sequence_length=512,
+    device=None,
+    dtype=None,
+):
+    del negative_prompt, do_classifier_free_guidance, device, dtype
+    batch_size = 1 if isinstance(prompt, str) else len(prompt)
+    n = batch_size * num_videos_per_prompt
+    hidden_size = 8
+    prompt_embeds = torch.zeros(n, max_sequence_length, hidden_size)
+    return prompt_embeds, None
+
+
 def _make_pipeline() -> Wan22Pipeline:
     pipeline = object.__new__(Wan22Pipeline)
     nn.Module.__init__(pipeline)
     pipeline.device = torch.device("cpu")
     pipeline.transformer = _StubTransformer()
     pipeline.transformer_2 = None
+    pipeline.text_encoder = _StubTextEncoder()
     pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4)
     pipeline.scheduler = _StubScheduler([9, 5])
     pipeline.od_config = SimpleNamespace(flow_shift=5.0)
@@ -61,15 +85,14 @@ def _make_pipeline() -> Wan22Pipeline:
     pipeline._num_timesteps = None
     pipeline._current_timestep = None
     pipeline.check_inputs = lambda **kwargs: None
+    pipeline.encode_prompt = _stub_encode_prompt  # type: ignore[method-assign]
     pipeline.prepare_latents = lambda **kwargs: torch.zeros((1, 4, 1, 8, 8), dtype=torch.float32)
     pipeline.progress_bar = _noop_progress_bar
     return pipeline
 
 
-def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
+def test_forward_delegates_denoising_to_diffuse() -> None:
     pipeline = _make_pipeline()
-
-    prompt_embeds = torch.randn(1, 8)
     captured: dict[str, object] = {}
 
     def _fake_diffuse(**kwargs):
@@ -86,8 +109,8 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
             width=None,
             num_frames=1,
             num_inference_steps=2,
-            guidance_scale_provided=False,
-            guidance_scale=None,
+            guidance_scale_provided=True,
+            guidance_scale=1.0,
             guidance_scale_2=None,
             boundary_ratio=None,
             generator=None,
@@ -95,14 +118,16 @@ def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None:
             num_outputs_per_prompt=1,
             max_sequence_length=32,
             latents=None,
+            output_type="latent",
             extra_args={},
         ),
     )
     batch = DiffusionRequestBatch(requests=[mock_req])
 
-    output = pipeline.forward(batch, prompt_embeds=prompt_embeds, output_type="latent", guidance_scale=1.0)
+    output = pipeline.forward(batch)
 
     assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
+    assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
     assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
     assert captured["guidance_low"] == 1.0
     assert captured["guidance_high"] == 1.0
diff --git a/tests/diffusion/offloader/test_module_collector.py b/tests/diffusion/offloader/test_module_collector.py
index bc7fd97d1..e97c1953a 100644
--- a/tests/diffusion/offloader/test_module_collector.py
+++ b/tests/diffusion/offloader/test_module_collector.py
@@ -149,6 +149,40 @@ class MultiVaePipeline(nn.Module, SupportsComponentDiscovery):
         self.audio_vae = nn.Linear(10, 10)
 
 
+class SiblingDiTsPipeline(nn.Module, SupportsComponentDiscovery):
+    """Two independent DiTs (e.g. a MoE transformer pair)."""
+
+    _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Linear(10, 10)
+        self.transformer_2 = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
+class NestedDiTPipeline(nn.Module, SupportsComponentDiscovery):
+    """Cosmos3-like: a DiT plus one of its own submodules listed as a DiT.
+
+    The nested entry exists so offloading can treat the inner stack as an
+    independent ring; the outer DiT's _hsdp_shard_conditions still own all
+    nested blocks for sharding. Order is nested-first, matching Cosmos3.
+    """
+
+    _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = nn.Module()
+        self.transformer.language_model = nn.Linear(10, 10)
+        self.transformer.gen_layers = nn.Linear(10, 10)
+        self.vae = nn.Linear(10, 10)
+
+
 # ---------------------------------------------------------------------------
 # Tests
 # ---------------------------------------------------------------------------
@@ -238,3 +272,35 @@ class TestProtocolDiscovery:
         assert len(result.vaes) == 2
         assert result.vaes[0] is pipeline.vae
         assert result.vaes[1] is pipeline.audio_vae
+
+
+class TestOutermostDits:
+    """PipelineModules.outermost_dits() drops DiTs nested inside another DiT.
+
+    HSDP shards each returned DiT via its own _hsdp_shard_conditions, so a DiT
+    nested inside another discovered DiT must be excluded: it is already covered
+    by its ancestor's conditions and may not declare conditions of its own.
+    """
+
+    def test_keeps_independent_sibling_dits(self):
+        result = ModuleDiscovery.discover(SiblingDiTsPipeline())
+
+        names, modules = result.outermost_dits()
+
+        assert names == ["transformer", "transformer_2"]
+        assert [id(m) for m in modules] == [id(m) for m in result.dits]
+
+    def test_drops_nested_dit_keeping_only_ancestor(self):
+        pipeline = NestedDiTPipeline()
+        result = ModuleDiscovery.discover(pipeline)
+
+        # Discovery sees both, nested-first (matches Cosmos3's _dit_modules order).
+        assert result.dit_names == ["transformer.language_model", "transformer"]
+
+        names, modules = result.outermost_dits()
+
+        # Only the ancestor survives; sharding the nested condition-less stack
+        # again would double-wrap its blocks and raise in apply_hsdp_to_model.
+        assert names == ["transformer"]
+        assert len(modules) == 1
+        assert modules[0] is pipeline.transformer
diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py
index 2539cc068..dc74248e5 100644
--- a/tests/diffusion/offloader/test_sequential_backend.py
+++ b/tests/diffusion/offloader/test_sequential_backend.py
@@ -7,7 +7,8 @@ import pytest
 import torch
 from torch import nn
 
-from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
@@ -41,6 +42,40 @@ def _track_pin_memory_calls():
     return tracker, mock
 
 
+def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None:
+    class CustomPipeline(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.enable_args = None
+            self.disable_called = False
+
+        def enable_omni_model_cpu_offload(self, **kwargs) -> None:
+            self.enable_args = kwargs
+
+        def disable_omni_model_cpu_offload(self) -> None:
+            self.disable_called = True
+
+    pipeline = CustomPipeline()
+    backend = ModelLevelOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False),
+        torch.device("cpu"),
+    )
+
+    backend.enable(pipeline)
+
+    assert backend.enabled is True
+    assert pipeline.enable_args == {
+        "device": torch.device("cpu"),
+        "pin_memory": False,
+        "use_hsdp": False,
+    }
+
+    backend.disable()
+
+    assert backend.enabled is False
+    assert pipeline.disable_called is True
+
+
 class TestMoveParamsPinMemory:
     def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """DTensor should skip pin_memory to avoid RuntimeError."""
diff --git a/tests/diffusion/quantization/test_bitsandbytes_config.py b/tests/diffusion/quantization/test_bitsandbytes_config.py
new file mode 100644
index 000000000..8da2cd521
--- /dev/null
+++ b/tests/diffusion/quantization/test_bitsandbytes_config.py
@@ -0,0 +1,258 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for BitsAndBytes quantization config."""
+
+import importlib.util
+import sys
+import types
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+from torch.nn import Module, Parameter
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
+
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="GPU platform not available.")
+
+bitsandbytes_available = pytest.mark.skipif(
+    importlib.util.find_spec("bitsandbytes") is None,
+    reason="bitsandbytes package not installed",
+)
+
+
+def _ensure_bitsandbytes_importable(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Inject a stub bitsandbytes package when the optional dep is absent.
+
+    ``mocker.patch("bitsandbytes...")`` imports the target module; without a
+    stub, core-model CI without bitsandbytes raises ModuleNotFoundError.
+    """
+    if importlib.util.find_spec("bitsandbytes") is not None:
+        return
+    bnb = types.ModuleType("bitsandbytes")
+    bnb_functional = types.ModuleType("bitsandbytes.functional")
+    bnb.functional = bnb_functional
+    monkeypatch.setitem(sys.modules, "bitsandbytes", bnb)
+    monkeypatch.setitem(sys.modules, "bitsandbytes.functional", bnb_functional)
+
+
+def test_bitsandbytes_config_creation():
+    """Test that BitsAndBytes config can be created."""
+    config = build_quant_config("bitsandbytes")
+    assert config is not None
+    assert config.get_name() == "bitsandbytes"
+
+
+def test_bitsandbytes_config_with_custom_params():
+    """Test BitsAndBytes config with custom parameters."""
+    config = build_quant_config(
+        "bitsandbytes",
+        quant_type="fp4",
+        compress_statistics=False,
+        ignored_layers=["to_out"],
+    )
+    assert config is not None
+    assert config.quant_type == "fp4"
+    assert config.compress_statistics is False
+    assert "to_out" in config.ignored_layers
+
+
+def test_supported_methods():
+    """Test that supported methods list includes bitsandbytes."""
+    assert "bitsandbytes" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_quantization_integration():
+    """Test end-to-end quantization flow through OmniDiffusionConfig."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    config = OmniDiffusionConfig(model="test", quantization_config="bitsandbytes")
+    assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "bitsandbytes"
+
+    config2 = OmniDiffusionConfig(
+        model="test",
+        quantization_config={
+            "method": "bitsandbytes",
+            "quant_type": "nf4",
+            "compress_statistics": True,
+        },
+    )
+    assert config2.quantization_config is not None
+    assert config2.quantization_config.get_name() == "bitsandbytes"
+    assert config2.quantization_config.quant_type == "nf4"
+    assert config2.quantization_config.compress_statistics is True
+
+
+def test_quantization_dict_not_mutated():
+    """Test that passing a dict to quantization_config doesn't mutate it."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    original_dict = {"method": "bitsandbytes", "quant_type": "nf4"}
+    dict_copy = original_dict.copy()
+
+    OmniDiffusionConfig(model="test", quantization_config=original_dict)
+
+    assert original_dict == dict_copy
+
+
+def test_get_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    """Test get_quant_method routing for CUDA."""
+    from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+    config = build_quant_config("bitsandbytes")
+
+    def _fake_init(self, quant_config):
+        pass
+
+    layer = mocker.Mock(spec=LinearBase)
+    mocker.patch.object(BnBOnlineLinearMethod, "__init__", _fake_init)
+
+    prefix = "test_layer"
+
+    monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
+    method = config.get_quant_method(layer, prefix)
+    assert isinstance(method, BnBOnlineLinearMethod)
+
+    config.ignored_layers = [prefix]
+    method = config.get_quant_method(layer, prefix)
+    assert isinstance(method, UnquantizedLinearMethod)
+
+
+class TestBnBOnlineLinearMethod:
+    @pytest.fixture
+    def mock_quant_config(self, mocker):
+        config = mocker.Mock()
+        config.quant_type = "nf4"
+        config.compress_statistics = True
+        return config
+
+    @pytest.fixture
+    def mock_deps(self, mocker, monkeypatch: pytest.MonkeyPatch):
+        _ensure_bitsandbytes_importable(monkeypatch)
+        mock_qweight = torch.ones((64, 32), dtype=torch.uint8)
+        mock_quant_state = mocker.Mock()
+        mock_quant = mocker.patch(
+            "bitsandbytes.functional.quantize_4bit",
+            return_value=(mock_qweight, mock_quant_state),
+            create=True,
+        )
+        # x is (2, 16, 32) -> flattened (32, 32); matmul must return (32, 64)
+        # so apply can reshape back to (2, 16, 64).
+        mock_matmul = mocker.patch(
+            "bitsandbytes.matmul_4bit",
+            return_value=torch.randn(32, 64),
+            create=True,
+        )
+        return {
+            "quant": mock_quant,
+            "matmul": mock_matmul,
+            "mock_qweight": mock_qweight,
+            "mock_quant_state": mock_quant_state,
+        }
+
+    def test_process_weights_after_loading(self, mock_deps, mock_quant_config, mocker):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(mock_quant_config)
+        layer = Module()
+        layer.weight = Parameter(torch.randn(64, 32))
+        mocker.patch.object(torch.Tensor, "cuda", return_value=layer.weight.data)
+
+        method.process_weights_after_loading(layer)
+
+        mock_deps["quant"].assert_called_once()
+        assert torch.equal(layer.weight.data, mock_deps["mock_qweight"])
+        assert layer.quant_state is mock_deps["mock_quant_state"]
+        assert layer.bnb_shape == (64, 32)
+
+    def test_apply(self, mock_deps, mock_quant_config):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(mock_quant_config)
+        layer = Module()
+        layer.weight = Parameter(torch.ones((64, 32), dtype=torch.uint8), requires_grad=False)
+        layer.quant_state = mock_deps["mock_quant_state"]
+
+        x = torch.randn(2, 16, 32, dtype=torch.float16)
+        bias = torch.randn(64)
+        output = method.apply(layer, x, bias)
+
+        mock_deps["matmul"].assert_called_once()
+        assert output.shape == (2, 16, 64)
+        assert output.dtype == torch.float16
+
+
+@pytest.fixture
+def quant_config():
+    from vllm_omni.quantization.bitsandbytes_config import DiffusionBitsAndBytesConfig
+
+    return DiffusionBitsAndBytesConfig(
+        quant_type="nf4",
+        compress_statistics=True,
+    )
+
+
+@cuda_available
+@bitsandbytes_available
+class TestCudaBnBSmoke:
+    """Smoke tests using real bitsandbytes CUDA kernels."""
+
+    @pytest.fixture
+    def real_layer(self):
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(
+            torch.randn(128, 64, dtype=torch.float16, device="cuda"),
+            requires_grad=False,
+        )
+        layer.logical_widths = [128]
+        layer.input_size_per_partition = 64
+        layer.output_size_per_partition = 128
+        layer.orig_dtype = torch.float16
+        return layer
+
+    def test_real_cuda_quantize_4bit_shape_contract(self, quant_config):
+        import bitsandbytes.functional as bnb_F
+
+        weight = torch.randn(128, 64, dtype=torch.float16, device="cuda")
+        qweight, quant_state = bnb_F.quantize_4bit(
+            weight,
+            quant_type=quant_config.quant_type,
+            compress_statistics=quant_config.compress_statistics,
+        )
+
+        assert qweight.dtype == torch.uint8
+        assert quant_state is not None
+
+    def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(quant_config)
+        method.process_weights_after_loading(real_layer)
+
+        assert real_layer.weight.dtype == torch.uint8
+        assert hasattr(real_layer, "quant_state")
+        assert real_layer.bnb_shape == (128, 64)
+
+    def test_real_cuda_bnb_apply_forward(self, quant_config):
+        from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod
+
+        method = BnBOnlineLinearMethod(quant_config)
+
+        layer = torch.nn.Module()
+        layer.weight = torch.nn.Parameter(
+            torch.randn(128, 64, dtype=torch.float16, device="cuda"),
+            requires_grad=False,
+        )
+        method.process_weights_after_loading(layer)
+
+        x = torch.randn(2, 16, 64, dtype=torch.float16, device="cuda")
+        output = method.apply(layer, x)
+
+        assert output.shape == (2, 16, 128)
+        assert output.dtype == torch.float16
diff --git a/tests/diffusion/quantization/test_fp8_config.py b/tests/diffusion/quantization/test_fp8_config.py
index a1f79030e..9bad16792 100644
--- a/tests/diffusion/quantization/test_fp8_config.py
+++ b/tests/diffusion/quantization/test_fp8_config.py
@@ -154,21 +154,6 @@ def test_component_config_with_default():
     assert config.resolve("vae.encoder") is None
 
 
-def test_gguf_config():
-    from vllm_omni.quantization import build_quant_config
-    from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
-
-    config = build_quant_config(
-        {
-            "method": "gguf",
-            "gguf_model": "path/to/model.gguf",
-        }
-    )
-    assert isinstance(config, DiffusionGGUFConfig)
-    assert config.gguf_model == "path/to/model.gguf"
-    assert config.get_name() == "gguf"
-
-
 def test_integration_string():
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
@@ -236,7 +221,7 @@ def test_transformer_config_auto_detects_modelopt_fp8():
 def test_supported_methods_includes_vllm():
     from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
 
-    for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]:
+    for method in ["fp8", "awq", "gptq", "bitsandbytes", "modelopt"]:
         assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing"
 
 
diff --git a/tests/diffusion/quantization/test_gguf_config.py b/tests/diffusion/quantization/test_gguf_config.py
deleted file mode 100644
index 201bcdaba..000000000
--- a/tests/diffusion/quantization/test_gguf_config.py
+++ /dev/null
@@ -1,128 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for diffusion GGUF quantization config and linear method."""
-
-from __future__ import annotations
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
-
-from vllm_omni.quantization import build_quant_config
-from vllm_omni.quantization.gguf_config import (
-    UNQUANTIZED_TYPES,
-    DiffusionGGUFConfig,
-    DiffusionGGUFLinearMethod,
-    dequant_gemm_gguf,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def test_gguf_config_creation_and_delegation():
-    config = DiffusionGGUFConfig(
-        gguf_model="weights.gguf",
-        unquantized_modules=["proj_out"],
-    )
-
-    assert config.gguf_model == "weights.gguf"
-    assert config.unquantized_modules == ["proj_out"]
-    assert config.get_name() == "gguf"
-
-
-def test_build_quant_config_builds_gguf_config():
-    config = build_quant_config(
-        "gguf",
-        gguf_model="repo/model-Q4_0.gguf",
-        unquantized_modules=["norm_out"],
-    )
-
-    assert isinstance(config, DiffusionGGUFConfig)
-    assert config.gguf_model == "repo/model-Q4_0.gguf"
-    assert config.unquantized_modules == ["norm_out"]
-
-
-def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers():
-    linear = object.__new__(LinearBase)
-    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in")
-
-    assert isinstance(method, DiffusionGGUFLinearMethod)
-
-
-def test_gguf_vllm_config_respects_unquantized_modules():
-    linear = object.__new__(LinearBase)
-    method = DiffusionGGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out")
-
-    assert isinstance(method, UnquantizedLinearMethod)
-
-
-def test_gguf_vllm_config_returns_none_for_non_linear_layers():
-    method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm")
-
-    assert method is None
-
-
-def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    x = torch.tensor([[1.0, 2.0]], dtype=torch.float32)
-    qweight = torch.tensor([[3.0, 4.0], [5.0, 6.0]], dtype=torch.float32)
-
-    out = dequant_gemm_gguf(x, qweight, qweight_type)
-
-    assert torch.allclose(out, x @ qweight.T)
-
-
-def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
-    layer = SimpleNamespace(
-        qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32),
-        qweight_type=SimpleNamespace(weight_type=qweight_type),
-    )
-    x = torch.tensor([[[1.0, 0.5]]], dtype=torch.float32)
-    bias = torch.tensor([0.25, -0.5], dtype=torch.float32)
-
-    out = method.apply(layer, x, bias)
-
-    expected = x @ layer.qweight.T + bias
-    assert torch.allclose(out, expected)
-
-
-def test_diffusion_gguf_linear_method_concatenates_sharded_outputs():
-    qweight_type = next(iter(UNQUANTIZED_TYPES))
-    method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[]))
-    qweight = torch.nn.Parameter(
-        torch.tensor(
-            [
-                [1.0, 0.0],
-                [0.0, 1.0],
-                [2.0, 0.0],
-                [0.0, 2.0],
-            ],
-            dtype=torch.float32,
-        ),
-        requires_grad=False,
-    )
-    qweight.shard_id = ["left", "right"]
-    qweight.shard_offset_map = {
-        "left": (0, 2, 2),
-        "right": (2, 4, 2),
-    }
-    layer = SimpleNamespace(
-        qweight=qweight,
-        qweight_type=SimpleNamespace(
-            shard_weight_type={
-                "left": qweight_type,
-                "right": qweight_type,
-            }
-        ),
-    )
-    x = torch.tensor([[1.5, 2.0]], dtype=torch.float32)
-
-    out = method.apply(layer, x)
-
-    expected_q = x @ qweight[:2].T
-    expected_k = x @ qweight[2:4].T
-    assert torch.allclose(out, torch.cat([expected_q, expected_k], dim=-1))
diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py
index af3a9c9c4..240ee2509 100644
--- a/tests/diffusion/quantization/test_mxfp4_config.py
+++ b/tests/diffusion/quantization/test_mxfp4_config.py
@@ -537,3 +537,118 @@ def test_single_scale_row_parallel_tp2_loader_simulation():
 
     for rank in range(_TP2):
         assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape
+
+
+# ---------------------------------------------------------------------------
+# ROCm MXFP4 (gfx950) — get_quant_method dispatch
+#
+# These run on CPU: the platform check, gcnArchName probe and aiter op
+# registration are all mocked (stdlib unittest.mock + monkeypatch, so no
+# pytest-mock dependency).  They cover only the dispatch + weight-allocation
+# logic added by the ROCm PR — the AITER GEMM / quant kernels require real
+# gfx950 hardware and are intentionally NOT exercised here.
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def _rocm_platform(monkeypatch: pytest.MonkeyPatch):
+    """Make current_omni_platform report ROCm, and stub the aiter
+    custom-op registration so ROCmMxfp4*Method can be constructed without aiter."""
+    from vllm_omni.platforms import current_omni_platform
+    from vllm_omni.quantization import mxfp4_config
+
+    monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
+    monkeypatch.setattr(current_omni_platform, "is_rocm", lambda: True)
+    monkeypatch.setattr(mxfp4_config, "_register_rocm_mxfp4_op", lambda: None)
+
+
+def _patch_gcn_arch(monkeypatch: pytest.MonkeyPatch, arch: str) -> None:
+    """Patch torch.cuda.get_device_properties(...).gcnArchName to return `arch`."""
+    from types import SimpleNamespace
+
+    monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0)
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda *a, **k: SimpleNamespace(gcnArchName=arch),
+    )
+
+
+def _fake_linear_layer():
+    """A stand-in that passes isinstance(layer, LinearBase) without a real layer."""
+    from unittest.mock import MagicMock
+
+    from vllm.model_executor.layers.linear import LinearBase
+
+    return MagicMock(spec=LinearBase)
+
+
+def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch):
+    """ROCm + gfx950 + online checkpoint must return ROCmMxfp4OnlineLinearMethod."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False)
+
+    method = cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
+    assert isinstance(method, ROCmMxfp4OnlineLinearMethod)
+
+
+def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch):
+    """A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe."""
+    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"])
+
+    assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod)
+
+
+def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch):
+    """MXFP4 on ROCm requires gfx950; any other arch must raise."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
+
+    _patch_gcn_arch(monkeypatch, "gfx942:sramecc+:xnack-")
+    cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False)
+
+    with pytest.raises(NotImplementedError, match="gfx950"):
+        cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q")
+
+
+# ---------------------------------------------------------------------------
+# ROCm MXFP4 — create_weights (online lazy meta-device placeholder)
+#
+# The base ROCmMxfp4LinearMethod is abstract (no create_weights); the online
+# subclass gets create_weights from _LazyWeightMixin, which registers a BF16
+# weight on the meta device to be materialised at load time.
+# ---------------------------------------------------------------------------
+
+
+def test_rocm_create_weights_column_parallel_tp2(_rocm_platform):
+    """Column-parallel TP=2: meta BF16 weight has the output halved, input full."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
+
+    assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K)
+    assert layer.weight.dtype == torch.bfloat16
+    assert layer.weight.device.type == "meta"
+    assert layer.weight.input_dim == 1
+    assert layer.weight.output_dim == 0
+    assert layer.logical_widths == [_TP2_N // _TP2]
+    assert layer.weight_block_size is None
+
+
+def test_rocm_create_weights_row_parallel_tp2(_rocm_platform):
+    """Row-parallel TP=2: meta BF16 weight has the input halved, output full."""
+    from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod
+
+    method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config())
+    layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
+
+    assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
+    assert layer.input_size_per_partition == _TP2_K // _TP2
+    assert layer.output_size_per_partition == _TP2_N
diff --git a/tests/diffusion/quantization/test_trajectory_similarity_tool.py b/tests/diffusion/quantization/test_trajectory_similarity_tool.py
index 45d51d356..252e6f1c1 100644
--- a/tests/diffusion/quantization/test_trajectory_similarity_tool.py
+++ b/tests/diffusion/quantization/test_trajectory_similarity_tool.py
@@ -99,7 +99,6 @@ def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantiza
         candidate_quantization_config_json=None,
         candidate_ignored_layers=None,
         ignored_layers=None,
-        candidate_gguf_model=None,
         candidate_load_format="default",
     )
 
diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py
index 55584609d..5d4ef5d72 100644
--- a/tests/diffusion/test_diffusion_model_runner.py
+++ b/tests/diffusion/test_diffusion_model_runner.py
@@ -509,6 +509,25 @@ def test_split_diffusion_output_by_request_slices_tuple_outputs():
     assert torch.equal(outputs[1].output[1], torch.tensor([20]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_model_runs_forward_after_kv_receive(monkeypatch):
+    """execute_model runs the pipeline forward after the KV receive step."""
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.kv_transfer_manager.receive_multi_kv_cache_distributed = lambda *a, **k: True
+    req = _make_request(skip_cache_refresh=True)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "ok"
+    assert runner.pipeline.forward_calls == 1
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py
index 185d51b54..22b305eab 100644
--- a/tests/diffusion/test_diffusion_output_formatter.py
+++ b/tests/diffusion/test_diffusion_output_formatter.py
@@ -51,21 +51,23 @@ def _timings() -> DiffusionStepTimings:
     )
 
 
-def test_formatter_preserves_single_video_audio_actions_and_custom_output(
+def test_formatter_preserves_single_video_audio_actions_and_metadata(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
-    custom_output = {"from_output": "kept"}
     postprocess_output = normalize_diffusion_postprocess_output(
         {
-            "video": ["frame-0"],
-            "audio": "audio-0",
-            "actions": "action-0",
-            "custom_output": {"from_postprocess": "merged"},
-            "audio_sample_rate": 48000,
-            "fps": 24.0,
-        },
-        custom_output,
+            "payload": {
+                "video": ["frame-0"],
+                "audio": "audio-0",
+                "actions": "action-0",
+            },
+            "metadata": {
+                "audio": {"sample_rate": 48000},
+                "video": {"fps": 24.0},
+                "actions": {"action_mode": "policy"},
+            },
+        }
     )
 
     results = format_diffusion_outputs(
@@ -73,7 +75,6 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
         od_config=_config(),
         diffusion_output=DiffusionOutput(
             output=None,
-            custom_output=custom_output,
             stage_durations={"execute": 1.25},
             peak_memory_mb=321.0,
         ),
@@ -87,12 +88,12 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
     assert result.images == ["frame-0"]
     assert result.prompt == "prompt-0"
     assert result.final_output_type == "image"
-    assert result.custom_output == {
-        "from_output": "kept",
-        "from_postprocess": "merged",
-    }
-    assert custom_output == {"from_output": "kept"}
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 48000},
+            "video": {"fps": 24.0},
+            "actions": {"action_mode": "policy"},
+        },
         "audio": "audio-0",
         "audio_sample_rate": 48000,
         "fps": 24.0,
@@ -110,18 +111,123 @@ def test_formatter_preserves_single_video_audio_actions_and_custom_output(
     }
 
 
-def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch) -> None:
+def test_formatter_normalizes_payload_metadata_envelope(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
     postprocess_output = normalize_diffusion_postprocess_output(
-        "caption",
-        {"text_output": "caption"},
+        {
+            "payload": {
+                "video": ["frame-0"],
+                "actions": "action-0",
+            },
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": 2,
+                    "action_mode": "policy",
+                    "domain_id": 7,
+                },
+                "common": {
+                    "action_only_output": True,
+                },
+                "internal": {
+                    "robolab_action_postprocess": object(),
+                },
+            },
+        }
+    )
+
+    assert postprocess_output.outputs == {
+        "video": ["frame-0"],
+        "actions": "action-0",
+    }
+    assert postprocess_output.primary_key == "video"
+    assert postprocess_output.metadata == {
+        "actions": {
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        },
+        "common": {
+            "action_only_output": True,
+        },
+    }
+    assert "internal" not in postprocess_output.metadata
+
+    [result] = format_diffusion_outputs(
+        request=_request("prompt-0"),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.multimodal_output == {
+        "metadata": postprocess_output.metadata,
+        "actions": "action-0",
+    }
+
+
+def test_formatter_maps_trajectory_payload_to_request_output(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    latents = torch.zeros(2, 1, 4)
+    timesteps = torch.tensor([1.0, 0.5])
+    log_probs = torch.tensor([-0.1, -0.2])
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "payload": {
+                "image": "image-0",
+                "trajectory": {
+                    "latents": latents,
+                    "timesteps": timesteps,
+                    "log_probs": log_probs,
+                },
+            },
+            "metadata": {"trajectory": {"type": "denoising"}},
+        }
+    )
+
+    [result] = format_diffusion_outputs(
+        request=_request("prompt-0"),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"raw": "output"},
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert result.images == ["image-0"]
+    assert result.latents is latents
+    assert result.trajectory_latents is latents
+    assert result.trajectory_timesteps is timesteps
+    assert result.trajectory_log_probs is log_probs
+    assert result.multimodal_output == {
+        "metadata": {"trajectory": {"type": "denoising"}},
+        "trajectory": {
+            "latents": latents,
+            "timesteps": timesteps,
+            "log_probs": log_probs,
+        },
+    }
+
+
+def test_formatter_preserves_text_envelope_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "payload": {"text": "caption"},
+            "metadata": {"text": {"text_output": "caption"}},
+        }
     )
 
     [result] = format_diffusion_outputs(
         request=_request("describe this"),
         od_config=_config(),
-        diffusion_output=DiffusionOutput(output="caption"),
-        output_data="caption",
+        diffusion_output=DiffusionOutput(output=None),
+        output_data={"payload": {"text": "caption"}},
         postprocess_output=postprocess_output,
         timings=_timings(),
     )
@@ -129,7 +235,10 @@ def test_formatter_preserves_text_custom_output(monkeypatch: pytest.MonkeyPatch)
     assert result.images == []
     assert result.prompt == "describe this"
     assert result.final_output_type == "text"
-    assert result.custom_output == {"text_output": "caption"}
+    assert result.multimodal_output == {
+        "metadata": {"text": {"text_output": "caption"}},
+        "text": "caption",
+    }
 
 
 def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
@@ -144,7 +253,7 @@ def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
         "_try_load_model_cls",
         lambda _: AudioModel,
     )
-    postprocess_output = normalize_diffusion_postprocess_output(["waveform"], {})
+    postprocess_output = normalize_diffusion_postprocess_output(["waveform"])
 
     [result] = format_diffusion_outputs(
         request=_request("speak"),
@@ -175,8 +284,7 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
             "actions": "action-0",
             "audio_sample_rate": 16000,
             "fps": 30.0,
-        },
-        {},
+        }
     )
 
     [result] = format_diffusion_outputs(
@@ -192,6 +300,10 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
     assert result.prompt == "watch and listen"
     assert result.final_output_type == "image"
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 16000},
+            "video": {"fps": 30.0},
+        },
         "audio": "audio-0",
         "audio_sample_rate": 16000,
         "fps": 30.0,
@@ -207,8 +319,7 @@ def test_formatter_preserves_audio_only_postprocess_dict(
         {
             "audio": "waveform",
             "audio_sample_rate": 24000,
-        },
-        {},
+        }
     )
 
     [result] = format_diffusion_outputs(
@@ -224,6 +335,9 @@ def test_formatter_preserves_audio_only_postprocess_dict(
     assert result.prompt == "speak"
     assert result.final_output_type == "audio"
     assert result.multimodal_output == {
+        "metadata": {
+            "audio": {"sample_rate": 24000},
+        },
         "audio": "waveform",
         "audio_sample_rate": 24000,
     }
@@ -238,7 +352,7 @@ def test_formatter_preserves_single_prompt_multiple_audio_outputs(
         "_try_load_model_cls",
         lambda _: None,
     )
-    postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"], {})
+    postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"])
 
     [result] = format_diffusion_outputs(
         request=_request("speak", num_outputs_per_prompt=2),
@@ -264,10 +378,8 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads(
             "video": ["frame-0", "frame-1"],
             "audio": ["audio-0", "audio-1"],
             "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
-            "custom_output": {"shared": True},
             "fps": 12.5,
-        },
-        {},
+        }
     )
 
     results = format_diffusion_outputs(
@@ -282,7 +394,7 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads(
     assert len(results) == 1
     assert results[0].images == ["frame-0", "frame-1"]
     assert results[0].prompt == "prompt-0"
-    assert results[0].custom_output == {"shared": True}
+    assert results[0].multimodal_output["metadata"] == {"video": {"fps": 12.5}}
     assert results[0].multimodal_output["audio"] == ["audio-0", "audio-1"]
     assert results[0].multimodal_output["fps"] == 12.5
     torch.testing.assert_close(
diff --git a/tests/diffusion/test_diffusion_output_metadata.py b/tests/diffusion/test_diffusion_output_metadata.py
new file mode 100644
index 000000000..ee3559f47
--- /dev/null
+++ b/tests/diffusion/test_diffusion_output_metadata.py
@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.outputs.output_metadata import (
+    strip_internal_metadata,
+    validate_diffusion_metadata,
+    validate_public_diffusion_metadata,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_validate_accepts_canonical_metadata_fields() -> None:
+    validate_diffusion_metadata(
+        {
+            "video": {
+                "fps": 24.0,
+                "video_fps_multiplier": 2,
+                "shape": [1, 3, 16, 64, 64],
+            },
+            "audio": {
+                "sample_rate": 48000,
+                "audiox_task": "tts",
+            },
+            "actions": {
+                "raw_action_dim": 7,
+                "action_mode": "policy",
+                "domain_id": 3,
+            },
+            "transfer": {
+                "controls": {"edge": torch.zeros(1, 3, 4, 8, 8)},
+                "hints": ["edge"],
+            },
+            "text": {
+                "text_output": "caption",
+                "think_text": "reasoning",
+            },
+            "common": {
+                "action_only_output": True,
+            },
+        }
+    )
+
+
+@pytest.mark.parametrize(
+    ("metadata", "error_type"),
+    [
+        ({"video": {"fps": 0}}, ValueError),
+        ({"video": {"fps": "24"}}, TypeError),
+        ({"video": {"video_fps_multiplier": 0}}, ValueError),
+        ({"video": {"shape": [1, 0, 64]}}, ValueError),
+        ({"audio": {"sample_rate": 0}}, ValueError),
+        ({"audio": {"audiox_task": ""}}, ValueError),
+        ({"actions": {"raw_action_dim": 0}}, ValueError),
+        ({"actions": {"action_mode": ""}}, ValueError),
+        ({"actions": {"domain_id": "7"}}, TypeError),
+        ({"transfer": {"controls": {"edge": "not-a-tensor"}}}, TypeError),
+        ({"transfer": {"hints": ["edge", 1]}}, TypeError),
+        ({"text": {"text_output": 1}}, TypeError),
+        ({"text": {"think_text": 1}}, TypeError),
+        ({"common": {"action_only_output": "true"}}, TypeError),
+    ],
+)
+def test_validate_rejects_invalid_canonical_metadata(metadata: dict[str, object], error_type: type[Exception]) -> None:
+    with pytest.raises(error_type):
+        validate_diffusion_metadata(metadata)
+
+
+def test_validate_allows_unknown_model_metadata() -> None:
+    validate_diffusion_metadata(
+        {
+            "model_specific": {"opaque": object()},
+            "video": {"model_key": object()},
+        }
+    )
+
+
+def test_public_metadata_rejects_internal_group() -> None:
+    with pytest.raises(ValueError, match="internal"):
+        validate_public_diffusion_metadata({"internal": {"secret": object()}})
+
+
+def test_strip_internal_metadata_removes_private_group() -> None:
+    metadata = {
+        "actions": {"action_mode": "policy"},
+        "internal": {"robolab_action_postprocess": object()},
+    }
+
+    public_metadata = strip_internal_metadata(metadata)
+
+    assert public_metadata == {"actions": {"action_mode": "policy"}}
+    assert "internal" not in public_metadata
diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py
index 6804d550d..a437a572a 100644
--- a/tests/diffusion/test_diffusion_plugin_hooks.py
+++ b/tests/diffusion/test_diffusion_plugin_hooks.py
@@ -16,7 +16,6 @@ from unittest.mock import Mock, patch
 import pytest
 
 from vllm_omni.diffusion.registry import (
-    _DIFFUSION_ACTION_POST_PROCESS_FUNCS,
     _DIFFUSION_IR_OP_PRIORITY_FUNCS,
     _DIFFUSION_MODELS,
     _DIFFUSION_POST_PROCESS_FUNCS,
@@ -62,7 +61,6 @@ class TestRegisterDiffusionModel:
         original_models = _DIFFUSION_MODELS.copy()
         original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
         original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
-        original_action_post = _DIFFUSION_ACTION_POST_PROCESS_FUNCS.copy()
         original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
         yield
         _DIFFUSION_MODELS.clear()
@@ -71,8 +69,6 @@ class TestRegisterDiffusionModel:
         _DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre)
         _DIFFUSION_POST_PROCESS_FUNCS.clear()
         _DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS.clear()
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS.update(original_action_post)
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
 
@@ -84,7 +80,6 @@ class TestRegisterDiffusionModel:
             class_name="TestPipeline",
             pre_process_func_name="test_pre_process",
             post_process_func_name="test_post_process",
-            action_post_process_func_name="test_action_post_process",
             ir_op_priority_func_name="test_ir_op_priority",
         )
         assert "TestPipeline" in _DIFFUSION_MODELS
@@ -95,9 +90,21 @@ class TestRegisterDiffusionModel:
         )
         assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
         assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
-        assert _DIFFUSION_ACTION_POST_PROCESS_FUNCS["TestPipeline"] == "test_action_post_process"
         assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
 
+    def test_register_model_accepts_deprecated_action_postprocess_keyword(self):
+        """Deprecated action hook keyword is accepted but not registered."""
+        register_diffusion_model(
+            model_arch="LegacyActionPipeline",
+            module_name="test_plugin.diffusion.pipeline",
+            class_name="LegacyActionPipeline",
+            post_process_func_name="test_post_process",
+            action_post_process_func_name="test_action_post_process",
+        )
+
+        assert "LegacyActionPipeline" in _DIFFUSION_MODELS
+        assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process"
+
 
 class TestWorkerUsesHook:
     """Test that DiffusionWorker resolves model runner via platform hook."""
diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py
index bd2289283..7e18996ed 100644
--- a/tests/diffusion/test_diffusion_step_pipeline.py
+++ b/tests/diffusion/test_diffusion_step_pipeline.py
@@ -30,6 +30,9 @@ from vllm_omni.diffusion.ipc import (
     pack_diffusion_output_shm,
     unpack_diffusion_output_shm,
 )
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import StepScheduler
 from vllm_omni.diffusion.sched.interface import (
@@ -97,6 +100,51 @@ class _StepPipeline:
         return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
 
 
+class _ProfilingStepPipeline(_StepPipeline):
+    enable_diffusion_pipeline_profiler = True
+
+    def __init__(self):
+        super().__init__()
+        self._stage_durations: dict[str, float] = {}
+
+    @property
+    def stage_durations(self) -> dict[str, float]:
+        return dict(self._stage_durations)
+
+    def clear_profiler_records(self) -> None:
+        self._stage_durations.clear()
+
+    def prepare_encode(self, state, **kwargs):
+        result = super().prepare_encode(state, **kwargs)
+        self._stage_durations["QwenImagePipeline.text_encoder.forward"] = 1.0
+        return result
+
+    def denoise_step(self, input_batch, **kwargs):
+        result = super().denoise_step(input_batch, **kwargs)
+        self._stage_durations["QwenImagePipeline.diffuse"] = 2.0
+        return result
+
+    def post_decode(self, state, **kwargs):
+        result = super().post_decode(state, **kwargs)
+        self._stage_durations["QwenImagePipeline.vae.decode"] = 3.0
+        return result
+
+
+class _AutoDenoiseProfilerPipeline(DiffusionPipelineProfilerMixin):
+    _PROFILER_TARGETS: list[str] = []
+
+    def __init__(self):
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=True,
+        )
+
+    def forward(self):
+        return None
+
+    def denoise_step(self):
+        return "ok"
+
+
 class _InterruptingStepPipeline(_StepPipeline):
     interrupt = True
 
@@ -130,6 +178,9 @@ class _FakePeakMemoryPlatform:
         index = min(self.reset_calls - 1, len(self._reserved_mb) - 1)
         return int((self._reserved_mb[index] - 100) * 1024**2)
 
+    def is_available(self) -> bool:
+        return True
+
 
 class _IdentityNoiseTransformer(torch.nn.Module):
     def forward(self, x: torch.Tensor, **kwargs):
@@ -437,6 +488,16 @@ def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_compositi
     torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3))
 
 
+@pytest.mark.cpu
+def test_step_profiler_reports_denoise_step_as_diffuse():
+    pipeline = _AutoDenoiseProfilerPipeline()
+
+    assert pipeline.denoise_step() == "ok"
+
+    assert any(key.endswith(".diffuse") for key in pipeline.stage_durations)
+    assert not any(key.endswith(".denoise_step") for key in pipeline.stage_durations)
+
+
 @pytest.mark.cpu
 class TestRunner:
     """DiffusionModelRunner.execute_stepwise"""
@@ -469,6 +530,53 @@ class TestRunner:
         assert runner.pipeline.scheduler_calls == 2
         assert runner.pipeline.decode_calls == 1
 
+    def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch):
+        runner = _make_runner()
+        runner.pipeline = _ProfilingStepPipeline()
+        req = _make_step_request()
+        reset_calls = []
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "is_available",
+            lambda: True,
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "reset_peak_memory_stats",
+            lambda: reset_calls.append(True),
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "max_memory_reserved",
+            lambda: 2 * 1024**2,
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "max_memory_allocated",
+            lambda: 1024**2,
+        )
+
+        DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_scheduler_output(req, step_id=0),
+        )
+        result = DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_cached_scheduler_output(step_id=1),
+        )
+
+        output = result.get_request_output("req-1")
+        assert output.finished is True
+        assert output.result is not None
+        assert output.result.peak_memory_mb == 2
+        assert output.result.stage_durations == {
+            "QwenImagePipeline.text_encoder.forward": 1.0,
+            "QwenImagePipeline.diffuse": 4.0,
+            "QwenImagePipeline.vae.decode": 3.0,
+        }
+        assert reset_calls == [True]
+
     def test_carries_peak_memory_across_stepwise_request_lifecycle(self, monkeypatch):
         runner = _make_runner()
         req = _make_step_request()
diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py
index 50ed772c1..ae5f9eee8 100644
--- a/tests/diffusion/test_diffusion_streaming_output.py
+++ b/tests/diffusion/test_diffusion_streaming_output.py
@@ -85,19 +85,20 @@ class _PipelineBackedEngine:
             outputs = self.pipeline.step_outputs(request)
         except Exception as exc:
             outputs = [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
+                _streaming_diffusion_output(chunk=0, finished=False),
                 DiffusionOutput(error=str(exc), finished=True),
             ]
         for output in outputs:
             if output.error is not None:
                 yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)]
                 continue
-            custom_output = output.custom_output or {}
+            payload, metadata = _streaming_payload_and_metadata(output)
+            stream_metadata = metadata.get("stream", {})
             yield [
                 OmniRequestOutput.from_diffusion(
                     request_id=request.request_id,
-                    images=list(custom_output.get("images") or []),
-                    custom_output=custom_output,
+                    images=list(payload.get("image") or []),
+                    multimodal_output={"metadata": {"stream": stream_metadata}},
                     finished=output.finished,
                 )
             ]
@@ -106,6 +107,30 @@ class _PipelineBackedEngine:
         del request_id
 
 
+def _streaming_diffusion_output(
+    *,
+    chunk: int,
+    finished: bool,
+    images: list | None = None,
+) -> DiffusionOutput:
+    payload = {"image": images or []}
+    metadata = {"stream": {"chunk": chunk}}
+    return DiffusionOutput(
+        output={
+            "payload": payload,
+            "metadata": metadata,
+        },
+        finished=finished,
+    )
+
+
+def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]:
+    envelope = output.output if isinstance(output.output, dict) else {}
+    payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {}
+    metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {}
+    return payload, metadata
+
+
 class TestPipelineStreamingOutputToStageDiffusionClient:
     """Streaming pipeline output over ZMQ into ``StageDiffusionClient``."""
 
@@ -114,8 +139,8 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
         """Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags."""
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1}, finished=True),
+                _streaming_diffusion_output(chunk=0, finished=False),
+                _streaming_diffusion_output(chunk=1, finished=True),
             ]
         )
 
@@ -123,7 +148,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
 
         assert [output.request_id for output in outputs] == ["req-stream", "req-stream"]
         assert [output.finished for output in outputs] == [False, True]
-        assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
+        assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
         assert pipeline.requests[0].request_id == "req-stream"
 
     @pytest.mark.asyncio
@@ -136,7 +161,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient:
         assert len(outputs) == 2
         assert outputs[0].request_id == "req-error"
         assert outputs[0].finished is False
-        assert outputs[0].custom_output["chunk"] == 0
+        assert outputs[0].multimodal_output["metadata"]["stream"]["chunk"] == 0
         assert outputs[1].request_id == "req-error"
         assert outputs[1].finished is True
         assert outputs[1].error is not None
@@ -227,8 +252,8 @@ class TestPipelineStreamingOutputToEntrypoint:
         """Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator."""
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1}, finished=True),
+                _streaming_diffusion_output(chunk=0, finished=False),
+                _streaming_diffusion_output(chunk=1, finished=True),
             ]
         )
         inline_client = self._make_inline_pipeline_client(pipeline)
@@ -247,7 +272,7 @@ class TestPipelineStreamingOutputToEntrypoint:
 
             await _wait_for(lambda: len(pipeline.requests) == 1)
             assert pipeline.requests[0].request_id.startswith("req-omni-")
-            assert [output.custom_output["chunk"] for output in outputs] == [0, 1]
+            assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1]
             assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True]
         finally:
             await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client)
@@ -258,8 +283,8 @@ class TestPipelineStreamingOutputToEntrypoint:
         frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)]
         pipeline = _StepStreamingPipeline(
             [
-                DiffusionOutput(custom_output={"chunk": 0, "images": [frames]}, finished=False),
-                DiffusionOutput(custom_output={"chunk": 1, "images": [frames]}, finished=True),
+                _streaming_diffusion_output(chunk=0, images=[frames], finished=False),
+                _streaming_diffusion_output(chunk=1, images=[frames], finished=True),
             ]
         )
         inline_client = self._make_inline_pipeline_client(pipeline)
diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py
index 8f8783db2..261de0b32 100644
--- a/tests/distributed/omni_connectors/test_tp_rank_aware.py
+++ b/tests/distributed/omni_connectors/test_tp_rank_aware.py
@@ -937,6 +937,76 @@ class TestDistributedReceive:
 # ── TP auto-detect ───────────────────────────────────────────────────
 
 
+class TestBuildCfgRankLocalPayloads:
+    """Guard _build_cfg_rank_local_payloads padding logic for the
+    cfg_size==2 / no branch_roles scenario (hunyuan-image shared-prefix reuse)."""
+
+    def test_cfg_size2_no_branch_roles_copies_main_payload_to_follower(self):
+        """When AR doesn't split KV by branch, cfg follower (rank 1) should
+        receive a copy of the positive KV, not None."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        # Rank 0: main payload (always present)
+        assert payloads[0] is not None
+        assert "past_key_values" in payloads[0]
+        # Rank 1: must also be non-None — same positive KV for shared-prefix reuse
+        assert payloads[1] is not None
+        assert payloads[1]["past_key_values"] is req.past_key_values
+        assert payloads[1]["kv_metadata"] == {"seq_len": 10}
+
+    def test_cfg_size3_no_branch_roles_keeps_none_padding(self):
+        """When cfg_size > 2 and no branch_roles, padding stays None (safety guard)."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=3)
+
+        assert len(payloads) == 3
+        assert payloads[0] is not None
+        assert payloads[1] is None
+        assert payloads[2] is None
+
+    def test_empty_main_payload_keeps_none_padding(self):
+        """When main_payload is empty, padding stays None even with cfg_size==2."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            sampling_params=SimpleNamespace(),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        # No past_key_values → main_payload is empty dict → falsy → padding stays None
+        assert payloads[1] is None
+
+    def test_with_branch_roles_uses_original_behavior(self):
+        """When branch_roles exist, follower payload comes from branch KV, not padding."""
+        req = SimpleNamespace(
+            request_id="req-1",
+            past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]),
+            kv_metadata={"seq_len": 10},
+            sampling_params=SimpleNamespace(
+                cfg_text_past_key_values=SimpleNamespace(key_cache=[torch.tensor([2.0])]),
+                cfg_text_kv_metadata={"source": "cfg_text"},
+            ),
+        )
+        payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2)
+
+        assert len(payloads) == 2
+        assert payloads[0] is not None
+        assert payloads[1] is not None
+        assert payloads[1]["sp.cfg_active_branch"] == "cfg_text"
+
+
 class TestAutoDetectTP:
     def test_auto_detect_when_config_defaults(self):
         """When config from_tp/to_tp == 1 (default), manager should auto-detect."""
diff --git a/tests/docs/test_generate_examples.py b/tests/docs/test_generate_examples.py
new file mode 100644
index 000000000..0b29837a9
--- /dev/null
+++ b/tests/docs/test_generate_examples.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib.util
+from pathlib import Path
+
+import pytest
+
+ROOT_DIR = Path(__file__).parents[2]
+HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_examples.py"
+SPEC = importlib.util.spec_from_file_location("generate_examples", HOOK_PATH)
+assert SPEC and SPEC.loader
+generate_examples = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(generate_examples)
+
+
+def test_model_example_titles_match_shared_display_names():
+    for slug, display_name in generate_examples.MODEL_DISPLAY_NAMES.items():
+        for category, mode_title in generate_examples.SERVING_MODE_TITLES.items():
+            example_path = ROOT_DIR / "examples" / category / slug
+            example = generate_examples.Example(example_path, category)
+
+            assert example.title == f"{display_name}: {mode_title}"
+            assert example.nav_title == display_name
+
+
+def test_model_example_title_mismatch_is_rejected(tmp_path):
+    example_path = tmp_path / "mimo_audio"
+    example_path.mkdir()
+    (example_path / "README.md").write_text("# A different title\n", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="Model example title mismatch"):
+        generate_examples.Example(example_path, "offline_inference")
+
+
+def test_model_display_names_require_string_mapping(tmp_path, monkeypatch):
+    display_names_path = tmp_path / "model_display_names.yml"
+    display_names_path.write_text("- not a mapping\n", encoding="utf-8")
+    monkeypatch.setattr(generate_examples, "MODEL_DISPLAY_NAMES_FILE", display_names_path)
+
+    with pytest.raises(ValueError, match="MODEL_DISPLAY_NAMES_FILE"):
+        generate_examples.load_model_display_names()
diff --git a/tests/dreamzero/test_pipeline_state.py b/tests/dreamzero/test_pipeline_state.py
index a3a5aaec0..f00caf53c 100644
--- a/tests/dreamzero/test_pipeline_state.py
+++ b/tests/dreamzero/test_pipeline_state.py
@@ -4,7 +4,6 @@
 from collections import OrderedDict
 
 import pytest
-import torch
 
 from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
 from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState
@@ -46,14 +45,11 @@ def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None:
     assert "session-b" not in pipeline._states
 
 
-def test_dreamzero_state_cache_access_requires_initialization() -> None:
+def test_dreamzero_state_owns_no_kv_caches() -> None:
+    """KV (self- and cross-attention) is engine-owned since the AR-Diffusion
+    paged backend: the model-local cache accessors must be gone so nothing can
+    silently bypass the engine pool."""
     state = DreamZeroState()
 
-    with pytest.raises(RuntimeError, match="KV caches not initialized"):
-        state.get_kv_caches()
-
-    with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"):
-        state.get_crossattn_caches()
-
-    with pytest.raises(RuntimeError, match="create_kv_caches first"):
-        state.update_kv_cache(0, torch.empty(0))
+    for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"):
+        assert not hasattr(state, removed)
diff --git a/tests/e2e/accuracy/test_cosmos3_similarity.py b/tests/e2e/accuracy/test_cosmos3_similarity.py
index ff2350096..efa095fb1 100644
--- a/tests/e2e/accuracy/test_cosmos3_similarity.py
+++ b/tests/e2e/accuracy/test_cosmos3_similarity.py
@@ -42,8 +42,6 @@ def _server_args() -> list[str]:
     return [
         "--num-gpus",
         "1",
-        "--model-class-name",
-        "Cosmos3OmniDiffusersPipeline",
         "--stage-init-timeout",
         "900",
         "--init-timeout",
diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py
index 753050c41..00e77fc8f 100644
--- a/tests/e2e/accuracy/test_hunyuan_image3.py
+++ b/tests/e2e/accuracy/test_hunyuan_image3.py
@@ -331,7 +331,10 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag
         if ro and getattr(ro, "outputs", None):
             cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
         if not cot_text:
-            ar_text = getattr(out, "custom_output", {}).get("ar_generated_text")
+            multimodal_output = getattr(out, "multimodal_output", {}) or {}
+            metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {}
+            text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {}
+            ar_text = text_metadata.get("ar_generated_text") if isinstance(text_metadata, dict) else None
             if isinstance(ar_text, list):
                 cot_text = "\n".join(text for text in ar_text if text)
             else:
diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
index 6ac930f46..3f799c038 100644
--- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
+++ b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
@@ -48,6 +48,7 @@ _DEPLOY_CONFIG = {
             "trust_remote_code": True,
             "devices": "0,1,2,3",  # set dynamically by _write_deploy_config
             "vae_use_slicing": False,
+            "moe_backend": "flashinfer_cutlass",
             "vae_use_tiling": False,
             "parallel_config": {
                 "pipeline_parallel_size": 1,
diff --git a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py b/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py
index 97efab0fe..193d1b1e2 100644
--- a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py
+++ b/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py
@@ -288,7 +288,7 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
         else:
             # Both prompt_ids and prompt_embeds are None (e.g. during warmup/dummy run).
             # Return a minimal dummy output to avoid crashing.
-            return DiffusionOutput(output=None, custom_output={})
+            return DiffusionOutput(output=None)
 
         if isinstance(negative_prompt_ids, list):
             negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)
@@ -399,14 +399,20 @@ class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
             image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
 
         return DiffusionOutput(
-            output=_maybe_to_cpu(image),
+            output={
+                "payload": {
+                    "image": _maybe_to_cpu(image),
+                },
+                "metadata": {
+                    "prompt_embeddings": {
+                        "prompt_embeds": _maybe_to_cpu(prompt_embeds),
+                        "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
+                        "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
+                        "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
+                    },
+                },
+            },
             trajectory_latents=_maybe_to_cpu(all_latents),
             trajectory_log_probs=_maybe_to_cpu(all_log_probs),
             trajectory_timesteps=_maybe_to_cpu(all_timesteps),
-            custom_output={
-                "prompt_embeds": _maybe_to_cpu(prompt_embeds),
-                "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask),
-                "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds),
-                "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask),
-            },
         )
diff --git a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
index 9bcaf0e40..69952435e 100644
--- a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+++ b/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
@@ -390,20 +390,23 @@ class vLLMOmniHttpServerLocal:
 
         diffusion_output = self._to_tensor(final_res.images[0]).float() / 255.0
 
-        mm_output = final_res.custom_output or {}
+        metadata = (
+            final_res.multimodal_output.get("metadata", {}) if isinstance(final_res.multimodal_output, dict) else {}
+        )
+        prompt_embeddings = metadata.get("prompt_embeddings", {}) if isinstance(metadata, dict) else {}
 
         if sampling_params.get("logprobs", False):
-            all_log_probs = mm_output.get("all_log_probs")
+            all_log_probs = final_res.trajectory_log_probs
             log_probs = all_log_probs[0] if all_log_probs is not None else None
         else:
             log_probs = None
 
-        all_latents = mm_output.get("all_latents")
-        all_timesteps = mm_output.get("all_timesteps")
-        prompt_embeds = mm_output.get("prompt_embeds")
-        prompt_embeds_mask = mm_output.get("prompt_embeds_mask")
-        negative_prompt_embeds = mm_output.get("negative_prompt_embeds")
-        negative_prompt_embeds_mask = mm_output.get("negative_prompt_embeds_mask")
+        all_latents = final_res.trajectory_latents
+        all_timesteps = final_res.trajectory_timesteps
+        prompt_embeds = prompt_embeddings.get("prompt_embeds")
+        prompt_embeds_mask = prompt_embeddings.get("prompt_embeds_mask")
+        negative_prompt_embeds = prompt_embeddings.get("negative_prompt_embeds")
+        negative_prompt_embeds_mask = prompt_embeddings.get("negative_prompt_embeds_mask")
 
         extra_fields = {
             "all_latents": all_latents[0] if all_latents is not None else None,
diff --git a/tests/e2e/offline_inference/test_flux2_klein.py b/tests/e2e/offline_inference/test_flux2_klein.py
index 3bd815229..9656772c9 100644
--- a/tests/e2e/offline_inference/test_flux2_klein.py
+++ b/tests/e2e/offline_inference/test_flux2_klein.py
@@ -18,6 +18,8 @@ from tests.helpers.runtime import DiffusionResponse, OmniRunnerHandler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
+# NOTE: Most of the flux2klein tests have migrated to the common diffusion tests.
+# This model will also use the tiny model builder when run at the core model level.
 MODEL = "black-forest-labs/FLUX.2-klein-4B"
 
 _OMNI_RUNNER_PARAM = (MODEL, None)
@@ -70,41 +72,6 @@ def _send_inpaint_with_generator(
     )
 
 
-# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-def test_flux2_klein_can_accept_text_inputs(omni_runner_handler: OmniRunnerHandler):
-    omni_runner_handler.send_diffusion_request(
-        {
-            "model": MODEL,
-            "prompt": "a cup of coffee on the table",
-            "sampling_params": OmniDiffusionSamplingParams(num_inference_steps=2, seed=42),
-        }
-    )
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_flux2_klein_inpaint_basic(omni_runner_handler: OmniRunnerHandler):
-    input_image, mask_image = _create_test_inputs()
-    omni_runner_handler.send_diffusion_request(
-        {
-            "model": MODEL,
-            "prompt": "Fill in the masked area with a beautiful garden",
-            "multi_modal_data": {"image": input_image, "mask_image": mask_image},
-            "sampling_params": OmniDiffusionSamplingParams(
-                height=_HEIGHT,
-                width=_WIDTH,
-                num_inference_steps=_NUM_INFERENCE_STEPS,
-                guidance_scale=0.0,
-                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
-                num_outputs_per_prompt=1,
-            ),
-        }
-    )
-
-
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
diff --git a/tests/e2e/offline_inference/test_krea2_expansion.py b/tests/e2e/offline_inference/test_krea2_expansion.py
new file mode 100644
index 000000000..f00b4bb0b
--- /dev/null
+++ b/tests/e2e/offline_inference/test_krea2_expansion.py
@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""L4 functionality tests for the Krea 2 text-to-image diffusion pipeline (``Krea2Pipeline``).
+
+Runs against the public few-step distilled checkpoint ``krea/Krea-2-Turbo`` by default (override with the
+``KREA2_MODEL`` environment variable, e.g. ``krea/Krea-2-Raw`` for the Raw checkpoint or a local diffusers
+directory). They cover a basic functional smoke plus the layerwise-CPU-offload path that the pipeline declares
+via ``SupportsComponentDiscovery`` / ``_layerwise_offload_blocks_attrs``.
+"""
+
+import os
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunnerHandler
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
+
+MODEL = os.environ.get("KREA2_MODEL", "krea/Krea-2-Turbo")
+# vLLM-Omni-compatible PEFT repackaging of krea/Krea-2-LoRA-darkbrush (264 modules, r=alpha=32).
+LORA = os.environ.get("KREA2_LORA", "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA")
+PROMPT = "a fox in the snow, photorealistic"
+# darkbrush is a trigger-word LoRA: the trained monochrome ink-wash style only manifests when
+# its trigger phrase is present, so the LoRA test prompt appends it for a strong, stable signal.
+LORA_TRIGGER = "monochrome ink wash style"
+LORA_PROMPT = f"a fox in the snow, {LORA_TRIGGER}"
+
+pytestmark = [
+    pytest.mark.diffusion,
+    pytest.mark.full_model,
+]
+
+
+def _sampling() -> OmniDiffusionSamplingParams:
+    # Small resolution + few steps to keep the L4 case light. guidance_scale resolution is checkpoint-aware inside
+    # the pipeline (distilled -> no-CFG, Raw -> CFG), so this stays agnostic to which checkpoint KREA2_MODEL points at.
+    return OmniDiffusionSamplingParams(
+        height=512,
+        width=512,
+        num_inference_steps=8,
+        guidance_scale=0.0,
+        seed=42,
+    )
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
+def test_krea2_text_to_image_001(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Basic functional smoke: a single prompt produces a decoded image."""
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize(
+    "omni_runner",
+    [(MODEL, None, {"enable_layerwise_offload": True})],
+    indirect=True,
+)
+def test_krea2_layerwise_offload(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Exercise layerwise CPU offload on the DiT (SupportsComponentDiscovery + _layerwise_offload_blocks_attrs)."""
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.cache
+@pytest.mark.parametrize("omni_runner", [(MODEL, None, {"cache_backend": "cache_dit"})], indirect=True)
+def test_krea2_cache_dit(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Exercise Cache-DiT on Krea 2 via the custom Krea2Pipeline enabler (``enable_cache_for_krea2``).
+
+    Validates the docs' Cache-DiT support claim for Krea 2. ``has_separate_cfg`` is checkpoint-aware
+    (False for the distilled Turbo no-CFG path, True for the Raw CFG path); the default Turbo checkpoint
+    exercises the no-CFG branch.
+    """
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize(
+    "omni_runner",
+    # Standalone HSDP: use_hsdp=True with an explicit hsdp_shard_size shards the DiT weights
+    # across 2 ranks (replicate_size defaults to 1 -> world_size = 1 * 2 = 2). HSDP must keep
+    # tensor_parallel_size/data_parallel_size at 1 (see vllm_omni/diffusion/data.py __post_init__).
+    # Also validated locally at 8 GPUs (hsdp_shard_size=8); CI shard provisions 2.
+    [(MODEL, None, {"use_hsdp": True, "hsdp_shard_size": 2})],
+    indirect=True,
+)
+def test_krea2_hsdp(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Exercise Hybrid Sharded Data Parallel (HSDP) weight sharding on the DiT across 2 GPUs.
+
+    Validates the docs' HSDP-support claim for Krea 2. Scheduled on the multi-GPU (distributed_cuda)
+    shard by ``num_cards=2``; skipped automatically on boxes with fewer than 2 CUDA devices.
+    """
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize(
+    "omni_runner",
+    # VAE patch-parallel decode: partition the VAE decode across the TP world. Mirrors the online
+    # "tp_vae_patch" case (tensor_parallel_size=2 + vae_patch_parallel_size=2) from the wan2.2/z-image
+    # expansion suites, so vae_patch_parallel_size must equal the world size established by TP.
+    # Also validated locally at 8 GPUs (tp=8, vae_patch=8); CI shard provisions 2.
+    [(MODEL, None, {"tensor_parallel_size": 2, "vae_patch_parallel_size": 2})],
+    indirect=True,
+)
+def test_krea2_vae_patch_parallel(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Exercise VAE patch-parallel (decode) across 2 GPUs.
+
+    Validates the docs' VAE-patch-parallel (decode) support claim for Krea 2. Scheduled on the
+    multi-GPU (distributed_cuda) shard by ``num_cards=2``; skipped automatically on boxes with
+    fewer than 2 CUDA devices.
+    """
+    omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()})
+
+
+def _generate(handler: OmniRunnerHandler, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> np.ndarray:
+    sp = _sampling()
+    sp.lora_request = lora_request
+    sp.lora_scale = lora_scale
+    resp = handler.send_diffusion_request({"model": MODEL, "prompt": LORA_PROMPT, "sampling_params": sp})
+    return np.asarray(resp.images[0].convert("RGB"), dtype=np.int16)
+
+
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True)
+def test_krea2_lora(omni_runner_handler: OmniRunnerHandler) -> None:
+    """Validate diffusion LoRA on Krea 2: visible effect, scale sensitivity, clean deactivation.
+
+    Uses ``LORA`` (a vLLM-Omni-compatible PEFT repackaging of ``krea/Krea-2-LoRA-darkbrush``,
+    264 modules matched via the ReplicatedLinear projections) passed per-request through
+    ``OmniDiffusionSamplingParams.lora_request`` / ``lora_scale`` — the same fields the
+    diffusion LoRA manager reads to activate/deactivate the adapter.
+    """
+    lora_request = LoRARequest(lora_name="darkbrush", lora_int_id=stable_lora_int_id(LORA), lora_path=LORA)
+
+    baseline = _generate(omni_runner_handler, None)
+    img_1x = _generate(omni_runner_handler, lora_request, lora_scale=1.0)
+    img_2x = _generate(omni_runner_handler, lora_request, lora_scale=2.0)
+    restored = _generate(omni_runner_handler, None)
+
+    diff_1x = np.abs(baseline - img_1x).mean()
+    diff_2x = np.abs(baseline - img_2x).mean()
+    diff_restored = np.abs(baseline - restored).mean()
+
+    # (a) Adapter has a visible effect at both scales.
+    assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
+    assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
+    # (b) Scale changes the output (stronger adapter -> larger perturbation).
+    assert not np.isclose(diff_1x, diff_2x, atol=1.0), f"LoRA scale had no effect: 1x={diff_1x:.2f}, 2x={diff_2x:.2f}"
+    # (c) Passing lora_request=None cleanly deactivates: same seed -> back to baseline (modulo fp drift).
+    assert diff_restored < 5.0, f"LoRA did not deactivate cleanly: diff_restored={diff_restored:.2f}"
diff --git a/tests/e2e/offline_inference/test_magi_human_expansion.py b/tests/e2e/offline_inference/test_magi_human_expansion.py
index 5ce016988..e6b3c6a42 100644
--- a/tests/e2e/offline_inference/test_magi_human_expansion.py
+++ b/tests/e2e/offline_inference/test_magi_human_expansion.py
@@ -22,6 +22,7 @@ _OMNI_RUNNER_PARAM = (
     {
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
+        "vae_use_tiling": True,
     },
 )
 
@@ -32,6 +33,7 @@ _OMNI_RUNNER_PARAM_OFFLOAD = (
         "init_timeout": 1200,
         "tensor_parallel_size": 2,
         "enable_cpu_offload": True,
+        "vae_use_tiling": True,
     },
 )
 
@@ -45,10 +47,6 @@ pytestmark = [
     ),
 ]
 
-_SKIP_ISSUE_3236 = pytest.mark.skip(
-    reason="https://github.com/vllm-project/vllm-omni/issues/3236",
-)
-
 
 def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     """Validate that the MP4 contains meaningful video and audio tracks."""
@@ -74,7 +72,6 @@ def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None:
     container.close()
 
 
-@_SKIP_ISSUE_3236
 @hardware_test(res={"cuda": "H100"}, num_cards=2)
 def test_magi_human_e2e(omni_runner: OmniRunner):
     """End-to-end test for MagiHuman generating video and audio."""
diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py
index b0d82dbc4..04308c4a7 100644
--- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py
+++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py
@@ -5,7 +5,7 @@ Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
 match a golden reference.
 
 Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
-Stage config: ``get_deploy_config_path("mammoth_moda2.yaml")`` → ``vllm_omni/deploy/mammoth_moda2.yaml``
+Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
 
 Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
   Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
@@ -36,9 +36,9 @@ _VISION_END_TOKEN_ID = 151653  # "<|vision_end|>"
 _AR_PATCH_SIZE = 16
 
 MODEL_PATH = "bytedance-research/MammothModa2-Preview"
-T2I_STAGE_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
+T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
 
-_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_STAGE_CONFIG)
+_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG})
 
 # Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
 _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
diff --git a/tests/e2e/offline_inference/test_omnivoice_expansion.py b/tests/e2e/offline_inference/test_omnivoice_expansion.py
index 692278ef7..e3a138112 100644
--- a/tests/e2e/offline_inference/test_omnivoice_expansion.py
+++ b/tests/e2e/offline_inference/test_omnivoice_expansion.py
@@ -18,13 +18,15 @@ from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "k2-fsa/OmniVoice"
-STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml")
+DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml")
 
-# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime
+# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs.
+# The migrated test passes deploy_config through extra Omni kwargs.
 _OMNI_RUNNER_PARAM = (
     MODEL,
-    STAGE_CONFIG,
+    None,
     {
+        "deploy_config": DEPLOY_CONFIG,
         "trust_remote_code": True,
         "log_stats": True,
     },
diff --git a/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py b/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
index 416a5009c..4fa2d4b6f 100644
--- a/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
+++ b/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
@@ -34,12 +34,10 @@ def get_cuda_graph_config():
             "stages": {
                 0: {
                     "max_num_seqs": 1,
-                    "gpu_memory_utilization": 0.2,
                     "enforce_eager": True,
                     "async_scheduling": False,
                 },
                 1: {
-                    "gpu_memory_utilization": 0.2,
                     "enforce_eager": True,
                     "async_scheduling": False,
                 },
diff --git a/tests/e2e/online_serving/test_flux2_klein.py b/tests/e2e/online_serving/test_flux2_klein.py
deleted file mode 100644
index f8c705df2..000000000
--- a/tests/e2e/online_serving/test_flux2_klein.py
+++ /dev/null
@@ -1,62 +0,0 @@
-"""
-Tests for Flux2 Klein; currently Dev is implemented separately,
-but ideally these models will fold together in the future.
-"""
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.runtime import (
-    OmniServer,
-    OmniServerParams,
-    OpenAIClientHandler,
-    dummy_messages_from_mix_data,
-)
-
-FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4)
-POSITIVE_PROMPT = "A cat sitting on a windowsill"
-NEGATIVE_PROMPT = "blurry, low quality"
-
-
-# For now, we only test a good configuration for this model to keep the CI
-# light, but in the near future may explore using random models to check
-# cross-feature compatibility more generally.
-def _get_diffusion_feature_cases(model: str):
-    return [
-        # FP8 / Hybrid sequence parallelism
-        pytest.param(
-            OmniServerParams(
-                model=model,
-            ),
-            id="default",
-            marks=FOUR_CARD_FEATURE_MARKS,
-        ),
-    ]
-
-
-@pytest.mark.advanced_model
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@pytest.mark.parametrize(
-    "omni_server",
-    _get_diffusion_feature_cases(
-        model="black-forest-labs/FLUX.2-klein-4B",
-    ),
-    indirect=True,
-)
-def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
-    messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT)
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "extra_body": {
-            "height": 512,
-            "width": 512,
-            "num_inference_steps": 2,
-            "negative_prompt": NEGATIVE_PROMPT,
-            "true_cfg_scale": 4.0,
-            "seed": 42,
-        },
-    }
-
-    openai_client.send_diffusion_request(request_config)
diff --git a/tests/e2e/online_serving/test_flux2_klein_expansion.py b/tests/e2e/online_serving/test_flux2_klein_expansion.py
deleted file mode 100644
index 98ccf0014..000000000
--- a/tests/e2e/online_serving/test_flux2_klein_expansion.py
+++ /dev/null
@@ -1,163 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-End-to-end tests for Flux2 Klein inpainting in online serving mode.
-
-Uses /v1/images/edits endpoint which is the correct API for image inpainting.
-"""
-
-import base64
-from io import BytesIO
-
-import httpx
-import pytest
-from PIL import Image, ImageDraw
-
-from tests.helpers.runtime import OmniServer, OmniServerParams
-
-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
-
-MODEL = "black-forest-labs/FLUX.2-klein-4B"
-
-_HEIGHT = 512
-_WIDTH = 512
-_NUM_INFERENCE_STEPS = 4
-
-
-def _get_diffusion_feature_cases(model: str):
-    return [
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--tensor-parallel-size", "2", "--enable-cpu-offload"],
-            ),
-            id="tp2_cpu_offload",
-        ),
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=["--tensor-parallel-size", "2"],
-            ),
-            id="tp2_basic",
-        ),
-    ]
-
-
-def _image_to_base64_jpeg(image: Image.Image) -> str:
-    buffer = BytesIO()
-    image.save(buffer, format="JPEG")
-    buffer.seek(0)
-    return base64.b64encode(buffer.read()).decode("utf-8")
-
-
-def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str:
-    mask = Image.new("L", (width, height), 0)
-    draw = ImageDraw.Draw(mask)
-    draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255)
-    return _image_to_base64_jpeg(mask)
-
-
-def _compare_images(img1: Image.Image, img2: Image.Image) -> bool:
-    return list(img1.getdata()) == list(img2.getdata())
-
-
-def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs):
-    url = f"http://{host}:{port}/v1/images/edits"
-    files = {
-        "image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"),
-        "mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"),
-    }
-    data = {"prompt": prompt, "model": model, **kwargs}
-    with httpx.Client(timeout=60.0) as client:
-        response = client.post(url, files=files, data=data)
-        response.raise_for_status()
-        return response.json()
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_basic(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-
-    result = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt="Fill in the masked area with a beautiful garden",
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=42,
-    )
-
-    assert "data" in result and len(result["data"]) == 1
-    img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1]
-    img = Image.open(BytesIO(base64.b64decode(img_data)))
-    assert img.size == (_WIDTH, _HEIGHT)
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-    prompt = "A red flower in a field"
-
-    result1 = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt=prompt,
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=12345,
-    )
-
-    result2 = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt=prompt,
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=1,
-        seed=12345,
-    )
-
-    img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1]
-    img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1]
-
-    img1 = Image.open(BytesIO(base64.b64decode(img1_data)))
-    img2 = Image.open(BytesIO(base64.b64decode(img2_data)))
-
-    assert _compare_images(img1, img2), (
-        "Same input with same seed should produce identical output. This is critical for offline/online consistency."
-    )
-
-
-@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True)
-def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer):
-    input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128)))
-    mask_b64 = _create_test_mask_base64()
-
-    result = _send_edit_request(
-        host=omni_server.host,
-        port=omni_server.port,
-        model=MODEL,
-        image_b64=input_image_b64,
-        mask_b64=mask_b64,
-        prompt="A beautiful landscape",
-        guidance_scale=1.0,
-        num_inference_steps=_NUM_INFERENCE_STEPS,
-        n=2,
-        seed=42,
-    )
-
-    assert "data" in result and len(result["data"]) == 2
diff --git a/tests/e2e/online_serving/test_higgs_audio_v3.py b/tests/e2e/online_serving/test_higgs_audio_v3.py
index 2ead1b5da..d26e6bff4 100644
--- a/tests/e2e/online_serving/test_higgs_audio_v3.py
+++ b/tests/e2e/online_serving/test_higgs_audio_v3.py
@@ -127,7 +127,7 @@ class TestHiggsAudioV3OnlineHappyPath:
         windows stitch into a coherent PCM stream. The byte-count gate is the same as
         the sync paths; per-chunk audio content is verified offline against Whisper.
 
-        NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
+        NOTE: ``min_hnr_db=0.0`` sit below the typical speech-noise floor so the
         check still catches catastrophic codec failure (silence, white noise, sample
         scramble all give HNR << 0) while allowing for the sliding-window codec's
         slightly-noisier-than-sync output. The default 1.0 dB threshold has only
diff --git a/tests/e2e/online_serving/test_krea2_expansion.py b/tests/e2e/online_serving/test_krea2_expansion.py
new file mode 100644
index 000000000..10aa520fb
--- /dev/null
+++ b/tests/e2e/online_serving/test_krea2_expansion.py
@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""L4 online-serving tests for the Krea 2 text-to-image diffusion pipeline.
+
+Exercises the OpenAI-compatible ``/v1/images/generations`` route (served via
+``vllm serve <model> --omni --enforce-eager``) rather than the offline engine path
+covered by ``tests/e2e/offline_inference/test_krea2_expansion.py``.
+
+Runs against the public few-step distilled checkpoint ``krea/Krea-2-Turbo`` by default
+(override with the ``KREA2_MODEL`` environment variable, e.g. ``krea/Krea-2-Raw`` for the
+Raw checkpoint or a local diffusers directory). Krea 2 shares the Qwen-Image VAE and a
+similar single-stream MMDiT, so this test mirrors the structure/fixtures/markers of
+``test_qwen_image_expansion.py``.
+
+Coverage mirrors the offline suite: a basic functional smoke plus the layerwise-CPU-offload
+path that the pipeline declares via ``SupportsComponentDiscovery`` /
+``_layerwise_offload_blocks_attrs``. Sizing is kept light (512x512, few steps) so the case
+fits single-GPU H100 CI rather than the full 2048x2048 distilled resolution.
+"""
+
+import base64
+import os
+from io import BytesIO
+
+import numpy as np
+import pytest
+import requests
+from PIL import Image
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = os.environ.get("KREA2_MODEL", "krea/Krea-2-Turbo")
+# vLLM-Omni-compatible PEFT repackaging of krea/Krea-2-LoRA-darkbrush (264 modules, r=alpha=32).
+LORA = os.environ.get("KREA2_LORA", "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA")
+PROMPT = "a fox in the snow, photorealistic"
+# darkbrush is a trigger-word LoRA; include its trigger phrase so the LoRA case has a strong signal.
+LORA_PROMPT = "a fox in the snow, monochrome ink wash style"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+def _get_diffusion_feature_cases(model: str):
+    """Return single-card L4 online-serving cases for Krea 2.
+
+    Mirrors the offline suite: a basic functional smoke plus layerwise CPU offload.
+    """
+    return [
+        # Basic functional smoke (single-card).
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--enforce-eager"],
+            ),
+            id="single_card_basic",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        # Layerwise CPU offload (single-card): SupportsComponentDiscovery +
+        # _layerwise_offload_blocks_attrs.
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=["--enforce-eager", "--enable-layerwise-offload"],
+            ),
+            id="single_card_layerwise_offload",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(MODEL),
+    indirect=True,
+)
+def test_krea2(omni_server: OmniServer, openai_client: OpenAIClientHandler):
+    """Serve Krea 2 and validate the OpenAI-compatible image-generation route.
+
+    guidance_scale resolution is checkpoint-aware inside the pipeline (distilled -> no-CFG,
+    Raw -> CFG), so guidance_scale=0.0 stays agnostic to which checkpoint KREA2_MODEL points at.
+    Validation is delegated to assert_diffusion_response, which checks output dimensions and
+    basic correctness.
+    """
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 8,
+            "guidance_scale": 0.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+def _basic_lora_payload(lora: dict | None = None) -> dict:
+    payload = {
+        "prompt": LORA_PROMPT,
+        "n": 1,
+        "size": "512x512",
+        "num_inference_steps": 8,
+        "guidance_scale": 0.0,
+        "seed": 42,
+    }
+    if lora is not None:
+        payload["lora"] = lora
+    return payload
+
+
+def _post_image(server: OmniServer, payload: dict) -> np.ndarray:
+    url = f"http://{server.host}:{server.port}/v1/images/generations"
+    resp = requests.post(url, json=payload, headers={"Authorization": "Bearer EMPTY"}, timeout=900)
+    resp.raise_for_status()
+    b64 = resp.json()["data"][0]["b64_json"]
+    img = Image.open(BytesIO(base64.b64decode(b64)))
+    img.load()
+    return np.asarray(img.convert("RGB"), dtype=np.int16)
+
+
+@pytest.mark.parametrize(
+    "omni_server",
+    [
+        pytest.param(
+            OmniServerParams(model=MODEL, server_args=["--enforce-eager"]),
+            id="single_card_lora",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        )
+    ],
+    indirect=True,
+)
+def test_krea2_lora(omni_server: OmniServer):
+    """Serve Krea 2 and validate per-request LoRA over ``/v1/images/generations``.
+
+    Passes ``LORA`` via the top-level ``lora`` object (``{name, path, scale}``) the diffusion
+    image route parses into a ``LoRARequest`` — no server-side LoRA flag needed. Confirms the
+    adapter has a visible effect and that omitting ``lora`` cleanly restores the base output.
+    """
+    baseline = _post_image(omni_server, _basic_lora_payload())
+    with_lora = _post_image(omni_server, _basic_lora_payload({"name": "darkbrush", "path": LORA, "scale": 1.0}))
+    restored = _post_image(omni_server, _basic_lora_payload())
+
+    diff_lora = np.abs(baseline - with_lora).mean()
+    diff_restored = np.abs(baseline - restored).mean()
+
+    assert diff_lora > 0.5, f"LoRA had no visible effect over the serving path: diff={diff_lora}"
+    assert diff_restored < 5.0, f"LoRA did not deactivate cleanly: diff_restored={diff_restored:.2f}"
diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py
index 0893ef1d2..ffddafbc0 100644
--- a/tests/e2e/online_serving/test_qwen3_omni.py
+++ b/tests/e2e/online_serving/test_qwen3_omni.py
@@ -226,3 +226,27 @@ def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None
     }
 
     _run_prefix_cache_check(openai_client, request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_completions_rejected_for_thinker_talker(omni_server, openai_client) -> None:
+    """Ensure Thinker-talker models reject /v1/completions; we do this because the
+    thinker-talker handoff implementations currently use ChatML <|im_start|> and
+    <|im_end|> markers to segment the input sequence; when we don't have them,
+    the talker does not get any embeddings, which breaks the server.
+    """
+    responses = openai_client.send_completions_http_request(
+        {
+            "json": {
+                "model": omni_server.model,
+                "prompt": "Hello, how are you?",
+                "max_tokens": 10,
+            },
+        },
+        err_code=400,
+    )
+    assert not responses[0].success
diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py
index db3d7a8d1..3ded285ce 100644
--- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py
+++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py
@@ -7,6 +7,7 @@ E2E Online tests for Qwen3-Omni model.
 import os
 
 import pytest
+from openai import BadRequestError
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
@@ -410,6 +411,44 @@ def test_audio_in_video_002(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None:
+    """
+    Regression: ``use_audio_in_video`` with default video loader sampling (no
+    ``media_io_kwargs.video.fps``) after a prior request that pinned fps.
+
+    Before the fix the second request could partial-hit mm-cache (audio cached,
+    video re-processed without audio) and fail with ``StopIteration`` in
+    ``replace_multimodal_special_tokens``.
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 128, embed_audio=True)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        content_text="Reply with one word: OK.",
+    )
+
+    base_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "use_audio_in_video": True,
+        "modalities": ["text"],
+    }
+
+    # Prime mm-cache with explicit loader fps (historical failure mode).
+    openai_client.send_omni_request(
+        {
+            **base_config,
+            "extra_body": {"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}},
+        }
+    )
+
+    # Default loader sampling: must succeed without explicit media_io fps.
+    openai_client.send_omni_request(dict(base_config))
+
+
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_one_word_prompt_001(omni_server, openai_client) -> None:
@@ -572,3 +611,21 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
     text = responses[0].text_content if responses else ""
     word_count = len(text.split())
     assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True)
+def test_invalid_audio_format_rejected(omni_server, openai_client) -> None:
+    """Ensure invalid audio format is rejected with 400 before streaming starts."""
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt(),
+    )
+    with pytest.raises(BadRequestError, match="audio format"):
+        openai_client.client.chat.completions.create(
+            model=omni_server.model,
+            messages=messages,
+            modalities=["text", "audio"],
+            audio={"voice": "alloy", "format": "aac"},
+            stream=True,
+        )
diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py
index c26e21039..315fba6ef 100644
--- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py
+++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py
@@ -7,7 +7,7 @@ and are supported by Qwen-Image-Layered model.
 Kept cases (maximum feature coverage, no cpu-offload):
   sp_001              : cache_dit + Ulysses-SP 2          (2×H100)
   cfg_parallel_002    : cache_dit + CFG-Parallel 2        (2×H100)
-  layers_guard_001    : layerwise offload + layers=3      (1×H100, issue #1969 guard)
+  layers_guard_001    : layerwise offload + layers=2      (1×H100, issue #1969 guard)
 
 Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload.
 """
@@ -116,8 +116,8 @@ LAYERS_GUARD_CASES = [
             model=MODEL,
             server_args=["--enable-layerwise-offload"],
         ),
-        3,
-        id="layers_guard_001_layers3",
+        2,
+        id="layers_guard_001_layers2",
         marks=SINGLE_CARD_FEATURE_MARKS,
     ),
 ]
diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py
index 31850401b..d5a50cf29 100644
--- a/tests/engine/test_async_omni_engine_stage_init.py
+++ b/tests/engine/test_async_omni_engine_stage_init.py
@@ -739,7 +739,7 @@ def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(m
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(
@@ -772,7 +772,7 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch):
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(
@@ -798,7 +798,7 @@ def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stag
     monkeypatch.setattr(
         engine_mod,
         "load_and_resolve_stage_configs",
-        lambda *args, **kwargs: ("dummy-config", [stage_cfg]),
+        lambda *args, **kwargs: ("dummy-config", [stage_cfg], None),
     )
 
     _config_path, stage_configs = engine._resolve_stage_configs(
@@ -874,3 +874,54 @@ def test_build_engine_args_dict_uses_diffusion_attention_config_key():
 
     assert "attention_config" not in engine_args_dict
     assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN"
+
+
+def test_omni_master_server_allocates_globally_unique_route_ports(monkeypatch):
+    """Regression: two stages must never draw the same ZMQ port.
+
+    ``get_open_ports_list`` only dedups within a single call, so per-route
+    allocation used to let a later stage reuse an earlier stage's port. The
+    second engine to ``bind()`` then died with ``zmq.error.ZMQError: Address
+    already in use`` (flaky multi-stage startup, e.g. Qwen3-Omni thinker/talker/
+    code2wav). ``OmniMasterServer`` now dedups every port it hands out.
+    """
+    import itertools
+
+    from vllm_omni.engine import stage_engine_startup as ses
+
+    # A colliding prefix (repeats + the master port 9000) followed by an endless
+    # fresh stream, so the only way to succeed is to redraw the collisions.
+    supply = itertools.chain([9000, 9000, 9001, 9001, 9000], itertools.count(9002))
+
+    def fake_get_open_ports_list(count):
+        return [next(supply) for _ in range(count)]
+
+    monkeypatch.setattr(ses, "get_open_ports_list", fake_get_open_ports_list)
+
+    server = ses.OmniMasterServer(
+        master_address="127.0.0.1",
+        master_port=9000,  # seed: a route must not reuse the registration port
+        stage_ids=[0, 1, 2],
+    )
+
+    ports = []
+    for sid in (0, 1, 2):
+        alloc = server.get_allocation(sid)
+        for addr in (
+            alloc.handshake_bind_address,
+            alloc.input_bind_address,
+            alloc.output_bind_address,
+        ):
+            ports.append(ses._port_from_zmq_address(addr))
+
+    assert len(ports) == len(set(ports)), f"duplicate route ports allocated: {ports}"
+    assert 9000 not in ports, "route reused the master registration port"
+
+
+def test_port_from_zmq_address_parsing():
+    from vllm_omni.engine.stage_engine_startup import _port_from_zmq_address
+
+    assert _port_from_zmq_address("tcp://127.0.0.1:34277") == 34277
+    assert _port_from_zmq_address(None) is None
+    assert _port_from_zmq_address("ipc:///tmp/sock") is None
+    assert _port_from_zmq_address("tcp://host:not-a-port") is None
diff --git a/tests/engine/test_output_modality.py b/tests/engine/test_output_modality.py
index 0abfe6ea0..ff15e5824 100644
--- a/tests/engine/test_output_modality.py
+++ b/tests/engine/test_output_modality.py
@@ -14,7 +14,7 @@ import torch
 # ── Load modules without triggering vllm_omni.__init__ ─────────────
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-_ENGINE_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "engine"
+_OUTPUTS_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "outputs"
 
 
 def _load_module(name: str, filepath: Path):
@@ -26,12 +26,16 @@ def _load_module(name: str, filepath: Path):
 
 
 _om_mod = _load_module(
-    "vllm_omni.engine.output_modality",
-    _ENGINE_DIR / "output_modality.py",
+    "vllm_omni.outputs.output_modality",
+    _OUTPUTS_DIR / "output_modality.py",
+)
+_load_module(
+    "vllm_omni.outputs.utils",
+    _OUTPUTS_DIR / "utils.py",
 )
 _mm_mod = _load_module(
-    "vllm_omni.engine.mm_outputs",
-    _ENGINE_DIR / "mm_outputs.py",
+    "vllm_omni.outputs.mm_outputs",
+    _OUTPUTS_DIR / "mm_outputs.py",
 )
 
 OutputModality = _om_mod.OutputModality
diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py
index f9b08ca15..86e71fcbc 100644
--- a/tests/engine/test_output_processor.py
+++ b/tests/engine/test_output_processor.py
@@ -11,8 +11,8 @@ from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.v1.engine import FinishReason
 
-from vllm_omni.engine.output_modality import OutputModalityNames
-from vllm_omni.engine.output_processor import OmniRequestState
+from vllm_omni.outputs.output_modality import OutputModalityNames
+from vllm_omni.outputs.output_processor import OmniRequestState
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
diff --git a/tests/engine/test_strategy_lb_policy.py b/tests/engine/test_strategy_lb_policy.py
new file mode 100644
index 000000000..8531c2a99
--- /dev/null
+++ b/tests/engine/test_strategy_lb_policy.py
@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for strategy-derived omni_lb_policy threading into AsyncOmniEngine.
+
+``AsyncOmniEngine._apply_strategy_lb_policy`` only touches ``self._omni_lb_policy``
+and module-level logging, so we exercise it with a lightweight stub ``self`` to
+avoid standing up a full engine (which needs GPUs/model weights).
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Stub:
+    def __init__(self, current: str = "random") -> None:
+        self._omni_lb_policy = current
+
+
+_apply = AsyncOmniEngine._apply_strategy_lb_policy
+
+
+def test_no_derived_policy_is_noop():
+    stub = _Stub("random")
+    _apply(stub, None, {})
+    assert stub._omni_lb_policy == "random"
+
+
+def test_derived_overrides_default_random():
+    # "random" is the engine default == "unset", so the strategy value wins.
+    stub = _Stub("random")
+    _apply(stub, "round-robin", {})
+    assert stub._omni_lb_policy == "round-robin"
+
+
+def test_derived_overrides_when_flag_absent():
+    stub = _Stub("random")
+    _apply(stub, "least-queue-length", {"omni_lb_policy": None})
+    assert stub._omni_lb_policy == "least-queue-length"
+
+
+def test_explicit_matching_flag_is_fine():
+    stub = _Stub("round-robin")
+    _apply(stub, "round-robin", {"omni_lb_policy": "round-robin"})
+    assert stub._omni_lb_policy == "round-robin"
+
+
+def test_explicit_conflicting_flag_raises():
+    stub = _Stub("round-robin")
+    with pytest.raises(ValueError, match="Conflicting load-balancer policy"):
+        _apply(stub, "least-queue-length", {"omni_lb_policy": "round-robin"})
+
+
+def test_explicit_random_is_treated_as_unset():
+    # An explicit --omni-lb-policy=random is indistinguishable from the default,
+    # so a strategy value still overrides it (documented Phase-1 semantics).
+    stub = _Stub("random")
+    _apply(stub, "round-robin", {"omni_lb_policy": "random"})
+    assert stub._omni_lb_policy == "round-robin"
diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py
new file mode 100644
index 000000000..2f1f487d9
--- /dev/null
+++ b/tests/entrypoints/openai_api/test_audio_format.py
@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for audio output format handling in chat completions.
+
+Covers:
+- #4716: audio.format parameter must be respected, not hardcoded to WAV
+- Format validation rejects unsupported formats
+- pcm16 is mapped to pcm for soundfile compatibility
+- Default format is WAV when not specified
+- create_audio encodes correctly for each supported format
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    DEFAULT_AUDIO_FORMAT,
+    SUPPORTED_AUDIO_FORMATS,
+    SUPPORTED_CHAT_AUDIO_FORMATS,
+    CreateAudio,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestAudioFormatConstants:
+    def test_default_format_is_wav(self):
+        assert DEFAULT_AUDIO_FORMAT == "wav"
+
+    def test_supported_formats_no_aac(self):
+        assert "aac" not in SUPPORTED_AUDIO_FORMATS
+
+    def test_chat_formats_include_pcm16(self):
+        assert "pcm16" in SUPPORTED_CHAT_AUDIO_FORMATS
+
+    def test_chat_formats_superset_of_audio_formats(self):
+        assert SUPPORTED_AUDIO_FORMATS <= SUPPORTED_CHAT_AUDIO_FORMATS
+
+
+class TestCreateAudio:
+    @pytest.fixture
+    def mixin(self):
+        return AudioMixin()
+
+    @pytest.fixture
+    def audio_tensor(self):
+        return np.sin(np.linspace(0, 2 * np.pi, 24000)).astype(np.float32)
+
+    @pytest.mark.parametrize("fmt", ["wav", "mp3", "flac", "opus", "pcm"])
+    def test_create_audio_supported_formats(self, mixin, audio_tensor, fmt):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format=fmt,
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert len(response.audio_data) > 0
+
+    def test_wav_magic_bytes(self, mixin, audio_tensor):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="wav",
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert response.audio_data[:4] == b"RIFF"
+        assert response.media_type == "audio/wav"
+
+    def test_mp3_encoding(self, mixin, audio_tensor):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="mp3",
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert response.media_type == "audio/mpeg"
+        assert response.audio_data[:4] != b"RIFF"
+
+    def test_flac_magic_bytes(self, mixin, audio_tensor):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="flac",
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert response.audio_data[:4] == b"fLaC"
+        assert response.media_type == "audio/flac"
+
+    def test_unsupported_format_falls_back_to_default(self, mixin, audio_tensor):
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="aac",
+            speed=1.0,
+            base64_encode=False,
+        )
+        response = mixin.create_audio(audio_obj)
+        assert response.audio_data[:4] == b"RIFF"
+        assert response.media_type == "audio/wav"
+
+    def test_base64_encoding(self, mixin, audio_tensor):
+        import base64
+
+        audio_obj = CreateAudio(
+            audio_tensor=audio_tensor,
+            sample_rate=24000,
+            response_format="wav",
+            speed=1.0,
+            base64_encode=True,
+        )
+        response = mixin.create_audio(audio_obj)
+        decoded = base64.b64decode(response.audio_data)
+        assert decoded[:4] == b"RIFF"
+
+
+class TestResolveAudioFormat:
+    """Test _resolve_audio_format via the serving chat class."""
+
+    @pytest.fixture
+    def serving_chat(self):
+        from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+        return object.__new__(OmniOpenAIServingChat)
+
+    def _make_request(self, audio_params=None):
+        from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+
+        req = ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "hello"}],
+        )
+        if audio_params is not None:
+            req.audio = audio_params
+        return req
+
+    def test_default_format_when_no_audio_params(self, serving_chat):
+        request = self._make_request()
+        result = serving_chat._resolve_audio_format(request)
+        assert result == "wav"
+
+    def test_extracts_mp3_format(self, serving_chat):
+        request = self._make_request({"format": "mp3", "voice": "alloy"})
+        result = serving_chat._resolve_audio_format(request)
+        assert result == "mp3"
+
+    def test_pcm16_mapped_to_pcm(self, serving_chat):
+        request = self._make_request({"format": "pcm16", "voice": "alloy"})
+        result = serving_chat._resolve_audio_format(request)
+        assert result == "pcm"
+
+    def test_invalid_format_returns_error(self, serving_chat):
+        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+
+        request = self._make_request({"format": "aac", "voice": "alloy"})
+        result = serving_chat._resolve_audio_format(request)
+        assert isinstance(result, ErrorResponse)
+        assert "aac" in result.error.message
+
+    def test_all_supported_formats_accepted(self, serving_chat):
+        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+
+        for fmt in SUPPORTED_CHAT_AUDIO_FORMATS:
+            request = self._make_request({"format": fmt, "voice": "alloy"})
+            result = serving_chat._resolve_audio_format(request)
+            assert not isinstance(result, ErrorResponse), f"Format {fmt} should be accepted"
diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py
index 31a82102a..5a7ef7bf0 100644
--- a/tests/entrypoints/openai_api/test_image_server.py
+++ b/tests/entrypoints/openai_api/test_image_server.py
@@ -1244,7 +1244,7 @@ def test_parameter_validation():
 
     # Invalid layers for layered models (must stay within the backend-supported range)
     with pytest.raises(ValueError):
-        ImageGenerationRequest(prompt="test", layers=2)
+        ImageGenerationRequest(prompt="test", layers=1)
 
     with pytest.raises(ValueError):
         ImageGenerationRequest(prompt="test", layers=11)
@@ -1663,13 +1663,13 @@ def test_image_edit_invalid_layers(async_omni_test_client):
         files=[("image", img_bytes)],
         data={
             "prompt": "test",
-            "layers": 2,
+            "layers": 1,
         },
     )
     assert response.status_code == 400
     detail = response.json()["detail"]
     assert "Invalid layers" in detail
-    assert "layers must be between 3 and 10 inclusive" in detail
+    assert "layers must be between 2 and 10 inclusive" in detail
 
     # Test layers above the supported range
     response = async_omni_test_client.post(
@@ -1683,7 +1683,7 @@ def test_image_edit_invalid_layers(async_omni_test_client):
     assert response.status_code == 400
     detail = response.json()["detail"]
     assert "Invalid layers" in detail
-    assert "layers must be between 3 and 10 inclusive" in detail
+    assert "layers must be between 2 and 10 inclusive" in detail
 
 
 def test_image_edit_resolution_and_size_conflict(async_omni_test_client):
diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
index 15e257959..50f8eef91 100644
--- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
+++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py
@@ -198,12 +198,28 @@ def _assert_realtime_smoke(result: dict) -> None:
     assert result["output_sample_rate"] > 0
 
 
-def _assert_realtime_accuracy(result: dict) -> None:
+def _assert_realtime_accuracy(
+    result: dict,
+    whisper_model_size: str = "large-v3",
+) -> None:
+    """Assert that whisper transcription of audio output matches model text.
+
+    Args:
+        result: Roundtrip result dict from ``_run_realtime_audio_roundtrip``.
+        whisper_model_size: Whisper model used to transcribe the generated audio
+                   for the accuracy check. Defaults to ``large-v3``: the default
+                   ``small`` model mishears short Chinese TTS clips (observed:
+                   北京→韦京 and a dropped leading sentence, sim=0.443), which
+                   caused spurious sim<0.8 failures under async_chunk codec
+                   variability even though audio generation was correct. large-v3
+                   transcribes these clips reliably, so a failure here now points
+                   at the model, not the ASR grader.
+    """
     final_text = (result["transcription_text"] or "").strip()
     assert final_text, "Expected non-empty transcription (model text stream)"
 
     wav_out = _wav_bytes_from_pcm16(result["output_pcm"], result["output_sample_rate"])
-    whisper_text = convert_audio_bytes_to_text(wav_out).strip()
+    whisper_text = convert_audio_bytes_to_text(wav_out, model_size=whisper_model_size).strip()
     assert whisper_text, "Whisper returned empty string for synthesized output audio"
 
     sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
diff --git a/tests/entrypoints/openai_api/test_serving_chat_speaker.py b/tests/entrypoints/openai_api/test_serving_chat_speaker.py
index 27839ee23..b0977c0c2 100644
--- a/tests/entrypoints/openai_api/test_serving_chat_speaker.py
+++ b/tests/entrypoints/openai_api/test_serving_chat_speaker.py
@@ -87,8 +87,8 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M
     serving_chat.chat_template = None
     serving_chat.chat_template_content_format = "string"
     serving_chat.default_chat_template_kwargs = {}
-    serving_chat.openai_serving_render = mocker.MagicMock()
-    serving_chat.openai_serving_render.validate_chat_template.return_value = None
+    serving_chat.online_renderer = mocker.MagicMock()
+    serving_chat.online_renderer.validate_chat_template.return_value = None
     serving_chat._effective_chat_template_kwargs = mocker.MagicMock(return_value={})
     serving_chat._preprocess_chat = mocker.AsyncMock(
         side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian")
diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py
index c3e4bd960..904e5735c 100644
--- a/tests/entrypoints/openai_api/test_serving_speech.py
+++ b/tests/entrypoints/openai_api/test_serving_speech.py
@@ -39,6 +39,8 @@ from vllm_omni.entrypoints.openai.serving_speech import (
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
+from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest, SpeechServingContext
+from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     FISH_TEXT_ONLY_SYSTEM_PROMPT,
     build_fish_voice_clone_prompt_ids,
@@ -972,6 +974,29 @@ class TestTTSMethods:
                 )
             )
 
+    def test_upload_ming_voice_embedding_wrong_dims_rejected_without_replacing_existing(self, speech_server):
+        """Ming embedding uploads must be 192-dim before replacing stored voices."""
+        speech_server._tts_model_type = "ming_tts"
+        existing = {
+            "name": "bad_emb_voice",
+            "file_path": "/tmp/voice_samples/bad_emb_voice.safetensors",
+            "mime_type": "application/x-safetensors",
+            "embedding_source": "direct",
+            "embedding_dim": 192,
+        }
+        speech_server.uploaded_speakers = {"bad_emb_voice": existing.copy()}
+
+        with pytest.raises(ValueError, match="Ming speaker embedding must have 192 dims, got 191"):
+            asyncio.run(
+                speech_server.upload_voice_embedding(
+                    embedding_json=json.dumps([0.0] * 191),
+                    consent="consent",
+                    name="bad_emb_voice",
+                )
+            )
+
+        assert speech_server.uploaded_speakers["bad_emb_voice"] == existing
+
     def test_base_task_requires_ref_audio_or_speaker_embedding(self, speech_server):
         """Base task without ref_audio or speaker_embedding is rejected."""
         req = OpenAICreateSpeechRequest(input="Hello", task_type="Base")
@@ -1556,6 +1581,86 @@ class TestTTSMethods:
         assert params["x_vector_only_mode"] == [True]
         assert "ref_audio" not in params
 
+    def test_ming_adapter_uploaded_direct_embedding_uses_embedding_not_audio(
+        self, speech_server, mocker: MockerFixture
+    ):
+        """Ming uploaded direct embeddings are loaded into prompt construction."""
+        speech_server.uploaded_speakers = {
+            "emb_voice": {
+                "name": "emb_voice",
+                "file_path": "/tmp/voice_samples/emb_voice.safetensors",
+                "mime_type": "application/x-safetensors",
+                "embedding_source": "direct",
+                "embedding_dim": 192,
+            }
+        }
+        fake_embedding = [0.1] * 192
+        mock_get_emb = mocker.patch.object(
+            speech_server, "_get_uploaded_speaker_embedding", return_value=fake_embedding
+        )
+        mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data")
+        mock_prompt = mocker.patch.object(
+            speech_server,
+            "_build_ming_dense_prompt",
+            return_value={"additional_information": {"speaker_count": 1}},
+        )
+
+        adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
+        req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice")
+        prepared = asyncio.run(adapter.build(req, [], False))
+
+        mock_get_emb.assert_called_once_with("emb_voice")
+        mock_get_audio.assert_not_called()
+        mock_prompt.assert_called_once_with(req, ref_audio_data=None)
+        prompt_request = mock_prompt.call_args.args[0]
+        assert prompt_request.speaker_embedding == fake_embedding
+        assert len(prompt_request.speaker_embedding) == 192
+        assert prepared.prompt["additional_information"]["speaker_count"] == 1
+        assert prepared.model_type == "ming_tts"
+
+    def test_ming_adapter_uploaded_audio_path_preserves_ref_text(self, speech_server, mocker: MockerFixture):
+        """Ming uploaded audio voices still resolve audio and stored ref_text."""
+        speech_server.uploaded_speakers = {
+            "audio_voice": {
+                "name": "audio_voice",
+                "file_path": "/tmp/voice_samples/audio_voice.wav",
+                "mime_type": "audio/wav",
+                "ref_text": "Reference transcript.",
+            }
+        }
+        ref_audio_source = "data:audio/wav;base64,ZmFrZWF1ZGlv"
+        ref_audio_data = ([0.0, 0.1], 16000)
+        fake_embedding = [0.2] * 192
+        mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data", return_value=ref_audio_source)
+        mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding")
+        mocker.patch.object(
+            speech_server,
+            "_resolve_ref_audio",
+            new=mocker.AsyncMock(return_value=ref_audio_data),
+        )
+        mock_extract = mocker.patch.object(
+            speech_server,
+            "_extract_ming_speaker_embeddings_from_ref_audio",
+            return_value=[fake_embedding],
+        )
+        mock_prompt = mocker.patch.object(
+            speech_server,
+            "_build_ming_dense_prompt",
+            return_value={"additional_information": {"speaker_count": 1}},
+        )
+
+        adapter = MingTTSAdapter(SpeechServingContext(server=speech_server))
+        req = OpenAICreateSpeechRequest(input="Hello", voice="audio_voice")
+        prepared = asyncio.run(adapter.build(req, [], False))
+
+        mock_get_audio.assert_called_once_with("audio_voice")
+        mock_get_emb.assert_not_called()
+        mock_extract.assert_called_once_with([ref_audio_data])
+        mock_prompt.assert_called_once_with(req, ref_audio_data=ref_audio_data)
+        assert req.ref_text == "Reference transcript."
+        assert req.speaker_embedding == fake_embedding
+        assert prepared.model_type == "ming_tts"
+
     # ── regression: full flow from issue #1603 ──
 
     def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server, mocker: MockerFixture):
@@ -2530,7 +2635,7 @@ def test_api_server_delete_voice_exception_returns_500(mocker: MockerFixture):
 def test_api_server_create_speech_without_handler_returns_404(mocker: MockerFixture):
     fake_base = _patch_api_server_base(mocker)
     raw_request = _make_api_server_request(None, path="/v1/audio/speech")
-    raw_request.app.state.openai_serving_tokenization = fake_base
+    raw_request.app.state.serving_tokenization = fake_base
     request = OpenAICreateSpeechRequest(input="Hello")
 
     response = asyncio.run(api_server_module.create_speech(request, raw_request))
@@ -2543,7 +2648,7 @@ def test_api_server_create_speech_without_handler_returns_404(mocker: MockerFixt
 def test_api_server_create_speech_batch_without_handler_returns_404(mocker: MockerFixture):
     fake_base = _patch_api_server_base(mocker)
     raw_request = _make_api_server_request(None, path="/v1/audio/speech/batch")
-    raw_request.app.state.openai_serving_tokenization = fake_base
+    raw_request.app.state.serving_tokenization = fake_base
     request = BatchSpeechRequest(items=[SpeechBatchItem(input="hi")])
 
     response = asyncio.run(api_server_module.create_speech_batch(request, raw_request))
@@ -2612,7 +2717,7 @@ def test_api_server_create_audio_generate_without_handler_returns_404(mocker: Mo
     fake_base = _patch_api_server_base(mocker)
     raw_request = _make_api_server_request(None, path="/v1/audio/generate")
     raw_request.app.state.openai_serving_audio_generate = None
-    raw_request.app.state.openai_serving_tokenization = fake_base
+    raw_request.app.state.serving_tokenization = fake_base
     request = OpenAICreateAudioGenerateRequest(input="a bird singing")
 
     response = asyncio.run(api_server_module.create_audio_generate(request, raw_request))
@@ -3343,6 +3448,61 @@ class TestTTSAsyncOffloading:
         qwen3_tts_server._build_tts_params.assert_called_once()
         qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
 
+    def test_prepare_speech_generation_qwen3_default_seed_sets_tts_local_seed(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """Deploy default seed should seed Qwen3 TTS residual MTP sampling."""
+        qwen3_tts_server.engine_client.default_sampling_params_list = [
+            SimpleNamespace(max_tokens=2048, seed=42, extra_args=None)
+        ]
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = mocker.MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+        request = OpenAICreateSpeechRequest(input="hello")
+
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        stage0_params = qwen3_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0]
+        assert stage0_params.seed == 42
+        assert stage0_params.extra_args["tts_local_seed"] == 42
+        assert qwen3_tts_server.engine_client.default_sampling_params_list[0].extra_args is None
+
+    def test_prepare_speech_generation_uses_adapter_model_type_label(
+        self,
+        voxtral_server,
+        mocker: MockerFixture,
+    ):
+        """Adapter model_type should replace the legacy _tts_model_type label ladder."""
+        legacy_tts_model_type = "dummy_tts"
+        adapter_model_type = "adapter_dummy_tts"
+
+        class FakeAdapter:
+            def validate(self, request):
+                return None
+
+            async def build(self, request, sampling_params_list, has_inline_ref_audio):
+                return PreparedRequest(
+                    prompt={"prompt": request.input},
+                    tts_params={},
+                    model_type=adapter_model_type,
+                )
+
+        voxtral_server._tts_model_type = legacy_tts_model_type
+        mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter())
+        log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info")
+
+        asyncio.run(voxtral_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="hello")))
+
+        assert adapter_model_type != legacy_tts_model_type
+        assert any(
+            call.args
+            and call.args[0] == "TTS speech request %s: text=%r, model=%s"
+            and call.args[3] == adapter_model_type
+            for call in log_info.call_args_list
+        )
+
     def test_prepare_speech_generation_treats_sse_as_streaming(self, qwen3_tts_server, mocker: MockerFixture):
         """stream_format=sse should request delta-style multimodal outputs."""
         qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
@@ -3394,6 +3554,48 @@ class TestTTSAsyncOffloading:
 
         mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True)
 
+    def test_prepare_speech_generation_no_async_chunk_stream_uses_final_only(
+        self, qwen3_tts_server, mocker: MockerFixture
+    ):
+        """Full-payload TTS streaming should not request delta multimodal outputs."""
+        qwen3_tts_server.engine_client.model_config.async_chunk = False
+        qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._build_tts_params = mocker.MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+        mock_coerce = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
+            return_value=qwen3_tts_server.engine_client.default_sampling_params_list,
+        )
+        request = OpenAICreateSpeechRequest(input="hello", stream_format="audio", response_format="pcm")
+
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, False)
+
+    def test_prepare_speech_generation_no_async_chunk_stream_keeps_delta_for_non_qwen3(
+        self, voxtral_server, mocker: MockerFixture
+    ):
+        """FINAL_ONLY streaming for async_chunk=False is scoped to qwen3_tts only."""
+        voxtral_server.engine_client.model_config.async_chunk = False
+        mocker.patch.object(voxtral_server._get_tts_adapter(), "validate", return_value=None)
+        voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock(
+            return_value={
+                "prompt_token_ids": [1, 2, 3],
+                "additional_information": {"voice": ["test"]},
+            }
+        )
+        mock_coerce = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types",
+            return_value=voxtral_server.engine_client.default_sampling_params_list,
+        )
+        request = OpenAICreateSpeechRequest(input="hello", voice="test", stream_format="audio", response_format="pcm")
+
+        asyncio.run(voxtral_server._prepare_speech_generation(request))
+
+        mock_coerce.assert_called_once_with(voxtral_server.engine_client.default_sampling_params_list, True)
+
     def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false(
         self, qwen3_tts_server, mocker: MockerFixture
     ):
diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py
index 3f1c5f8e5..ff69a2bf2 100644
--- a/tests/entrypoints/openai_api/test_video_server.py
+++ b/tests/entrypoints/openai_api/test_video_server.py
@@ -44,23 +44,18 @@ class MockVideoResult:
         videos,
         audios=None,
         sample_rate=None,
-        custom_output=None,
+        multimodal_output=None,
         stage_durations=None,
         peak_memory_mb=0.0,
     ):
-        self.multimodal_output = {"video": videos}
+        self.multimodal_output = dict(multimodal_output or {"video": videos})
         if audios is not None:
             self.multimodal_output["audio"] = audios
         if sample_rate is not None:
             self.multimodal_output["audio_sample_rate"] = sample_rate
-        self._custom_output = custom_output or {}
         self.stage_durations = stage_durations or {}
         self.peak_memory_mb = peak_memory_mb
 
-    @property
-    def custom_output(self):
-        return self._custom_output
-
 
 class FakeAsyncOmni:
     def __init__(self):
@@ -835,7 +830,13 @@ def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker:
         engine.captured_sampling_params_list = sampling_params_list
         import numpy as np
 
-        yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], custom_output={"video_fps_multiplier": 2})
+        yield MockVideoResult(
+            [np.zeros((1, 64, 64, 3), dtype=np.uint8)],
+            multimodal_output={
+                "video": [np.zeros((1, 64, 64, 3), dtype=np.uint8)],
+                "metadata": {"video": {"video_fps_multiplier": 2}},
+            },
+        )
 
     engine.generate = _generate
 
@@ -939,11 +940,16 @@ def test_video_generation_response_exposes_action_payload(mocker: MockerFixture)
 
         yield MockVideoResult(
             [object()],
-            custom_output={
-                "action": np.array([[[1.5, 2.5], [3.5, 4.5]]], dtype=np.float32),
-                "raw_action_dim": 2,
-                "action_mode": "policy",
-                "domain_id": 7,
+            multimodal_output={
+                "video": [object()],
+                "actions": np.array([[[1.5, 2.5], [3.5, 4.5]]], dtype=np.float32),
+                "metadata": {
+                    "actions": {
+                        "raw_action_dim": 2,
+                        "action_mode": "policy",
+                        "domain_id": 7,
+                    },
+                },
             },
         )
 
@@ -981,11 +987,16 @@ def test_video_job_persists_action_metadata(test_client, mocker: MockerFixture):
         engine.captured_sampling_params_list = sampling_params_list
         yield MockVideoResult(
             [object()],
-            custom_output={
-                "action": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32),
-                "raw_action_dim": 2,
-                "action_mode": "policy",
-                "domain_id": 7,
+            multimodal_output={
+                "video": [object()],
+                "actions": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32),
+                "metadata": {
+                    "actions": {
+                        "raw_action_dim": 2,
+                        "action_mode": "policy",
+                        "domain_id": 7,
+                    },
+                },
             },
         )
 
@@ -1019,19 +1030,51 @@ def test_action_extraction_accepts_unbatched_action():
 
     result = MockVideoResult(
         [object()],
-        custom_output={
-            "action": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
-            "raw_action_dim": 2,
-            "action_mode": "policy",
-            "domain_id": 7,
+        multimodal_output={
+            "video": [object()],
+            "actions": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": 2,
+                    "action_mode": "policy",
+                    "domain_id": 7,
+                },
+            },
         },
     )
 
     actions = OmniOpenAIServingVideo._extract_action_outputs(result, expected_count=1)
 
+    assert actions and actions[0] is not None
+    assert actions[0].data == [[1.0, 2.0], [3.0, 4.0]]
+    assert actions[0].shape == [2, 2]
+
+
+def test_action_extraction_accepts_multimodal_actions_payload():
+    import numpy as np
+
+    result = MockVideoResult([object()])
+    result.multimodal_output.update(
+        {
+            "actions": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32),
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": 2,
+                    "action_mode": "policy",
+                    "domain_id": 7,
+                },
+            },
+        }
+    )
+
+    actions = OmniOpenAIServingVideo._extract_action_outputs(result, expected_count=1)
+
     assert actions[0] is not None
     assert actions[0].data == [[1.0, 2.0], [3.0, 4.0]]
     assert actions[0].shape == [2, 2]
+    assert actions[0].raw_action_dim == 2
+    assert actions[0].action_mode == "policy"
+    assert actions[0].domain_id == 7
 
 
 def test_missing_handler_returns_503():
@@ -1840,7 +1883,13 @@ def test_worker_fps_multiplier_is_applied_to_sync_encoding(test_client, mocker:
     async def _generate(prompt, request_id, sampling_params_list):
         engine.captured_prompt = prompt
         engine.captured_sampling_params_list = sampling_params_list
-        yield MockVideoResult([object()], custom_output={"video_fps_multiplier": 2})
+        yield MockVideoResult(
+            [object()],
+            multimodal_output={
+                "video": [object()],
+                "metadata": {"video": {"video_fps_multiplier": 2}},
+            },
+        )
 
     engine.generate = _generate
 
diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py
index 61a420a4c..ba132dbd2 100644
--- a/tests/entrypoints/test_async_omni_diffusion_config.py
+++ b/tests/entrypoints/test_async_omni_diffusion_config.py
@@ -343,7 +343,7 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
     )
     mocker.patch(
         "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
-        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage]),
+        return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None),
     )
 
     engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
@@ -359,3 +359,39 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo
 
     assert not hasattr(stage_configs[0].engine_args, "additional_config")
     assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}}
+
+
+def test_default_stage_config_includes_quantization_config():
+    """Ensure structured quantization_config survives default diffusion-stage creation."""
+    quantization_config = {
+        "method": "example_quant",
+        "weights": "weights.bin",
+    }
+
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0]
+
+    assert stage_cfg["engine_args"]["quantization_config"] == quantization_config
+
+
+def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(mocker):
+    fake_diffusion_stage = SimpleNamespace(
+        stage_type="diffusion",
+        engine_args=SimpleNamespace(quantization_config=None),
+    )
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
+        return_value=("dummy.yaml", [fake_diffusion_stage], None),
+    )
+
+    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+    engine._strip_single_engine_args = lambda kwargs: kwargs
+
+    _, stage_configs = engine._resolve_stage_configs(
+        "dummy-model",
+        {
+            "stage_configs_path": "dummy.yaml",
+            "quantization_config": {"method": "bitsandbytes"},
+        },
+    )
+
+    assert stage_configs[0].engine_args.quantization_config == {"method": "bitsandbytes"}
diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py
index 8f647138b..4da82390b 100644
--- a/tests/entrypoints/test_omni_entrypoints.py
+++ b/tests/entrypoints/test_omni_entrypoints.py
@@ -107,6 +107,7 @@ class FakeAsyncOmniEngine:
         self.stage_vllm_configs = [None for _ in range(self.num_stages)]
         self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)]
         self.input_processor = None
+        self.endpoint_restrictions = ()
 
         self.output_q: queue.Queue[Any] = queue.Queue()
         self.submitted: list[dict[str, Any]] = []
diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py
index c2a5ff0a1..efbb19fff 100644
--- a/tests/entrypoints/test_omni_sleep_mode.py
+++ b/tests/entrypoints/test_omni_sleep_mode.py
@@ -179,7 +179,7 @@ class TestOmniSleepMode:
         try:
             device_id = 0
             used_before = get_device_global_memory_used_gib(device_id)
-            acks = await llm_engine.sleep(stage_ids=[0], level=2)
+            acks = await llm_engine.sleep(stage_ids=[0], level=1)
             await asyncio.sleep(1.5)
             used_after = get_device_global_memory_used_gib(device_id)
             drop_gib = used_before - used_after
@@ -209,7 +209,7 @@ class TestOmniSleepMode:
         """Diffusion Worker stage signal loop"""
         try:
             logger.info("Starting Diffusion Worker Handshake Test")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
 
             def _get_status(ack):
                 return ack.status if hasattr(ack, "status") else ack.get("status")
@@ -264,9 +264,9 @@ class TestOmniSleepMode:
                 base_output = output
             assert base_output is not None and len(base_output.images) > 0
             logger.info("Baseline Generation successful.")
-            # Sleep Level 2
+            # Sleep Level 1
             logger.info("Entering Deep Sleep (VRAM Scavenging)...")
-            await diffusion_engine.sleep(stage_ids=[0], level=2)
+            await diffusion_engine.sleep(stage_ids=[0], level=1)
             # Wake-up
             logger.info("Waking up (Reloading Weights)...")
             await diffusion_engine.wake_up(stage_ids=[0])
@@ -353,8 +353,8 @@ class TestOmniSleepMode:
             logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB")
 
             # Sleep
-            logger.info("Triggering Level 2 Deep Sleep (Full Weight Offloading)...")
-            acks = await diffusion_engine.sleep(stage_ids=[0], level=2)
+            logger.info("Triggering Level 1 Deep Sleep (Partial Weight Offloading)...")
+            acks = await diffusion_engine.sleep(stage_ids=[0], level=1)
 
             reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks)
             reported_freed_gib = reported_freed_bytes / 1024**3
@@ -407,6 +407,47 @@ class TestOmniSleepMode:
             diffusion_engine.shutdown()
             await asyncio.sleep(1)
 
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_level2_sleep_wake_raises(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro A: wake_up() after sleep(level=2) must raise
+        NotImplementedError instead of silently producing corrupted output."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=2)
+            with pytest.raises(NotImplementedError, match="sleep\\(level=2\\)"):
+                await llm_engine.wake_up(stage_ids=[0])
+        finally:
+            llm_engine.shutdown()
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro B: generate() must be rejected if kv_cache
+        is still asleep after wake_up(tags=["weights"]), instead of crashing with
+        CUDA illegal memory access."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=1)
+            await llm_engine.wake_up(stage_ids=[0], tags=["weights"])
+            with pytest.raises(RuntimeError, match="partially or fully asleep"):
+                async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)):
+                    pass
+        finally:
+            llm_engine.shutdown()
+
+    @pytest.mark.asyncio
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni):
+        """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a
+        safe no-op instead of raising a cumem CUDA invalid argument error."""
+        try:
+            await llm_engine.sleep(stage_ids=[0], level=1)
+            first_acks = await llm_engine.wake_up(stage_ids=[0])
+            assert len(first_acks) > 0, "First wake_up() should return ACKs"
+            second_acks = await llm_engine.wake_up(stage_ids=[0])
+            assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}"
+        finally:
+            llm_engine.shutdown()
+
 
 # ---------------------------------------------------------------------------
 # H100 (or MI325) e2e: BAGEL / pure diffusion, multi-TP
@@ -444,7 +485,7 @@ async def test_diffusion_model_sleep_tp(tp_size: int):
         async for _ in engine.generate("test", sampling_params_list=[llm_sp, diff_sp]):
             pass
 
-        acks = await engine.sleep(level=2)
+        acks = await engine.sleep(level=1)
         statuses = [get_ack_info(ack, "status") for ack in acks]
         assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}"
 
@@ -500,7 +541,7 @@ async def test_multistage_sleep_h100(tp_size: int):
         async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]):
             pass
 
-        acks = await engine.sleep(stage_ids=[0, 1], level=2)
+        acks = await engine.sleep(stage_ids=[0, 1], level=1)
         assert len(acks) == 2
 
         await engine.wake_up(stage_ids=[0, 1])
diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py
index 374dcf7fa..4dcaddd63 100644
--- a/tests/entrypoints/test_resolve_dreamzero_config.py
+++ b/tests/entrypoints/test_resolve_dreamzero_config.py
@@ -24,15 +24,15 @@ def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch):
 
 def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
     monkeypatch.setattr(
-        "vllm_omni.config.config_factory.StageConfigFactory._auto_detect_model_type",
-        classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)),
+        "vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type",
+        classmethod(lambda _cls, model, trust_remote_code=True: "vla"),
     )
     monkeypatch.setattr(
         "vllm_omni.diffusion.utils.hf_utils._looks_like_dreamzero",
         lambda _model: True,
     )
 
-    stage_configs = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
+    stage_configs, _ = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID")
     engine_args = stage_configs[0].engine_args
 
     assert engine_args.model_class_name == "DreamZeroPipeline"
diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py
index 6c4445306..61ffed3d5 100644
--- a/tests/entrypoints/test_serve.py
+++ b/tests/entrypoints/test_serve.py
@@ -9,6 +9,7 @@ import pytest
 from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless
+from vllm_omni.entrypoints.utils import parse_stage_overrides
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -31,6 +32,20 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None:
     assert not explicit["async_chunk"]
 
 
+def test_serve_parser_accepts_strategy_config() -> None:
+    """``--strategy-config`` must parse onto the ``strategy_config`` dest and be
+    forwarded as an explicit kwarg so the engine can overlay the strategy."""
+    parser = TrackingArgumentParser()
+    subparsers = parser.add_subparsers(dest="subcommand")
+    cmd = OmniServeCommand()
+    cmd.subparser_init(subparsers)
+
+    argv = ["serve", "fake-model", "--omni", "--strategy-config", "/tmp/strategy.yaml"]
+    args = parser.parse_args(argv)
+    assert args.strategy_config == "/tmp/strategy.yaml"
+    assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml"
+
+
 def _make_headless_args(**kwargs) -> TrackingNamespace:
     defaults = {
         "model": "fake-model",
@@ -80,13 +95,97 @@ def test_run_headless_rejects_non_multiprocess_worker_backend() -> None:
         run_headless(args)
 
 
+# ---------------------------------------------------------------------------
+# --stage-overrides parsing parity (headless vs standard path)
+# ---------------------------------------------------------------------------
+
+
+def test_parse_stage_overrides_valid_json() -> None:
+    """A valid JSON string is parsed into the nested per-stage dict."""
+    parsed = parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}')
+    assert parsed == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
+
+
+def test_parse_stage_overrides_none_and_empty_return_none() -> None:
+    """No overrides (None / empty string) resolve to ``None``."""
+    assert parse_stage_overrides(None) is None
+    assert parse_stage_overrides("") is None
+
+
+def test_parse_stage_overrides_empty_dict_returns_none() -> None:
+    """An empty dict is falsy and must resolve to ``None``, locking in parity
+    with the original standard-path ``if stage_overrides_json:`` falsy check."""
+    assert parse_stage_overrides({}) is None
+
+
+def test_parse_stage_overrides_passes_through_non_str() -> None:
+    """An already-parsed mapping is returned unchanged (identity)."""
+    overrides = {"0": {"devices": "0"}}
+    assert parse_stage_overrides(overrides) is overrides
+
+
+def test_parse_stage_overrides_invalid_json_raises() -> None:
+    """Invalid JSON raises ValueError whose message matches the standard path
+    verbatim: the ``--stage-overrides is not valid JSON:`` prefix AND the
+    ``Got: <repr>`` suffix echoing the raw input."""
+    bad = "{not valid json}"
+    with pytest.raises(ValueError) as excinfo:
+        parse_stage_overrides(bad)
+    message = str(excinfo.value)
+    assert message.startswith("--stage-overrides is not valid JSON:")
+    assert f"Got: {bad!r}" in message
+
+
+def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) -> None:
+    """Regression: the headless path must parse ``--stage-overrides`` (a JSON
+    string) and forward the parsed dict to ``load_and_resolve_stage_configs``,
+    mirroring the standard engine path. Previously it was dropped entirely,
+    silently producing a different per-stage device layout."""
+    captured: dict = {}
+
+    def _fake_resolve(*args, **kwargs):
+        captured.update(kwargs)
+        # Return a stage that does NOT match stage_id=0 so run_headless stops
+        # right after the resolver call (we only care about how it was called).
+        return ("/fake/stages.yaml", [SimpleNamespace(stage_id=99)], None)
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        side_effect=_fake_resolve,
+    )
+
+    args = _make_headless_args(
+        stage_id=0,
+        strategy_config="/tmp/strategy.yaml",
+        stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}',
+    )
+    with pytest.raises(ValueError, match="No stage config found for stage_id=0"):
+        run_headless(args)
+
+    assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}}
+    assert captured["strategy_config_path"] == "/tmp/strategy.yaml"
+
+
+def test_run_headless_invalid_stage_overrides_raises(mocker: MockerFixture) -> None:
+    """Invalid ``--stage-overrides`` JSON in headless mode fails fast with the
+    shared ValueError instead of being silently ignored."""
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        return_value=("/fake/stages.yaml", [SimpleNamespace(stage_id=0)], None),
+    )
+
+    args = _make_headless_args(stage_id=0, stage_overrides="{not valid json}")
+    with pytest.raises(ValueError, match="--stage-overrides is not valid JSON"):
+        run_headless(args)
+
+
 def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None:
     """Headless looks up its assigned stage_id in the loaded deploy YAML and
     fails fast when the launcher's --stage-id doesn't match any entry."""
     other_stage = SimpleNamespace(stage_id=99)
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [other_stage]),
+        return_value=("/fake/stages.yaml", [other_stage], None),
     )
 
     args = _make_headless_args(stage_id=0)
@@ -134,7 +233,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg]),
+        return_value=("/fake/stages.yaml", [stage_cfg], None),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -210,7 +309,7 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg]),
+        return_value=("/fake/stages.yaml", [stage_cfg], None),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -269,7 +368,7 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture)
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg]),
+        return_value=("/fake/stages.yaml", [stage_cfg], None),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
@@ -349,7 +448,7 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur
 
     mocker.patch(
         "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
-        return_value=("/fake/stages.yaml", [stage_cfg]),
+        return_value=("/fake/stages.yaml", [stage_cfg], None),
     )
     mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
     mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None)
diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py
index 52cc32215..ebbca1917 100644
--- a/tests/entrypoints/test_stream_finish_reason.py
+++ b/tests/entrypoints/test_stream_finish_reason.py
@@ -144,7 +144,7 @@ def _build_serving_chat():
         engine_client=mock_engine,
         models=models,
         response_role="assistant",
-        openai_serving_render=mock_render,
+        online_renderer=mock_render,
         request_logger=None,
         chat_template=None,
         chat_template_content_format="auto",
diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py
index 479ada963..b48ed2d91 100644
--- a/tests/entrypoints/test_utils.py
+++ b/tests/entrypoints/test_utils.py
@@ -318,7 +318,7 @@ class TestLoadAndResolveStageConfigs:
     def test_load_and_resolve_with_kwargs(self):
         """Ensure that dtype survives default stage creation."""
         kwargs = {"dtype": torch.float32}
-        config_path, stage_configs = load_and_resolve_stage_configs(
+        config_path, stage_configs, _ = load_and_resolve_stage_configs(
             model="black-forest-labs/FLUX.2-klein-4B",
             stage_configs_path=None,
             kwargs=kwargs,
@@ -349,10 +349,10 @@ class TestLoadAndResolveStageConfigs:
         ]
         load_stage_configs = mocker.patch(
             "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-            return_value=returned_stage_configs,
+            return_value=(returned_stage_configs, None),
         )
 
-        config_path, stage_configs = load_and_resolve_stage_configs(
+        config_path, stage_configs, _ = load_and_resolve_stage_configs(
             model="dummy-model",
             stage_configs_path=str(deploy_path),
             kwargs={},
@@ -363,6 +363,7 @@ class TestLoadAndResolveStageConfigs:
             base_engine_args={},
             deploy_config_path=str(deploy_path),
             stage_overrides=None,
+            strategy_config_path=None,
         )
         assert config_path == str(deploy_path)
         assert len(stage_configs) == 2
diff --git a/tests/examples/offline_inference/test_image_to_video.py b/tests/examples/offline_inference/test_image_to_video.py
new file mode 100644
index 000000000..b53648cb3
--- /dev/null
+++ b/tests/examples/offline_inference/test_image_to_video.py
@@ -0,0 +1,81 @@
+"""
+Offline inference tests: image-to-video.
+See examples/offline_inference/image_to_video/README.md
+"""
+
+import shutil
+from pathlib import Path
+
+import pytest
+import requests
+
+from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
+from tests.helpers.assertions import assert_video_valid
+from tests.helpers.mark import hardware_marks
+
+pytestmark = [
+    pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
+    pytest.mark.full_model,
+    pytest.mark.example,
+    *hardware_marks(res={"cuda": "H100"}),
+]
+
+I2V_SCRIPT = EXAMPLES / "offline_inference" / "image_to_video" / "image_to_video.py"
+README_PATH = I2V_SCRIPT.with_name("README.md")
+EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_i2v"
+
+_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
+_IMAGE_NAME = "cherry_blossom.jpg"
+
+_SKIP_SECTIONS = {
+    "Prerequisites",
+    "Advanced Features",
+    "FAQ",
+    # VACE conditional tasks are covered by test_vace_video_generation.py, which
+    # runs the same README snippets with synthetic assets and smoke settings. The
+    # prep block here downloads inputs via wget (not portable, not collected), so
+    # exercising this section from the shared runner would only duplicate and fail.
+    "Wan2.1 VACE Conditional Tasks",
+}
+
+
+def _ensure_test_image(run_dir: Path) -> None:
+    """Download or copy the example image into run_dir so CLI snippets can find it."""
+    dest = run_dir / _IMAGE_NAME
+    if dest.exists():
+        return
+    src = I2V_SCRIPT.parent / _IMAGE_NAME
+    if src.exists():
+        shutil.copy2(src, dest)
+        return
+    response = requests.get(_IMAGE_URL, timeout=60)
+    response.raise_for_status()
+    dest.write_bytes(response.content)
+
+
+def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, str]:
+    if h2_title in _SKIP_SECTIONS:
+        return True, f"README section '{h2_title}' is intentionally excluded for examples tests"
+    if language == "python":
+        return True, "Python API snippets produce video files that ExampleRunner does not auto-collect"
+    if "/path/to/" in code:
+        return True, "Snippet references a placeholder local model path"
+    return False, ""
+
+
+README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet)
+
+
+@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
+def test_image_to_video(snippet: ReadmeSnippet, example_runner: ExampleRunner):
+    should_skip, reason = snippet.skip
+    if should_skip:
+        pytest.skip(reason)
+
+    run_dir = example_runner.output_root / EXAMPLE_OUTPUT_SUBFOLDER / snippet.test_id
+    run_dir.mkdir(parents=True, exist_ok=True)
+    _ensure_test_image(run_dir)
+
+    result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
+    for asset in result.assets:
+        assert_video_valid(asset)
diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py
index b3abea735..11e86ed8c 100644
--- a/tests/helpers/assertions.py
+++ b/tests/helpers/assertions.py
@@ -19,19 +19,20 @@ from PIL import Image
 from tests.helpers.media import (
     convert_audio_bytes_to_text,
     cosine_similarity_text,
+    preprocess_text,
 )
 
 _GENDER_PIPELINE = None
 _GENDER_PIPELINE_LOCK = threading.Lock()
 # Transcript gates default to whisper ``small`` for speed. ``small`` mishears a
 # short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a
-# leading SFX token), which flakes the deterministic similarity gate. A test can
-# opt in to ASR escalation by setting ``transcript_escalation_model`` to a
-# whisper model name (e.g. ``"large-v3"``) in its request_config: on a failed
-# fast pass the clip is re-transcribed with that stronger ASR before the test
-# fails, so a weak-ASR mishear is rescued while a genuine model artifact still
-# fails (the strong ASR mismatches too). Tests that do not opt in keep the
-# strict behaviour.
+# leading SFX token), which flakes the deterministic similarity gate. Short
+# clips first get a conservative containment fallback for minor ASR repeats/noise.
+# A test can also opt in to ASR escalation by setting
+# ``transcript_escalation_model`` to a whisper model name (e.g. ``"large-v3"``)
+# in its request_config: on a failed fast pass the clip is re-transcribed with
+# that stronger ASR before the test fails, so a weak-ASR mishear is rescued while
+# a genuine model artifact still fails (the strong ASR mismatches too).
 _PCM_SPEECH_SAMPLE_RATE_HZ = 24_000
 _MIN_PCM_SPEECH_HNR_DB = 1.0
 _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
@@ -43,6 +44,23 @@ _PRESET_VOICE_GENDER_MAP: dict[str, str] = {
 }
 
 
+def _short_transcript_contains_expected(transcript: str, expected: str) -> bool:
+    """Allow minor ASR repeats/noise for very short speech clips."""
+    transcript_clean = preprocess_text(transcript)
+    expected_clean = preprocess_text(expected)
+    if not transcript_clean or not expected_clean:
+        return False
+
+    transcript_words = transcript_clean.split()
+    expected_words = expected_clean.split()
+    if not transcript_words or not expected_words:
+        return False
+
+    short_text = min(len(transcript_clean), len(expected_clean)) <= 15
+    small_word_delta = len(transcript_words) <= len(expected_words) + 2
+    return short_text and small_word_delta and expected_clean in transcript_clean
+
+
 def assert_image_diffusion_response(
     response: "DiffusionResponse",
     request_config: dict[str, Any],
@@ -620,13 +638,16 @@ def _assert_transcript_matches(
     ``transcript`` is the fast whisper-``small`` result. If it clears
     ``threshold`` the check passes immediately.
 
+    If the cosine check fails, very short clips get a conservative containment
+    fallback that accepts minor ASR repeats/noise only when the expected text is
+    still present and the transcript has few extra words.
+
     When ``escalation_model`` is set (opt-in via the ``transcript_escalation_model``
-    request_config key) and the fast check fails, the clip is re-transcribed with
-    that stronger ASR and the assertion is decided on its verdict -- so a weak
-    whisper-``small`` mishear on a short clip does not flake the gate, while a
-    genuine model artifact still fails (the strong ASR mismatches too). When
-    ``escalation_model`` is ``None`` the original strict behaviour is preserved,
-    so other tests are unaffected.
+    request_config key) and the fast check plus containment fallback fail, the
+    clip is re-transcribed with that stronger ASR and the assertion is decided on
+    its verdict -- so a weak whisper-``small`` mishear on a short clip does not
+    flake the gate, while a genuine model artifact still fails (the strong ASR
+    mismatches too).
     """
     expected = str(expected_text).strip().lower()
     similarity = cosine_similarity_text(transcript.strip().lower(), expected)
@@ -634,6 +655,10 @@ def _assert_transcript_matches(
     if similarity > threshold:
         return
 
+    if _short_transcript_contains_expected(transcript, expected):
+        print("short speech containment check passed")
+        return
+
     if escalation_model and audio_bytes:
         print(
             f"whisper-small below threshold ({similarity:.2f} <= {threshold}); "
diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py
index 51ee981cb..885a4654f 100644
--- a/tests/helpers/runtime.py
+++ b/tests/helpers/runtime.py
@@ -47,6 +47,7 @@ from tests.helpers.media import (
     _merge_base64_audio_to_segment,
     decode_b64_image,
 )
+from tests.model_tests.diffusion.utils import resolve_tiny_model_path
 from vllm_omni.config.stage_config import resolve_deploy_yaml
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
 from vllm_omni.outputs import OmniRequestOutput
@@ -1064,6 +1065,29 @@ class OpenAIClientHandler:
         )
         return [resp]
 
+    def send_completions_http_request(
+        self,
+        request_config: dict[str, Any],
+        *,
+        err_code: int | tuple[int, ...] | list[int] | None = None,
+        err_message: str | tuple[str, ...] | list[str] | None = None,
+    ) -> list[HttpResponse]:
+        """POST ``/v1/completions`` with ``json`` or ``raw_body``."""
+        # TODO (Alex): A lot of these helpers should be consolidated as they differ only by endpoint
+        cfg = _merge_http_expectation_kwargs(
+            request_config,
+            err_code=err_code,
+            err_message=err_message,
+        )
+        r = self._post_json_endpoint("/v1/completions", cfg, default_timeout=120.0)
+        resp = self._http_response_from_requests(r)
+        assert_http_error(
+            resp,
+            err_code=cfg.get("err_code"),
+            err_message=cfg.get("err_message"),
+        )
+        return [resp]
+
     def send_omni_sleep_http_request(
         self,
         request_config: dict[str, Any],
@@ -2792,11 +2816,24 @@ def iter_omni_server(
 
     with omni_fixture_lock:
         params: OmniServerParams = request.param
-        model = model_prefix + params.model
+        # For now, when a tiny model is substituted, we preserve the original model
+        # name via --served-model-name (so that the server still accepts requests with
+        # the original name). We also do the same for server.model so that tests reading
+        # server.model send the correct name in requests.
+        #
+        # TODO: core models on this path currently do not clean up tiny models, although
+        # tiny model paths are deterministic, so it's not a huge footprint. Still, it would
+        # be ideal to cleanup consistently everywhere.
+        original_model = model_prefix + params.model
+        model = original_model
+        if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
+            model = resolve_tiny_model_path(model)
         port = params.port
         stage_config_path = stage_config_path_for_run_level(params.stage_config_path, run_level)
 
         server_args = params.server_args or []
+        if model != original_model:
+            server_args = [*server_args, "--served-model-name", original_model]
         if params.use_omni and params.stage_init_timeout is not None:
             server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
         else:
@@ -2822,6 +2859,8 @@ def iter_omni_server(
                 port=port,
                 env_dict=params.env_dict,
             ) as server:
+                if model != original_model:
+                    server.model = original_model
                 print("OmniServer started successfully")
                 yield server
                 print("OmniServer stopping...")
@@ -2845,6 +2884,8 @@ def iter_omni_server(
                     use_omni=params.use_omni,
                 )
             ) as server:
+                if model != original_model:
+                    server.model = original_model
                 print("OmniServer started successfully")
                 yield server
                 print("OmniServer stopping...")
@@ -2876,6 +2917,8 @@ def iter_omni_runner(
             extra_omni_kwargs = dict(extra) if extra is not None else {}
         stage_config_path = stage_config_path_for_run_level(stage_config_path, run_level)
         model = model_prefix + model
+        if run_level == "core_model" and request.node.get_closest_marker("diffusion"):
+            model = resolve_tiny_model_path(model)
         with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner:
             print("OmniRunner started successfully")
             yield runner
diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py
index 4351eb27d..8706bb1bc 100644
--- a/tests/helpers/stage_config.py
+++ b/tests/helpers/stage_config.py
@@ -633,6 +633,18 @@ def get_deploy_config_path(rel_path: str) -> str:
     return str(_DEPLOY_DIR / rel_path)
 
 
+def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]:
+    """Return one stage entry from a deploy yaml by ``stage_id``."""
+    with open(get_deploy_config_path(rel_path), encoding="utf-8") as f:
+        cfg = yaml.safe_load(f) or {}
+
+    stage_key = "stages" if "stages" in cfg else "stage_args"
+    for stage in cfg.get(stage_key, []):
+        if stage.get("stage_id") == stage_id:
+            return stage
+    raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}")
+
+
 def _get_config_value_by_path(config_dict: dict, path: str) -> Any:
     """Read a dot-separated path from a nested dict (e.g. ``engine_args.load_format``)."""
     current: Any = config_dict
@@ -707,6 +719,7 @@ def stage_config_path_for_run_level(stage_config_path: str | None, run_level: st
 
 __all__ = [
     "get_deploy_config_path",
+    "get_deploy_config_stage",
     "modify_stage_config",
     "stage_config_path_for_run_level",
 ]
diff --git a/tests/helpers/test_assertions.py b/tests/helpers/test_assertions.py
new file mode 100644
index 000000000..0bfa7f593
--- /dev/null
+++ b/tests/helpers/test_assertions.py
@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from tests.helpers.assertions import _assert_transcript_matches
+
+
+def test_short_transcript_repeat_passes_containment_fallback():
+    _assert_transcript_matches(
+        " How... how are you?",
+        audio_bytes=None,
+        expected_text="how are you",
+        threshold=0.9,
+    )
+
+
+def test_short_transcript_unrelated_text_still_fails():
+    with pytest.raises(AssertionError, match="Transcript doesn't match input"):
+        _assert_transcript_matches(
+            " I don't know, sorry.",
+            audio_bytes=None,
+            expected_text="how are you",
+            threshold=0.9,
+        )
diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py
index 87281af82..5e31c2cbd 100644
--- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py
+++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py
@@ -18,8 +18,11 @@ import types
 
 import pytest
 import torch
+import torch.nn.functional as F
 from pytest_mock import MockerFixture
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 # Direct file import to avoid vllm_omni.__init__ patch dependencies.
 _MODELS = os.path.join(
     os.path.dirname(__file__),
@@ -46,6 +49,11 @@ def _load_module(name: str, filename: str):
 
 def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
     """Build the dict of modules to inject into sys.modules."""
+
+    class NativeCustomOp(torch.nn.Module):
+        def forward(self, *args, **kwargs):
+            return self.forward_native(*args, **kwargs)
+
     platforms_mock = mocker.MagicMock()
     platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False
     platforms_mock.current_omni_platform.is_npu.return_value = False
@@ -73,6 +81,8 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
 
     vllm_parallel_mock = mocker.MagicMock()
     vllm_parallel_mock.VocabParallelEmbedding = torch.nn.Embedding
+    custom_op_mock = types.ModuleType("vllm_omni.diffusion.layers.custom_op")
+    custom_op_mock.CustomOp = NativeCustomOp
 
     return {
         "vllm_omni": mocker.MagicMock(),
@@ -82,6 +92,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]:
         "vllm.config.vllm": vllm_config_mod,
         "vllm.model_executor.model_loader.weight_utils": weight_utils_mock,
         "vllm.model_executor.layers.vocab_parallel_embedding": vllm_parallel_mock,
+        "vllm_omni.diffusion.layers.custom_op": custom_op_mock,
         "vllm_omni.model_executor": types.ModuleType("vllm_omni.model_executor"),
         "vllm_omni.model_executor.models": models_pkg,
         "vllm_omni.model_executor.models.common": common_pkg,
@@ -165,6 +176,38 @@ def _make_vllm_config(mocker: MockerFixture, max_num_seqs: int = 4):
     return vllm_config
 
 
+def test_npu_custom_ops_use_fused_norm_and_cached_rope(mocker: MockerFixture, loaded_target_classes) -> None:
+    common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    cp_config, _ = _make_tiny_config(loaded_target_classes)
+    rms_calls = []
+
+    def npu_rms_norm(hidden_states, weight, epsilon):
+        rms_calls.append((hidden_states, weight, epsilon))
+        return hidden_states + 1, None
+
+    mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=True)
+    mocker.patch.object(
+        common_mod,
+        "torch_npu",
+        types.SimpleNamespace(npu_rms_norm=npu_rms_norm),
+        create=True,
+    )
+
+    hidden_states = torch.zeros(1, 2, cp_config.hidden_size, dtype=torch.float16)
+    norm = common_mod._RMSNorm(cp_config.hidden_size, eps=cp_config.rms_norm_eps)
+    torch.testing.assert_close(norm.forward_npu(hidden_states), hidden_states + 1)
+    assert len(rms_calls) == 1
+    assert rms_calls[0][2] == cp_config.rms_norm_eps
+
+    rotary = common_mod._RotaryEmbedding(cp_config)
+    assert rotary.cos_cached.shape == (cp_config.num_code_groups + 1, cp_config.head_dim)
+    assert rotary.sin_cached.shape == rotary.cos_cached.shape
+    position_ids = torch.tensor([[0, 2, 4]])
+    cos, sin = rotary.forward_npu(hidden_states, position_ids)
+    torch.testing.assert_close(cos, rotary.cos_cached[position_ids].to(torch.float16))
+    torch.testing.assert_close(sin, rotary.sin_cached[position_ids].to(torch.float16))
+
+
 class TestCodePredictorDtypeAlignment:
     """Test that code predictor buffers match model parameter dtype."""
 
@@ -339,6 +382,91 @@ class TestCodePredictorDtypeAlignment:
         assert not torch.equal(first[:, 1:], different[:, 1:])
 
 
+class TestCodePredictorPerRowGenerators:
+    """Seeded requests must stay deterministic inside a multi-row batch (#4883)."""
+
+    def _make_predictor(self, mocker: MockerFixture, loaded_target_classes):
+        _, _, code_predictor_wrapper, _, _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+        cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
+        predictor = code_predictor_wrapper(
+            vllm_config=vllm_config,
+            config=cp_config,
+            talker_config=talker_config,
+        )
+        predictor._wrapper_config.use_cuda_graphs = False
+        return predictor, talker_config
+
+    def test_multinomial_per_row_matches_single_row_draws(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        multinomial = common_mod.CodePredictorWrapper._multinomial
+        torch.manual_seed(7)
+        probs = torch.softmax(torch.randn(3, 64), dim=-1)
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=probs.device)
+            generator.manual_seed(seed)
+            return generator
+
+        batched = multinomial(probs, None, [seeded(11), None, seeded(22)])
+        assert batched.shape == (3, 1)
+        # Each seeded row must reproduce a standalone draw from the same seed,
+        # independent of what the other rows in the batch consume.
+        assert torch.equal(batched[0:1], torch.multinomial(probs[0:1], num_samples=1, generator=seeded(11)))
+        assert torch.equal(batched[2:3], torch.multinomial(probs[2:3], num_samples=1, generator=seeded(22)))
+
+    def test_forward_per_row_generators_are_row_independent(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        torch.manual_seed(123)
+        row_embed = torch.randn(1, hidden)
+        row_hidden = torch.randn(1, hidden)
+        # Three identical rows: two share a seed, one differs.
+        layer0_code = torch.zeros(3, dtype=torch.long)
+        layer0_embed = row_embed.expand(3, hidden).contiguous()
+        last_talker_hidden = row_hidden.expand(3, hidden).contiguous()
+
+        def seeded(seed: int) -> torch.Generator:
+            generator = torch.Generator(device=layer0_code.device)
+            generator.manual_seed(seed)
+            return generator
+
+        def run(generators):
+            return predictor(
+                layer0_code=layer0_code,
+                layer0_embed=layer0_embed,
+                last_talker_hidden=last_talker_hidden,
+                do_sample=True,
+                temperature=0.9,
+                top_k=50,
+                top_p=1.0,
+                generators=generators,
+            )
+
+        first = run([seeded(1234), seeded(1234), seeded(4321)])
+        second = run([seeded(1234), seeded(1234), seeded(4321)])
+
+        # Same per-row seeds -> the whole batch reproduces across calls.
+        assert torch.equal(first, second)
+        # Identical inputs + identical seeds -> identical rows within a call.
+        assert torch.equal(first[0], first[1])
+        # A different seed on the same inputs must diverge in the residual layers.
+        assert not torch.equal(first[0, 1:], first[2, 1:])
+
+    def test_forward_rejects_mismatched_generators_length(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        predictor, talker_config = self._make_predictor(mocker, loaded_target_classes)
+        hidden = talker_config.hidden_size
+        with pytest.raises(ValueError, match="one entry per row"):
+            predictor(
+                layer0_code=torch.zeros(2, dtype=torch.long),
+                layer0_embed=torch.randn(2, hidden),
+                last_talker_hidden=torch.randn(2, hidden),
+                generators=[torch.Generator()],
+            )
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 
@@ -457,3 +585,92 @@ class TestCodePredictorWrapperConfig:
             wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True),
         )
         assert graph_wrapper._prefix_graphs_enabled is True
+
+
+class TestGumbelMaxSampling:
+    """Sanity tests for the shared Gumbel-max sampling helper."""
+
+    @staticmethod
+    def _sample(logits: torch.Tensor, generator=None) -> torch.Tensor:
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        return common_mod.CodePredictorWrapper._sample_codes_gumbel(logits, generator=generator)
+
+    def test_gumbel_max_picks_unique_high_logit(self, loaded_target_classes) -> None:
+        """If one logit is dominant, Gumbel-max must always pick it."""
+        torch.manual_seed(0)
+        bsz, vocab = 4, 64
+        logits = torch.full((bsz, vocab), -1e4)
+        logits[:, 7] = 0.0  # only finite entry
+        out = self._sample(logits)
+        assert out.shape == (bsz, 1)
+        assert torch.equal(out.squeeze(-1), torch.full((bsz,), 7, dtype=torch.long))
+
+    def test_gumbel_max_handles_single_finite_entry_rows(self, loaded_target_classes) -> None:
+        """Degenerate rows with one surviving finite entry should still return
+        a valid index instead of tripping multinomial-style asserts."""
+        bsz, vocab = 2, 8
+        logits = torch.full((bsz, vocab), float("-inf"))
+        logits[0, 3] = 0.0
+        logits[1, 5] = 0.0
+        out = self._sample(logits)
+        assert out[0].item() == 3
+        assert out[1].item() == 5
+
+    def test_gumbel_max_generator_makes_results_reproducible(self, loaded_target_classes) -> None:
+        """Same ``Generator.manual_seed`` should produce identical samples."""
+        torch.manual_seed(0)
+        logits = torch.randn(8, 32)
+        g1 = torch.Generator()
+        g1.manual_seed(123)
+        g2 = torch.Generator()
+        g2.manual_seed(123)
+        g3 = torch.Generator()
+        g3.manual_seed(456)
+        s1 = self._sample(logits, generator=g1)
+        s2 = self._sample(logits, generator=g2)
+        s3 = self._sample(logits, generator=g3)
+        assert torch.equal(s1, s2)
+        # Different seeds must (almost surely) yield a different sequence.
+        assert not torch.equal(s1, s3)
+
+    def test_gumbel_max_per_row_generators_match_individual_rows(self, loaded_target_classes) -> None:
+        """A per-row generator list should preserve row independence."""
+        torch.manual_seed(11)
+        logits = torch.randn(3, 16)
+        seeds = [101, 202, 303]
+        batched_generators = []
+        expected_rows = []
+        for row, seed in enumerate(seeds):
+            row_generator = torch.Generator()
+            row_generator.manual_seed(seed)
+            batched_generators.append(row_generator)
+
+            solo_generator = torch.Generator()
+            solo_generator.manual_seed(seed)
+            expected_rows.append(self._sample(logits[row : row + 1], generator=solo_generator))
+
+        batched = self._sample(logits, generator=batched_generators)
+        expected = torch.cat(expected_rows, dim=0)
+        assert torch.equal(batched, expected)
+
+    def test_gumbel_max_rejects_generator_length_mismatch(self, loaded_target_classes) -> None:
+        logits = torch.randn(2, 8)
+        generators = [torch.Generator()]
+        with pytest.raises(ValueError, match="Expected 2 per-row generators"):
+            self._sample(logits, generator=generators)
+
+    def test_gumbel_max_distribution_matches_softmax(self, loaded_target_classes) -> None:
+        """Empirical frequency under Gumbel-max should be close to the
+        softmax probability mass.  Loose tolerance because we draw a
+        finite number of samples."""
+        torch.manual_seed(7)
+        vocab = 16
+        logits = torch.randn(vocab) * 2.0
+        target = F.softmax(logits, dim=-1, dtype=torch.float32)
+
+        n_draws = 20000
+        batched_logits = logits.unsqueeze(0).expand(n_draws, vocab).contiguous()
+        samples = self._sample(batched_logits).squeeze(-1)
+        empirical = torch.bincount(samples, minlength=vocab).float() / n_draws
+        # L1 distance under 0.03 with 20k draws is comfortable for vocab=16.
+        assert (empirical - target).abs().sum().item() < 0.05
diff --git a/tests/model_executor/models/test_kv_cache_scale_mapper.py b/tests/model_executor/models/test_kv_cache_scale_mapper.py
new file mode 100644
index 000000000..98f016bb4
--- /dev/null
+++ b/tests/model_executor/models/test_kv_cache_scale_mapper.py
@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for vLLM's KV-cache scale mapper load path."""
+
+from collections.abc import Iterable
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_STALE_API_FILES = [
+    "vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py",
+    "vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py",
+    "vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py",
+    "vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py",
+]
+
+_SOURCE_SCALE_NAME = "layers.0.self_attn.k_proj.output_scale"
+_MAPPED_SCALE_NAME = "layers.0.self_attn.attn.k_scale"
+
+
+class _RecordingCacheScaleMapper:
+    def __init__(self) -> None:
+        self.called = False
+
+    def apply(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> Iterable[tuple[str, torch.Tensor]]:
+        self.called = True
+        for name, weight in weights:
+            yield name.replace(".k_proj.output_scale", ".attn.k_scale"), weight
+
+
+class _QuantConfigWithoutOldCacheScale:
+    def __init__(self) -> None:
+        self.mapper = _RecordingCacheScaleMapper()
+
+    def get_cache_scale_mapper(self) -> _RecordingCacheScaleMapper:
+        return self.mapper
+
+
+class _FakeModel:
+    def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None:
+        self.quant_config = quant_config
+        self.param = torch.nn.Parameter(torch.zeros(()), requires_grad=False)
+
+    def named_parameters(
+        self,
+        remove_duplicate: bool = True,
+    ) -> Iterable[tuple[str, torch.nn.Parameter]]:
+        del remove_duplicate
+        return [(_MAPPED_SCALE_NAME, self.param)]
+
+
+class _FakeHunyuanModel(_FakeModel):
+    def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None:
+        super().__init__(quant_config)
+        self.config = SimpleNamespace(
+            tie_word_embeddings=False,
+            num_attention_heads=1,
+            num_key_value_heads=1,
+            use_cla=False,
+        )
+
+    def get_expert_mapping(self) -> tuple[list[object], dict[str, object]]:
+        return [], {}
+
+    def _split_qkv_weight(self, weight: torch.Tensor) -> torch.Tensor:
+        return weight
+
+
+@pytest.mark.parametrize("rel_path", _STALE_API_FILES)
+def test_model_loaders_do_not_call_removed_get_cache_scale(rel_path: str) -> None:
+    """Affected loaders must not call the removed vLLM get_cache_scale API."""
+    source = (_REPO_ROOT / rel_path).read_text(encoding="utf-8")
+    assert ".get_cache_scale(" not in source
+
+
+@pytest.mark.parametrize(
+    ("module_name", "class_name", "fake_model_cls"),
+    [
+        (
+            "vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3",
+            "HunyuanModel",
+            _FakeHunyuanModel,
+        ),
+        (
+            "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old",
+            "Qwen2Model",
+            _FakeModel,
+        ),
+    ],
+)
+def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale(
+    module_name: str,
+    class_name: str,
+    fake_model_cls: type[_FakeModel],
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Delegated loaders must not expect the removed get_cache_scale method."""
+    module = __import__(module_name, fromlist=[class_name])
+    monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False)
+
+    fake_model = fake_model_cls(_QuantConfigWithoutOldCacheScale())
+    load_weights = getattr(module, class_name).load_weights
+
+    loaded = load_weights(fake_model, [(_MAPPED_SCALE_NAME, torch.ones(()))])
+
+    assert loaded == {_MAPPED_SCALE_NAME}
+    torch.testing.assert_close(fake_model.param.data, torch.ones(()))
+
+
+@pytest.mark.parametrize(
+    ("module_name", "class_name"),
+    [
+        (
+            "vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2",
+            "MammothModa2Qwen2ForCausalLM",
+        ),
+        (
+            "vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm",
+            "MiMoAudioLLMForConditionalGeneration",
+        ),
+    ],
+)
+def test_direct_custom_loaders_apply_cache_scale_mapper(
+    module_name: str,
+    class_name: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Direct custom loaders must apply the new mapper before fall-through load."""
+    module = __import__(module_name, fromlist=[class_name])
+    monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False)
+
+    quant_config = _QuantConfigWithoutOldCacheScale()
+    fake_model = _FakeModel(quant_config)
+    load_weights = getattr(module, class_name).load_weights
+
+    loaded = load_weights(fake_model, [(_SOURCE_SCALE_NAME, torch.ones(()))])
+
+    assert quant_config.mapper.called
+    assert loaded == {_MAPPED_SCALE_NAME}
+    torch.testing.assert_close(fake_model.param.data, torch.ones(()))
diff --git a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py
new file mode 100644
index 000000000..8154ddd65
--- /dev/null
+++ b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py
@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Static regression checks for VoxCPM2 runner-assisted unified graph hooks.
+
+These tests intentionally avoid importing torch/vLLM so they can run in a
+lightweight local checkout. Runtime audio quality still requires a CUDA test.
+"""
+
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[4]
+RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py"
+RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py"
+TALKER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "voxcpm2_talker.py"
+VOXCPM2_RUNTIME_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "runtime_config.py"
+VOXCPM2_PIPELINE = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "pipeline.py"
+VOXCPM2_SCHEDULER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "scheduler.py"
+OMNI_AR_SCHEDULER = REPO_ROOT / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py"
+DEPLOY = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm2.yaml"
+
+
+def test_ar_runner_exposes_runner_assisted_full_metadata_hook():
+    source = RUNNER.read_text()
+
+    assert "Models without this hook keep the normal runner path" in source
+    assert "get_runner_assisted_full_attention_metadata_request" in source
+    assert "set_runner_assisted_full_attention_metadata_context" in source
+    assert "runner_assisted_full_attn" in source
+    assert "runner_assisted_full_attn_capture" in source
+    assert "pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn" in source
+    assert "for_cudagraph_capture=runner_assisted_full_attn_capture" in source
+    assert "cudagraph_runtime_mode=(" in source
+    assert "CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode" in source
+    refresh_source = source[source.index("def _refresh_runner_assisted_full_attention_metadata_buffers") :]
+    refresh_source = refresh_source[: refresh_source.index("def _set_runner_assisted_full_attention_metadata_context")]
+    assert "num_computed_tokens_cpu" in refresh_source
+    assert "num_scheduled_tokens_np" in refresh_source
+    assert "optimistic_seq_lens_cpu[:num_reqs]=1" not in "".join(refresh_source.split())
+    padding_source = source[
+        source.index("runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request") :
+    ]
+    padding_source = padding_source[: padding_source.index("ubatch_slices, ubatch_slices_padded")]
+    assert "runner_assisted_full_attn_request.num_reqs_padded" in padding_source
+    assert "runner_assisted_full_attn_request.for_cudagraph_capture" in padding_source
+    assert "num_reqs_padded * max_num_scheduled_tokens" in padding_source
+
+
+def test_ar_runner_without_model_hook_stays_on_normal_path():
+    source = RUNNER.read_text()
+    contract_source = RUNNER_ASSISTED_METADATA.read_text()
+    request_source = source[source.index("def _get_runner_assisted_full_attention_metadata_request") :]
+    request_source = request_source[
+        : request_source.index("def _refresh_runner_assisted_full_attention_metadata_buffers")
+    ]
+    compact_request_source = "".join(request_source.split())
+
+    assert "Models without this hook keep the normal runner path" in request_source
+    assert "class RunnerAssistedAttentionMetadataProvider(Protocol)" in contract_source
+    assert "class RunnerAssistedFullAttentionMetadataRequest(NamedTuple)" in contract_source
+    assert "tuple[int, bool]" not in contract_source
+    assert 'hook=getattr(self.model,"get_runner_assisted_full_attention_metadata_request",None)' in (
+        compact_request_source
+    )
+    assert "ifnotcallable(hook):returnNone" in compact_request_source
+    assert ".tolist()" not in compact_request_source
+    assert "isinstance(request,RunnerAssistedFullAttentionMetadataRequest)" in compact_request_source
+    assert "request.num_reqs_padded" in request_source
+    assert "request.for_cudagraph_capture" in request_source
+    assert "exceptException" not in compact_request_source
+
+    set_context_source = source[source.index("def _set_runner_assisted_full_attention_metadata_context") :]
+    set_context_source = set_context_source[: set_context_source.index("def _deferred_prefix_cache_mm_keys")]
+    assert "except Exception" not in set_context_source
+
+
+def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise():
+    source = TALKER.read_text()
+    runtime_config_source = VOXCPM2_RUNTIME_CONFIG.read_text()
+    compact_source = "".join(source.split())
+
+    assert "_voxcpm2_compile_without_inductor_cudagraphs" in source
+    assert "def _compile_without_inductor_cudagraphs" not in source
+    assert "_voxcpm2_unwrap_torch_compile" in source
+    assert "_orig_mod" in source
+    compile_source = source[source.index("def _voxcpm2_compile_without_inductor_cudagraphs") :]
+    compile_source = compile_source[: compile_source.index("def _setup_torch_compile")]
+    assert "torch._inductor.config" not in compile_source
+    assert '"triton.cudagraphs": False' in compile_source
+    assert '"triton.cudagraph_trees": False' in compile_source
+
+    assert "self._enable_unified_decode_graph=self._runtime_config.unified_decode_graph_available(" in compact_source
+    assert "andnotself.deterministic_cfm_noise" in "".join(runtime_config_source.split())
+    assert "decode_tail_graph" not in source
+
+    unified_source = source[source.index("def _forward_unified_decode") :]
+    unified_source = unified_source[: unified_source.index("# -------------------- vllm hooks")]
+    assert "except Exception:" not in unified_source
+    assert "_failed_unified_graph_sizes" not in source
+    assert "self._enable_unified_decode_graph = False" not in unified_source
+    assert "_forward_unified_decode_fallback" not in source
+    assert "capture_failed" not in compact_source
+    assert "record_fallback" not in source
+    assert "fallbacks" not in source
+    assert "**info_dict: Unpack[VoxCPM2PreprocessInput]" in source
+    assert "**info: Unpack[VoxCPM2PostprocessInput]" in source
+    assert "reference_audio: Any" not in source
+    assert "prompt_audio: Any" not in source
+    assert "def _nullify_volatile_metadata(ctx: _ForwardContextLike) -> _ForwardContextLike" in source
+    assert "class _PrefillInputs(NamedTuple)" in source
+    assert "-> _PrefillInputs" in source
+    assert 'inputs["audio_feat"]' not in source
+    assert 'inputs["text_token"]' not in source
+
+    capture_source = source[source.index("def _capture_unified_decode_graph") :]
+    capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")]
+    assert "_voxcpm2_compile_unified_capture_estimator" in capture_source
+    assert "_voxcpm2_compile_unified_capture_feat_encoder" in capture_source
+    graph_body = capture_source[capture_source.index("with torch.cuda.graph") :]
+    assert "g.cfm_noise.normal_()" not in graph_body
+    assert "g.cfm_noise.normal_()" in unified_source
+
+    assert "_pre_capture_unified_graphs" not in source
+    assert "_unified_graph_pre_capture_sizes" not in source
+    assert "unified_decode_graph_pre_capture_sizes" not in source
+
+
+def test_voxcpm2_batch_unified_graph_requires_runner_metadata_marker():
+    source = TALKER.read_text()
+    compact_source = "".join(source.split())
+
+    assert "enable_runner_assisted_unified_decode_graph" not in source
+    assert "allow_unified_decode_graph_batch_attention" not in source
+    assert "get_runner_assisted_full_attention_metadata_request" in source
+    assert "set_runner_assisted_full_attention_metadata_context" in source
+    assert "RunnerAssistedFullAttentionMetadataRequest" in source
+    assert "_runner_assisted_unified_decode_graph_active" in source
+    assert "runner_full_metadata_missing" in source
+    assert "_build_unified_graph_bucket_sizes" in source
+    assert "_select_unified_graph_bucket_size" in source
+
+    capture_source = source[source.index("def _capture_unified_decode_graph") :]
+    capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")]
+    assert "override_forward_context" in capture_source
+    assert "_nullify_volatile_metadata" in capture_source
+    assert "capture_context = override_forward_context(self._nullify_volatile_metadata(ctx))" in capture_source
+    assert "ifsize>1:continue" not in compact_source
+    forward_source = source[source.index("def _forward_unified_decode") :]
+    forward_source = forward_source[: forward_source.index("# -------------------- vllm hooks")]
+    assert "graph_size = self._select_unified_graph_bucket_size(num_reqs)" in forward_source
+    assert "self._unified_graphs[graph_size]" in forward_source
+    assert "g.input_embeds[num_reqs:graph_size].zero_()" not in forward_source
+    assert "padded rows as valid duplicates" in forward_source
+    assert "g.input_embeds[last : last + 1].expand" in forward_source
+    assert "ifnum_reqs>1andnotself._runner_assisted_unified_decode_graph_active" in compact_source
+    assert "andnotcfg.enable_runner_assisted_unified_decode_graph" not in compact_source
+    needs_source = source[source.index("def get_runner_assisted_full_attention_metadata_request") :]
+    needs_source = needs_source[: needs_source.index("def set_runner_assisted_full_attention_metadata_context")]
+    assert "_select_unified_graph_bucket_size(num_reqs)" in needs_source
+    assert "_should_use_decode_graph(num_reqs)" not in needs_source
+    assert "RunnerAssistedFullAttentionMetadataRequest(" in needs_source
+
+
+def test_voxcpm2_unified_skip_preserves_segmented_decode_graphs():
+    source = TALKER.read_text()
+    compact_source = "".join(source.split())
+    forward_source = source[source.index("def forward(") :]
+    forward_source = forward_source[: forward_source.index("# -------------------- prefill / decode helpers")]
+
+    assert "use_segmented_decode_graph" in forward_source
+    assert "unified_skip_reasonisnotNone" in compact_source
+    assert "notself._enable_unified_decode_graphorunified_skip_reasonisnotNone" in compact_source
+    assert "ifuse_segmented_decode_graph:" in compact_source
+    assert "anduse_segmented_decode_graph" in compact_source
+
+
+def test_voxcpm2_scheduler_policy_stays_model_local():
+    common_source = OMNI_AR_SCHEDULER.read_text()
+    voxcpm2_scheduler_source = VOXCPM2_SCHEDULER.read_text()
+    pipeline_source = VOXCPM2_PIPELINE.read_text()
+
+    assert "VoxCPM2TalkerForConditionalGeneration" not in common_source
+    assert "voxcpm2_runtime_config" not in common_source
+    assert "pure_decode_graph" not in common_source
+    assert "_schedule_with_optional_waiting_deferral" not in common_source
+    assert "def _should_defer_waiting_admission(self) -> bool:" in common_source
+
+    assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source
+    assert "from .runtime_config import _VoxCPM2RuntimeConfig" in voxcpm2_scheduler_source
+    assert "_VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)" in voxcpm2_scheduler_source
+    assert "unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())" in voxcpm2_scheduler_source
+    assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source
+    assert "def schedule(" not in voxcpm2_scheduler_source
+    assert "create_request_queue" not in voxcpm2_scheduler_source
+    assert "original_waiting.prepend_requests(deferred_waiting)" not in voxcpm2_scheduler_source
+    assert "return self._should_defer_waiting_for_unified_decode_graph()" in voxcpm2_scheduler_source
+    assert "vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler" in pipeline_source
+
+
+def test_voxcpm2_deploy_defaults_to_full_unified_graph_only():
+    source = DEPLOY.read_text()
+
+    assert "max_num_seqs: 8" in source
+    assert "enable_unified_decode_graph: true" in source
+    assert "unified_decode_graph_max_batch_size: 8" in source
+    assert "unified_decode_graph_pre_capture_sizes" not in source
+    assert "enable_runner_assisted_unified_decode_graph" not in source
+    assert "allow_unified_decode_graph_batch_attention" not in source
+    assert "decode_graph_capture_policy" not in source
diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py
index 28da0b1bc..f5a94376a 100644
--- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py
+++ b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py
@@ -52,6 +52,10 @@ def _make_bare_talker():
     talker._cfm_buffers = None
     talker._last_audio_output_req_ids = []
     talker._batched_fsq_fusion_max_batch = 32
+    # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to
+    # default (False) to avoid AttributeError when test exercises forward
+    # path that checks this flag.
+    talker._enable_unified_decode_graph = False
     return talker
 
 
@@ -329,23 +333,27 @@ class TestDecodeBatchContract:
 
         for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True):
             torch.testing.assert_close(actual_res, expected_res)
-            torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"])
+            torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden)
         assert [state.decode_step_count for state in batch_states] == [1, 1]
 
     def test_batch_decode_matches_per_request_state_updates(self) -> None:
+        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+            _DecodeResidualMeta,
+        )
+
         seq_talker = self._make_decode_talker()
         batch_talker = self._make_decode_talker()
 
         seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         metas = [
-            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
-            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
+            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
+            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
         ]
         batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
 
         for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True):
-            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu"))
+            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0))
 
         batch_talker._finish_decode_batch(
             [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)],
@@ -360,12 +368,16 @@ class TestDecodeBatchContract:
             torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu)
 
     def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None:
+        from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
+            _DecodeResidualMeta,
+        )
+
         talker = self._make_decode_talker()
         talker._enable_batched_cfm = False
         states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
         metas = [
-            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
-            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
+            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])),
+            _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])),
         ]
         batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
 
diff --git a/tests/model_executor/models/voxtral_tts/test_tts_preprocess.py b/tests/model_executor/models/voxtral_tts/test_tts_preprocess.py
new file mode 100644
index 000000000..399a7ecbd
--- /dev/null
+++ b/tests/model_executor/models/voxtral_tts/test_tts_preprocess.py
@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import functools
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@functools.lru_cache(maxsize=1)
+def _voxtral_tts_model_cls():
+    from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules
+
+    bootstrap_vllm_layer_custom_op_modules()
+    import vllm.model_executor.models.utils  # noqa: F401
+
+    from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import (
+        VoxtralTTSForConditionalGeneration,
+    )
+
+    return VoxtralTTSForConditionalGeneration
+
+
+class FakeAudioGeneration(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed_multimodal_calls = []
+
+    def embed_multimodal(self, **kwargs):
+        self.embed_multimodal_calls.append(kwargs)
+        return [torch.full((1, 4), 7.0)]
+
+
+def _make_voxtral_tts_model():
+    model_cls = _voxtral_tts_model_cls()
+    model = model_cls.__new__(model_cls)
+    nn.Module.__init__(model)
+    model.model_stage = "audio_generation"
+    model.model = FakeAudioGeneration()
+    model._audio_token_id = 42
+    return model
+
+
+def test_tts_preprocess_consumes_nested_codes_audio_feedback():
+    model = _make_voxtral_tts_model()
+    input_ids = torch.tensor([model._audio_token_id])
+    input_embeds = torch.zeros((1, 4))
+    audio_tokens = torch.tensor([[1, 2, 3, 4]])
+
+    _, output_embeds, _ = model.tts_preprocess(
+        input_ids=input_ids,
+        input_embeds=input_embeds,
+        codes={"audio": audio_tokens},
+    )
+
+    assert len(model.model.embed_multimodal_calls) == 1
+    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
+    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))
+
+
+def test_tts_preprocess_keeps_legacy_top_level_audio_feedback():
+    model = _make_voxtral_tts_model()
+    input_ids = torch.tensor([model._audio_token_id])
+    input_embeds = torch.zeros((1, 4))
+    audio_tokens = torch.tensor([[5, 6, 7, 8]])
+
+    _, output_embeds, _ = model.tts_preprocess(
+        input_ids=input_ids,
+        input_embeds=input_embeds,
+        audio=audio_tokens,
+    )
+
+    assert len(model.model.embed_multimodal_calls) == 1
+    torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens)
+    torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0))
diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
index 83168b4f5..5309c7675 100644
--- a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
+++ b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py
@@ -8,7 +8,6 @@ import torch
 
 from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     talker2code2wav_async_chunk,
-    text2flow,
     text2flow_full_payload,
     text2flow_token_only,
 )
@@ -49,13 +48,13 @@ def _transfer_manager(
     )
 
 
-def test_text2flow_supports_batched_source_outputs():
+def test_text2flow_token_only_supports_batched_source_outputs():
     source_outputs = [
         _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
         _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[6, 7]])}),
     ]
 
-    outputs = text2flow(source_outputs=source_outputs, prompt=None)
+    outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None)
 
     assert len(outputs) == 2
     assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
@@ -64,16 +63,6 @@ def test_text2flow_supports_batched_source_outputs():
     assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21]
 
 
-def test_text2flow_strips_reference_speech_prefix_from_cumulative_ids():
-    source_outputs = [
-        _source_output("req-0", [10, 11], [8, 9, 1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}),
-    ]
-
-    outputs = text2flow(source_outputs=source_outputs, prompt=None)
-
-    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
-
-
 def test_text2flow_token_only_strips_reference_speech_prefix_from_cumulative_ids():
     source_outputs = [
         _source_output(
diff --git a/tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py b/tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py
index 76f3e5006..19405f28f 100644
--- a/tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py
+++ b/tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py
@@ -45,6 +45,33 @@ def test_truncate_at_cot_end_strips_tail_after_recaption_marker():
     assert text == "body text</recaption>"
 
 
+def test_ar2diffusion_returns_one_request_payload_for_request_level_batching():
+    result = ar2diffusion(
+        [_source_output([100], text="thought")],
+        prompt={"prompt": "edit"},
+    )
+
+    assert isinstance(result, dict)
+    assert result["prompt"] == "edit"
+
+
+def test_ar2diffusion_uses_parent_output_when_companions_are_present():
+    result = ar2diffusion(
+        [
+            _source_output([100], text="parent thought"),
+            _source_output([200], text="companion thought"),
+        ],
+        prompt={"prompt": "edit"},
+    )
+
+    assert result is not None
+    assert result["extra"]["ar_generated_text"] == "parent thought"
+
+
+def test_ar2diffusion_returns_none_without_parent_output():
+    assert ar2diffusion([], prompt={"prompt": "edit"}) is None
+
+
 def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeypatch: pytest.MonkeyPatch):
     real_import = builtins.__import__
 
@@ -67,10 +94,10 @@ def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeyp
         prompt=[{"prompt": "edit", "height": 64, "width": 64}],
     )
 
-    assert len(result) == 1
-    assert (result[0]["height"], result[0]["width"]) == (512, 2048)
-    assert result[0]["extra"]["ar_generated_text"] == "decoded without special tokens"
-    assert "ar_token_ids" not in result[0]["extra"]
+    assert result is not None
+    assert (result["height"], result["width"]) == (512, 2048)
+    assert result["extra"]["ar_generated_text"] == "decoded without special tokens"
+    assert "ar_token_ids" not in result["extra"]
 
 
 def test_ar2diffusion_forwards_custom_system_prompt_body():
@@ -88,5 +115,6 @@ def test_ar2diffusion_forwards_custom_system_prompt_body():
         ],
     )
 
-    assert result[0]["use_system_prompt"] == "custom"
-    assert result[0]["system_prompt"] == marker
+    assert result is not None
+    assert result["use_system_prompt"] == "custom"
+    assert result["system_prompt"] == marker
diff --git a/tests/model_executor/stage_input_processors/test_indextts2.py b/tests/model_executor/stage_input_processors/test_indextts2.py
index 59f14b414..11468603e 100644
--- a/tests/model_executor/stage_input_processors/test_indextts2.py
+++ b/tests/model_executor/stage_input_processors/test_indextts2.py
@@ -7,7 +7,6 @@ import torch
 
 from vllm_omni.model_executor.stage_input_processors.indextts2 import (
     _strip_stop_token,
-    talker2s2mel,
     talker2s2mel_full_payload,
     talker2s2mel_token_only,
 )
@@ -108,118 +107,6 @@ def test_strip_stop_token_all_stop():
     assert lat.shape[1] == 0
 
 
-# ── talker2s2mel ──
-
-
-def test_talker2s2mel_basic():
-    mel_codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN])
-    latent = torch.randn(4, LATENT_DIM)
-    s_ref = torch.randn(1, 128)
-    ref_mel = torch.randn(1, 80, 200)
-    style = torch.randn(1, 256)
-
-    outputs = talker2s2mel(
-        source_outputs=[
-            _make_talker_output(
-                mel_codes=mel_codes,
-                latent=latent,
-                s_ref=s_ref,
-                ref_mel=ref_mel,
-                style=style,
-            )
-        ]
-    )
-
-    assert len(outputs) == 1
-    info = outputs[0]["additional_information"]
-    assert info["mel_codes"].shape == (1, 3)
-    assert info["latent"].shape == (1, 3, LATENT_DIM)
-    assert info["code_lens"].tolist() == [3]
-    assert "S_ref" in info
-    assert "ref_mel" in info
-    assert "style" in info
-    assert outputs[0]["prompt_token_ids"] == [0]
-
-
-def test_talker2s2mel_skips_unfinished():
-    mel_codes = torch.tensor([10, 20])
-    latent = torch.randn(2, LATENT_DIM)
-
-    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, finished=False)])
-
-    assert len(outputs) == 0
-
-
-def test_talker2s2mel_skips_empty_mel():
-    outputs = talker2s2mel(
-        source_outputs=[
-            _make_talker_output(
-                mel_codes=torch.tensor([]),
-                latent=torch.randn(0, LATENT_DIM),
-            )
-        ]
-    )
-
-    assert len(outputs) == 0
-
-
-def test_talker2s2mel_missing_optional_meta():
-    mel_codes = torch.tensor([10, 20])
-    latent = torch.randn(2, LATENT_DIM)
-
-    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent)])
-
-    assert len(outputs) == 1
-    info = outputs[0]["additional_information"]
-    assert "S_ref" not in info
-    assert "ref_mel" not in info
-    assert "style" not in info
-
-
-def test_talker2s2mel_batch():
-    out0 = _make_talker_output(
-        mel_codes=torch.tensor([1, 2, STOP_MEL_TOKEN]),
-        latent=torch.randn(3, LATENT_DIM),
-        s_ref=torch.randn(1, 64),
-    )
-    out1 = _make_talker_output(
-        mel_codes=torch.tensor([5, 6, 7, 8]),
-        latent=torch.randn(4, LATENT_DIM),
-        ref_mel=torch.randn(1, 80, 100),
-    )
-
-    outputs = talker2s2mel(source_outputs=[out0, out1])
-
-    assert len(outputs) == 2
-    assert outputs[0]["additional_information"]["code_lens"].tolist() == [2]
-    assert outputs[1]["additional_information"]["code_lens"].tolist() == [4]
-
-
-def test_talker2s2mel_tensors_on_cpu():
-    device = torch.device("cpu")
-    mel_codes = torch.tensor([10, 20], device=device)
-    latent = torch.randn(2, LATENT_DIM, device=device)
-    s_ref = torch.randn(1, 64, device=device)
-
-    outputs = talker2s2mel(source_outputs=[_make_talker_output(mel_codes=mel_codes, latent=latent, s_ref=s_ref)])
-
-    info = outputs[0]["additional_information"]
-    assert info["latent"].device.type == "cpu"
-    assert info["mel_codes"].device.type == "cpu"
-    assert info["S_ref"].device.type == "cpu"
-
-
-def test_talker2s2mel_no_multimodal_dict():
-    output = SimpleNamespace(
-        finished=True,
-        outputs=[SimpleNamespace(multimodal_output="not_a_dict")],
-    )
-
-    results = talker2s2mel(source_outputs=[output])
-
-    assert len(results) == 0
-
-
 # ── full-payload connector path ──
 
 
diff --git a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py
index d1cc0dbf7..ae1c8e2f9 100644
--- a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py
+++ b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py
@@ -11,59 +11,40 @@ import torch
 from vllm_omni.model_executor.stage_input_processors import mimo_audio as sip
 from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
     MAX_CODE2WAV_TOKENS,
-    llm2code2wav,
+    llm2code2wav_full_payload,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"):
-    output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
-    """Build minimal source_outputs with one talker output carrying codec codes."""
-    talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
-    return [talker_output]
-
-
-def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
-    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated, not passed through."""
-    # prepend_and_flatten_colmajor produces 36 ids per (8, 4) codec frame:
-    # pad adds one row -> (9, 4) per frame, permuted and flattened.
-    # Pick enough frames to comfortably exceed the cap.
+def test_llm2code2wav_full_payload_truncates_when_flat_exceeds_max(caplog):
+    """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated."""
     frames = (MAX_CODE2WAV_TOKENS // 36) + 100
     codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long)
+    request = SimpleNamespace(request_id="req-long")
 
-    source_outputs = _make_source_outputs(codec_codes, request_id="req-long")
-
-    # Attach caplog's handler directly to the module logger so the warning is
-    # captured regardless of propagation (vllm's logger configuration can
-    # interact badly with caplog.at_level's default root-handler path).
     target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.mimo_audio")
     target_logger.addHandler(caplog.handler)
     prev_level = target_logger.level
     target_logger.setLevel(logging.WARNING)
     try:
-        prompts = llm2code2wav(source_outputs)
+        payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
     finally:
         target_logger.removeHandler(caplog.handler)
         target_logger.setLevel(prev_level)
 
-    assert len(prompts) == 1
-    assert len(prompts[0]["prompt_token_ids"]) == MAX_CODE2WAV_TOKENS
-    assert any("truncating" in rec.getMessage() for rec in caplog.records), (
-        f"Expected a 'truncating' warning; captured records: {[r.getMessage() for r in caplog.records]}"
-    )
+    assert payload is not None
+    assert len(payload["codes"]["audio"]) == MAX_CODE2WAV_TOKENS
 
 
-def test_llm2code2wav_short_sequence_unchanged():
-    """Short codec sequences are returned without truncation."""
+def test_llm2code2wav_full_payload_short_sequence_unchanged():
     codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long)
-    source_outputs = _make_source_outputs(codec_codes, request_id="req-short")
+    request = SimpleNamespace(request_id="req-short")
 
-    prompts = llm2code2wav(source_outputs)
+    payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request)
 
-    assert len(prompts) == 1
-    # 4 frames + 1 pad row, flattened col-major → well below the cap
-    assert 0 < len(prompts[0]["prompt_token_ids"]) <= MAX_CODE2WAV_TOKENS
+    assert payload is not None
+    assert 0 < len(payload["codes"]["audio"]) <= MAX_CODE2WAV_TOKENS
 
 
 def test_llm2code2wav_truncation_boundary_constant_matches_yaml():
diff --git a/tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py
new file mode 100644
index 000000000..bff3e31e3
--- /dev/null
+++ b/tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py
@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.ming_tts.constants import LATENT_DIM, PATCH_SIZE
+from vllm_omni.model_executor.stage_input_processors.ming_tts import (
+    MING_EMIT_PATCH_COUNT_KEY,
+    llm2audio_vae_async_chunk,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_transfer_manager(*, chunk_size: int = 5, initial_chunk_size: int = 2):
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        put_req_chunk=defaultdict(int),
+        request_payload={},
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "latent_chunk_size": chunk_size,
+                    "initial_latent_chunk_size": initial_chunk_size,
+                    "latent_left_context": 0,
+                }
+            }
+        ),
+    )
+
+
+def _make_request(req_id: str = "req", *, finished: bool = False):
+    return SimpleNamespace(
+        external_req_id=req_id,
+        is_finished=lambda: finished,
+    )
+
+
+def _make_output(value: float):
+    return {
+        "ming_has_patch": torch.tensor([1], dtype=torch.bool),
+        "ming_latent_patch": torch.full((1, PATCH_SIZE, LATENT_DIM), value),
+    }
+
+
+def _append_patch(tm, req_id: str, idx: int, *, finished: bool = False):
+    return llm2audio_vae_async_chunk(
+        tm,
+        _make_output(float(idx)),
+        _make_request(req_id, finished=finished),
+        is_finished=finished,
+    )
+
+
+def test_ming_async_chunk_emits_small_initial_chunk_then_steady_chunks():
+    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
+
+    assert _append_patch(tm, "req", 0) is None
+
+    first = _append_patch(tm, "req", 1)
+    assert first is not None
+    assert first.latent.shape == (2, PATCH_SIZE, LATENT_DIM)
+    assert first.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 2
+    assert first.meta.finished.item() is False
+
+    for idx in range(2, 6):
+        assert _append_patch(tm, "req", idx) is None
+
+    second = _append_patch(tm, "req", 6)
+    assert second is not None
+    assert second.latent.shape == (5, PATCH_SIZE, LATENT_DIM)
+    assert second.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 5
+
+
+def test_ming_async_chunk_flushes_short_final_chunk():
+    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
+
+    payload = _append_patch(tm, "req", 0, finished=True)
+
+    assert payload is not None
+    assert payload.latent.shape == (1, PATCH_SIZE, LATENT_DIM)
+    assert payload.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 1
+    assert payload.meta.finished.item() is True
+
+
+def test_ming_async_chunk_flushes_leftover_after_initial_chunk():
+    tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2)
+
+    assert _append_patch(tm, "req", 0) is None
+    assert _append_patch(tm, "req", 1) is not None
+    assert _append_patch(tm, "req", 2) is None
+
+    final = _append_patch(tm, "req", 3, finished=True)
+
+    assert final is not None
+    assert final.latent.shape == (2, PATCH_SIZE, LATENT_DIM)
+    assert final.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 2
+    assert final.meta.finished.item() is True
diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py
index 69b612a05..8a6dfbe94 100644
--- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py
+++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py
@@ -163,7 +163,7 @@ def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
 
     assert payload is not None
     assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22]
-    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None:
@@ -204,7 +204,7 @@ def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> N
 
     assert payload is not None
     assert payload["codes"]["audio"] == [0, 0, 0]
-    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
@@ -224,7 +224,7 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
 
     assert payload is not None
     assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9]
-    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+    assert "code_predictor_codes" not in payload
 
 
 def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
@@ -247,11 +247,33 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
     assert payload["ids"]["all"] == [151644, 872, 3]
     assert payload["embed"]["prefill"].device.type == "cpu"
     assert payload["hidden_states"]["output"].device.type == "cpu"
-    assert payload["next_stage_prompt_len"] > 0
     assert payload["embed"]["prefill"].shape[0] == 2
     assert payload["hidden_states"]["output"].shape[0] == 2
 
 
+def test_thinker2talker_token_only_preserves_voice_metadata() -> None:
+    source_outputs = [
+        SimpleNamespace(
+            request_id="req-1",
+            prompt_token_ids=[1, 2],
+            outputs=[SimpleNamespace(cumulative_token_ids=[3])],
+        )
+    ]
+    prompt = {
+        "additional_information": {
+            "speaker": ["ethan"],
+            "language": ["English"],
+        }
+    }
+
+    [talker_prompt] = q3.thinker2talker_token_only(source_outputs, prompt)
+
+    assert talker_prompt["additional_information"] == {
+        "speaker": ["ethan"],
+        "language": ["English"],
+    }
+
+
 def test_accumulator_replaces_keys_in_replace_set() -> None:
     """REPLACE-key semantics: subsequent emissions of the same key replace, not append."""
     from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
index 82f7da711..bc6239e29 100644
--- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
+++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py
@@ -12,9 +12,9 @@ from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
     max_ic_for_chunk_size,
 )
 from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
-    talker2code2wav,
     talker2code2wav_async_chunk,
     talker2code2wav_full_payload,
+    talker2code2wav_token_only,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -368,7 +368,15 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     assert rid in tm.request_payload
 
 
-def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
+def test_non_async_token_only_sizes_placeholder_for_ref_and_audio_frames():
+    """``talker2code2wav_token_only`` only allocates placeholder prompt slots.
+
+    After the connector refactor, actual codec flattening (ref prepend +
+    codebook-major layout) is performed by ``talker2code2wav_full_payload`` on
+    the worker data plane.  The orchestrator hook still derives
+    ``left_context_size`` from stage-0 multimodal_output so Code2Wav can trim
+    reference frames once the connector payload arrives.
+    """
     ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
     audio_codes = torch.tensor(
         [
@@ -387,12 +395,38 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav(stage.engine_outputs)
+    prompts = talker2code2wav_token_only(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]
     assert prompt["additional_information"] == {"meta": {"left_context_size": 2}}
-    assert prompt["prompt_token_ids"] == [
+    # 2 ref frames + 2 valid audio frames (zero row filtered), 4 quantizers.
+    assert prompt["prompt_token_ids"] == [0] * (_Q * (2 + 2))
+
+
+def test_full_payload_prepends_ref_code_and_flattens_codebook_major():
+    """Worker producer is authoritative for ref prepend + codec flatten."""
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+    audio_codes = torch.tensor(
+        [
+            [0, 0, 0, 0],
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+        ],
+        dtype=torch.long,
+    )
+    pooling_output = {
+        "codes.audio": audio_codes,
+        "codes.ref": ref_code,
+        "meta.ref_code_len": 2,
+    }
+    request = SimpleNamespace(request_id="r", output_token_ids=list(range(3)))
+
+    payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request)
+
+    assert payload is not None
+    assert payload["meta"]["left_context_size"] == 2
+    assert payload["codes"]["audio"].tolist() == [
         9,
         8,
         1,
@@ -433,7 +467,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav(stage.engine_outputs)
+    prompts = talker2code2wav_token_only(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]
diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py
index 3fd5675ce..dc248f211 100644
--- a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py
+++ b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py
@@ -8,61 +8,12 @@ import pytest
 import torch
 
 from vllm_omni.model_executor.stage_input_processors.voxtral_tts import (
-    generator2tokenizer,
     generator2tokenizer_async_chunk,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-# ---- Helpers for generator2tokenizer (non-streaming) ----
-
-
-def _make_stage(engine_outputs):
-    return SimpleNamespace(engine_outputs=engine_outputs)
-
-
-def _make_output(audio_tensors):
-    """Build a single generator output with multimodal_output codes.audio."""
-    return SimpleNamespace(
-        outputs=[SimpleNamespace(multimodal_output={"codes": {"audio": audio_tensors}})],
-    )
-
-
-# ---- Tests for generator2tokenizer ----
-
-
-class TestGenerator2Tokenizer:
-    def test_single_output_single_chunk(self):
-        """Single output with one audio tensor produces correct prompt_token_ids."""
-        audio = torch.tensor([10, 20, 30, 40])
-        stage = _make_stage([_make_output([audio])])
-
-        result = generator2tokenizer(stage.engine_outputs)
-
-        assert len(result) == 1
-        assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
-        assert result[0]["multi_modal_data"] is None
-
-    def test_2d_audio_tensors_flattened(self):
-        """2D audio tensors (e.g., multi-codebook frames) are flattened."""
-        audio = torch.tensor([[1, 2, 3], [4, 5, 6]])  # (2, 3)
-        stage = _make_stage([_make_output([audio])])
-
-        result = generator2tokenizer(stage.engine_outputs)
-
-        assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
-
-    def test_no_outputs_yet_raises(self):
-        """Stage with engine_outputs=None raises RuntimeError."""
-        stage = _make_stage(engine_outputs=None)
-        with pytest.raises(TypeError):
-            generator2tokenizer(stage.engine_outputs)
-
-
-# ---- Helpers for generator2tokenizer_async_chunk (streaming) ----
-
-
 def _req(external_req_id: str, *, finished: bool):
     return SimpleNamespace(
         external_req_id=external_req_id,
diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py
index f3a09346c..76ef84e7e 100644
--- a/tests/model_extras/test_model_extras.py
+++ b/tests/model_extras/test_model_extras.py
@@ -60,8 +60,9 @@ def test_sensenova_extra_registry_declares_request_and_response_params() -> None
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_cosmos3_extra_registry_declares_request_and_response_params() -> None:
-    assert get_extra_body_params("Cosmos3OmniDiffusersPipeline") == frozenset(
+@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
+def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_name: str) -> None:
+    assert get_extra_body_params(pipeline_name) == frozenset(
         {
             "flow_shift",
             "max_sequence_length",
@@ -98,7 +99,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params() -> None:
             "session_id",
         }
     )
-    assert get_extra_output_params("Cosmos3OmniDiffusersPipeline") == frozenset(
+    assert get_extra_output_params(pipeline_name) == frozenset(
         {
             "action",
             "raw_action_dim",
@@ -106,7 +107,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params() -> None:
             "action_mode",
         }
     )
-    assert should_init_extra_args_for_non_diffusion_stages("Cosmos3OmniDiffusersPipeline") is False
+    assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False
 
 
 @pytest.mark.core_model
@@ -128,9 +129,106 @@ def test_magi_human_extra_registry_declares_request_and_response_params() -> Non
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:
+def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None:
+    assert get_extra_body_params("MingImagePipeline") == frozenset(
+        {
+            "height",
+            "width",
+            "steps",
+            "cfg",
+            "seed",
+            "byte5_text",
+            "negative_prompt",
+        }
+    )
+    assert get_extra_output_params("MingImagePipeline") == frozenset()
+    assert should_init_extra_args_for_non_diffusion_stages("MingImagePipeline") is True
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_ming_flash_omni_declared_extra_args_route_into_sampling_params() -> None:
+    params = OmniDiffusionSamplingParams()
+    declared = get_extra_body_params("MingImagePipeline")
+    apply_declared_extra_args(
+        params,
+        declared,
+        {
+            "steps": 6,
+            "cfg": 1.5,
+            "byte5_text": ["理解与生成统一"],
+            "negative_prompt": "ugly, blurry",
+            "unknown": "ignored",
+        },
+    )
+    assert params.extra_args == {
+        "steps": 6,
+        "cfg": 1.5,
+        "byte5_text": ["理解与生成统一"],
+        "negative_prompt": "ugly, blurry",
+    }
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_ming_flash_omni_text_to_image_prompt_builder() -> None:
     assert build_text_to_image_prompt(
-        "Cosmos3OmniDiffusersPipeline",
+        "MingImagePipeline",
+        prompt="Please draw a cute cat.",
+        negative_prompt="ugly, blurry",
+        height=512,
+        width=768,
+    ) == {
+        "prompt": "Please draw a cute cat.",
+        "modalities": ["image"],
+        "mm_processor_kwargs": {
+            "modalities": ["image"],
+            "target_h": 512,
+            "target_w": 768,
+        },
+        "negative_prompt": "ugly, blurry",
+    }
+    # target_h/w are omitted when height/width are not supplied so the
+    # pipeline's 1024x1024 default still applies.
+    assert build_text_to_image_prompt(
+        "MingImagePipeline",
+        prompt="Draw a poster.",
+        negative_prompt=None,
+    ) == {
+        "prompt": "Draw a poster.",
+        "modalities": ["image"],
+        "mm_processor_kwargs": {"modalities": ["image"]},
+    }
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_ming_flash_omni_image_to_image_prompt_builder() -> None:
+    dummy_image = Image.new("RGB", (64, 64))
+    result = build_image_to_image_prompt(
+        "MingImagePipeline",
+        prompt="Change the background to a sandy beach at sunset.",
+        negative_prompt=None,
+        input_image=dummy_image,
+        height=256,
+        width=256,
+    )
+    assert result["modalities"] == ["img2img"]
+    assert result["multi_modal_data"]["img2img"] is dummy_image
+    assert result["mm_processor_kwargs"] == {
+        "modalities": ["img2img"],
+        "target_h": 256,
+        "target_w": 256,
+    }
+    assert "negative_prompt" not in result
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
+def test_cosmos3_text_to_image_prompt_builder_selects_image_modality(pipeline_name: str) -> None:
+    assert build_text_to_image_prompt(
+        pipeline_name,
         prompt="a red sports car at golden hour",
         negative_prompt="blurry, distorted",
         height=1024,
@@ -141,7 +239,7 @@ def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:
         "negative_prompt": "blurry, distorted",
     }
     assert build_text_to_image_prompt(
-        "Cosmos3OmniDiffusersPipeline",
+        pipeline_name,
         prompt="a red sports car",
         negative_prompt=None,
     ) == {"prompt": "a red sports car", "modalities": ["image"]}
diff --git a/tests/model_tests/__init__.py b/tests/model_tests/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/model_tests/conftest.py b/tests/model_tests/conftest.py
new file mode 100644
index 000000000..ab71a02cd
--- /dev/null
+++ b/tests/model_tests/conftest.py
@@ -0,0 +1,41 @@
+from shutil import rmtree
+
+import pytest
+
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+
+
+@pytest.fixture(scope="session")
+def tiny_model_paths(request, run_level):
+    """Build or download the tiny models for the selected tests.
+
+    At core_model level, builds tiny models via the builder function.
+    At advanced_model / full_model level, uses the real HF model path.
+
+    NOTE: this is session scoped to avoid churn in tiny model creation,
+    but will ensure all the tiny models you need are created for the selected tests
+    before it starts to execute them."""
+    model_paths = {}
+    built_paths = []
+    print("Initializing tiny models...")
+    for item in request.session.items:
+        if not hasattr(item, "callspec"):
+            continue
+        model_name = item.callspec.params.get("model_name")
+        if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS:
+            continue
+        if model_name not in model_paths:
+            settings = DIFFUSION_TEST_SETTINGS[model_name]
+            if run_level == "core_model":
+                print(f"Calling tiny model builder for: {model_name}")
+                path = settings.builder()
+                built_paths.append(path)
+            else:
+                print(f"Run level is {run_level}; {model_name} will use the full model")
+                path = settings.model
+            model_paths[model_name] = path
+
+    yield model_paths
+    for path in built_paths:
+        print(f"Removing tiny model: {path}")
+        rmtree(path, ignore_errors=True)
diff --git a/tests/model_tests/diffusion/__init__.py b/tests/model_tests/diffusion/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/model_tests/diffusion/case_filtering.py b/tests/model_tests/diffusion/case_filtering.py
new file mode 100644
index 000000000..a4fa5dc13
--- /dev/null
+++ b/tests/model_tests/diffusion/case_filtering.py
@@ -0,0 +1,104 @@
+"""
+Analogous to: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/vlm_utils/case_filtering.py
+
+This follows a similar pattern so that we can use `get_parametrized_options` to expand our test
+configurations out, similar to the way in which vLLM expands supported input type tests. For now,
+this mostly just means parametrizing over the test groups, which may be unique per model, but
+doing so in this way lets us cleanly define test groups that may be heterogeneous + avoid having
+excessive skip marks for features that may not be supported in every pipeline.
+"""
+
+# Engine initialization is currently quite expensive. For efficiency,
+# when we run the tests, we should try to collapse compatible accelerations,
+# because even with tiny models, the overhead of spinning up a new instance
+# stacks up fast.
+import itertools
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.model_tests.diffusion.config_types import DiffusionAccs, DiffusionModelTestOpts, get_required_device_count
+from vllm_omni.platforms import current_omni_platform
+
+# These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs),
+# but are intended to be smoke tests that we can properly validate without real model
+# weights in the CI.
+#
+# As such, no DiffusionAccs may require more than 4 GPUs, since we would otherwise skip
+# them in the CI.
+MAX_CI_DEVICES = 4
+
+
+def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> list:
+    """Build the full set of pytest marks for a test group.
+
+    For now, single device groups default to core model, and multi device groups
+    default to advanced_model marks. We also validate and throw if a test group exceeds
+    the max number of GPUs for the CI to prevent accidentally adding configurations that
+    would be skipped in the CI."""
+    marks = list(model_marks) if model_marks is not None else []
+
+    required_devices = get_required_device_count(test_group)
+    if required_devices > MAX_CI_DEVICES:
+        raise ValueError(
+            f"Test group {test_group} for {model_name} requires {required_devices} devices, "
+            f"but the max CI device count is {MAX_CI_DEVICES}. "
+        )
+
+    marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=required_devices))
+    # hardware_marks only adds skipif for num_cards > 1, so we currently handle the single-device case
+    # directly here. This should probably be handled in a more common way later on.
+    assert current_omni_platform is not None and current_omni_platform.device_count is not None
+    device_count = current_omni_platform.device_count()
+    if current_omni_platform.is_cuda() and device_count < required_devices:
+        marks.append(pytest.mark.skip(reason=f"Need {required_devices} devices, got {device_count}"))
+    if required_devices > 1:
+        marks.append(pytest.mark.full_model)
+    else:
+        marks.append(pytest.mark.core_model)
+    return marks
+
+
+def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts, online: bool):
+    """Given a model & its corresponding test options, build the list of pytest params
+    to be run for this model. The base case (no accelerations) is always included. Extra
+    test groups are always appended for offline tests, but are only added to the online
+    tests if online_base_only is overridden to False to avoid redundant testing.
+    """
+    test_groups: list[list[DiffusionAccs] | None] = [None]
+    if not (online and test_info.online_base_only) and test_info.extra_test_groups:
+        test_groups.extend(test_info.extra_test_groups)
+
+    return [
+        pytest.param(
+            model_name,
+            test_group,
+            test_info.supported_tasks,
+            test_info.check_multi_output and test_group is None,
+            test_info.check_determinism and test_group is None,
+            id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name,
+            marks=get_test_group_marks(model_name, test_group, test_info.marks),
+        )
+        for test_group in test_groups
+    ]
+
+
+def get_parametrized_options(
+    test_settings: dict[str, DiffusionModelTestOpts],
+    online: bool = False,
+):
+    """Converts all the DiffusionModelTestOpts into an expanded list of parameters
+    based on which accelerations are available.
+
+    When online=True and a model has online_base_only=True, only the base case
+    (no accelerations) is included for that model's parametrization to minimize
+    redundancy between offline / online tests.
+    """
+    # Get a list per model type, where each entry contains a tuple of all of
+    # that model type's cases, then flatten them into the top level so that
+    # we can consume them in one mark.parametrize call.
+    parametrization = [
+        get_model_parametrization(model_name, test_info, online=online)
+        for model_name, test_info in test_settings.items()
+    ]
+    return list(itertools.chain(*parametrization))
diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py
new file mode 100644
index 000000000..4120900fe
--- /dev/null
+++ b/tests/model_tests/diffusion/config_types.py
@@ -0,0 +1,187 @@
+"""
+Common definitions for controlling what tests run where.
+"""
+
+from collections.abc import Callable
+from enum import StrEnum, auto
+from typing import NamedTuple, TypeAlias
+
+from pytest import MarkDecorator
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+
+# All builder funcs take no params and return a path
+TinyDiffusionBuilder: TypeAlias = Callable[[], str]
+
+
+class DiffusionAccs(StrEnum):
+    """Supported acceleration types / test settings for Diffusion Models."""
+
+    HSDP = auto()
+    TEA_CACHE = auto()
+    CACHE_DIT = auto()
+    SEQUENCE_PARALLEL = auto()
+    CFG_PARALLEL = auto()
+    TENSOR_PARALLEL = auto()
+    CPU_OFFLOAD = auto()
+    LAYERWISE_OFFLOAD = auto()
+    VAE_PATCH_PARALLEL = auto()
+
+
+class DiffusionTasks(StrEnum):
+    """Supported tasks for Diffusion Models."""
+
+    TEXT_TO_IMAGE = auto()
+    IMAGE_TO_IMAGE = auto()
+    # Text to video, text to audio, etc should be added here as needed
+
+
+class DiffusionModelTestOpts(NamedTuple):
+    """Configuration for one Diffusion model's tests."""
+
+    # HF model name for real-weight tests (advanced_model / full_model level).
+    # For now, whether we use the real weights vs tiny weights in the common tests
+    # depends on the run level.
+    model: str
+
+    # Creates a tiny model for the given architecture. We should always use tiny
+    # model weights for tests that do not require us to check the model quality.
+    builder: TinyDiffusionBuilder
+
+    # Actual tasks which controls the tests actually run
+    supported_tasks: list[DiffusionTasks]
+
+    # Additional acceleration groups to run beyond the base case (no acceleration).
+    # The base case is always run for every model in the test settings. None means
+    # we only run the base case.
+    extra_test_groups: list[list[DiffusionAccs]] | None = None
+
+    # Pytest Marks for this model. This may be useful for selecting which models
+    # we want to run where, similar to the way vLLM's multimodal tests mark some
+    # as core models to always run in the CI.
+    # Example: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/test_common.py#L131
+    marks: list[MarkDecorator] | None = None
+
+    # When True (default), online tests only run the base case (no accelerations).
+    # When False, online tests run all test_groups, same as offline. This should be
+    # True unless there is a good reason for it not to be, because the execution of
+    # the acceleration should be the same on both codepaths, and CLI parsing etc should
+    # be tested by tests adding the acceleration, and not per model.
+    online_base_only: bool = True
+
+    # Additional checks to run for the base case.
+    check_multi_output: bool = True  # Runs multiple generations in one request
+    check_determinism: bool = True  # Runs 2 generations with the same seed and check determinism
+
+
+### Mappings & utils for building offline Omni() instances given a list of enabled accelerations
+ACC_OMNI_KWARGS = {
+    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True},
+    DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True},
+    DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True},
+    DiffusionAccs.CACHE_DIT: {"cache_backend": "cache_dit"},
+    DiffusionAccs.TEA_CACHE: {"cache_backend": "tea_cache"},
+}
+
+ACC_PARALLEL_KWARGS = {
+    DiffusionAccs.HSDP: {"use_hsdp": True, "hsdp_shard_size": 2},
+    DiffusionAccs.TENSOR_PARALLEL: {"tensor_parallel_size": 2},
+    DiffusionAccs.CFG_PARALLEL: {"cfg_parallel_size": 2},
+    DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_patch_parallel_size": 2},
+    # For SP, we don't run ring here to conserve devices, since it is easy
+    # to blow up the number of needed devices fast. Compatibility for ring
+    # and ulysses together should be tested generically outside of these tests.
+    DiffusionAccs.SEQUENCE_PARALLEL: {"ulysses_degree": 2},
+}
+
+ACC_DEVICE_COUNT_KEYS = {
+    DiffusionAccs.TENSOR_PARALLEL: "tensor_parallel_size",
+    DiffusionAccs.CFG_PARALLEL: "cfg_parallel_size",
+    DiffusionAccs.SEQUENCE_PARALLEL: "ulysses_degree",
+    DiffusionAccs.VAE_PATCH_PARALLEL: "vae_patch_parallel_size",
+    DiffusionAccs.HSDP: "hsdp_shard_size",
+}
+
+### CLI args for launching an OmniServer subprocess per acceleration.
+# Follows the same pattern as e2e online serving tests, which hand-code
+# their server_args lists in OmniServerParams.
+ACC_SERVER_ARGS: dict[DiffusionAccs, list[str]] = {
+    DiffusionAccs.HSDP: ["--use-hsdp", "--hsdp-shard-size", "2"],
+    DiffusionAccs.TEA_CACHE: ["--cache-backend", "tea_cache"],
+    DiffusionAccs.CACHE_DIT: ["--cache-backend", "cache_dit"],
+    DiffusionAccs.SEQUENCE_PARALLEL: ["--usp", "2"],
+    DiffusionAccs.CFG_PARALLEL: ["--cfg-parallel-size", "2"],
+    DiffusionAccs.TENSOR_PARALLEL: ["--tensor-parallel-size", "2"],
+    DiffusionAccs.CPU_OFFLOAD: ["--enable-cpu-offload"],
+    DiffusionAccs.LAYERWISE_OFFLOAD: ["--enable-layerwise-offload"],
+    DiffusionAccs.VAE_PATCH_PARALLEL: ["--vae-use-tiling", "--vae-patch-parallel-size", "2"],
+}
+
+## TODO ^ These should be in the same object, it's getting messy.
+
+
+def get_required_device_count(accelerations: list[DiffusionAccs] | None) -> int:
+    """Compute the minimum number of devices needed for a set of accelerations.
+    The total is the product of all parallel dimensions (defaulting to 1).
+
+    If not enough devices are available for a test group's accelerations,
+    that test will be skipped."""
+    count = 1
+    if accelerations is None:
+        return count
+
+    for acc in accelerations:
+        key = ACC_DEVICE_COUNT_KEYS.get(acc)
+        if key is not None:
+            count *= ACC_PARALLEL_KWARGS[acc][key]
+    return count
+
+
+def build_parallel_config_from_diff_accelerations(accelerations: list[DiffusionAccs]) -> DiffusionParallelConfig | None:
+    """Given a list of accelerations pertaining to the current test group,
+    build the parallel config needed for the Omni() object (if any)."""
+    config_kwargs = {}
+    for acc in accelerations:
+        update_dict = ACC_PARALLEL_KWARGS.get(acc, {})
+        config_kwargs.update(update_dict)
+    if config_kwargs:
+        return DiffusionParallelConfig(**config_kwargs)
+    return None
+
+
+### Offline Omni() object builder
+def build_omni_from_diff_accelerations(accelerations: list[DiffusionAccs] | None, **kwargs) -> Omni:
+    """Given one or more acceleration types, build the corresponding Omni() object."""
+    # Coerce to a list and build the parallel config, since that depends on the accelerations
+    if accelerations is None:
+        accelerations = []
+    parallel_config = build_parallel_config_from_diff_accelerations(accelerations)
+
+    # Then add anything else that's a top-level kwarg
+    acc_kwargs = {}
+    if parallel_config is not None:
+        acc_kwargs["parallel_config"] = parallel_config
+    for acc in accelerations:
+        update_dict = ACC_OMNI_KWARGS.get(acc, {})
+        acc_kwargs.update(update_dict)
+
+    # Keys passed through should mostly be things like enforce_eager;
+    # if there's overlap, it's probably due to a misconfiguration
+    shared_keys = acc_kwargs.keys() & kwargs.keys()
+    if shared_keys:
+        raise ValueError(f"Explicit Omni kwargs and inferred Omni kwargs for accelerations overlap: {shared_keys}")
+    omni_kwargs = {**acc_kwargs, **kwargs}
+    return Omni(**omni_kwargs)
+
+
+### Online server flag builder
+def build_server_args_from_diff_accelerations(accelerations: list[DiffusionAccs] | None) -> list[str]:
+    """Given one or more acceleration types, build the corresponding CLI args
+    for launching an OmniServer subprocess."""
+    if accelerations is None:
+        return []
+    args = []
+    for acc in accelerations:
+        args.extend(ACC_SERVER_ARGS.get(acc, []))
+    return args
diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py
new file mode 100644
index 000000000..9bee71f9b
--- /dev/null
+++ b/tests/model_tests/diffusion/diff_model_builders.py
@@ -0,0 +1,63 @@
+import os
+import tempfile
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2
+from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
+from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM
+
+TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models")
+
+
+def _get_tiny_model_path(name: str) -> str:
+    path = os.path.join(TINY_MODEL_DIR, name)
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
+def tiny_flux2_klein_builder() -> str:
+    """Build a tiny Flux2Klein model."""
+    model_dir = _get_tiny_model_path("Flux2KleinPipeline")
+
+    pipe = Flux2KleinPipeline(
+        scheduler=FlowMatchEulerDiscreteScheduler(),
+        vae=AutoencoderKLFlux2(
+            down_block_types=("DownEncoderBlock2D",),
+            up_block_types=("UpDecoderBlock2D",),
+            block_out_channels=(32,),
+            layers_per_block=1,
+            latent_channels=16,
+            norm_num_groups=16,
+        ),
+        # NOTE: For now we need 28 layers because of hardcoded stuff in the model :(
+        text_encoder=Qwen3ForCausalLM(
+            Qwen3Config(
+                hidden_size=32,
+                intermediate_size=64,
+                num_hidden_layers=28,
+                num_attention_heads=2,
+                num_key_value_heads=2,
+                head_dim=16,
+                vocab_size=151936,
+                max_position_embeddings=512,
+            )
+        ),
+        tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"),
+        # NOTE: For now we need at least 2 layers for the transformer
+        # due to hardcoded hacks in CacheDiT for Flux2Klein specifically.
+        transformer=Flux2Transformer2DModel(
+            in_channels=64,
+            num_layers=2,
+            num_single_layers=2,
+            attention_head_dim=16,
+            num_attention_heads=2,
+            joint_attention_dim=96,
+            timestep_guidance_channels=32,
+            axes_dims_rope=(4, 4, 4, 4),
+        ),
+    )
+    # Need dtypes to be consistent; for now we just put it on bfloat16
+    pipe.to(torch.bfloat16).save_pretrained(model_dir)
+    return model_dir
diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py
new file mode 100644
index 000000000..3f44afcae
--- /dev/null
+++ b/tests/model_tests/diffusion/model_settings.py
@@ -0,0 +1,42 @@
+from tests.model_tests.diffusion import diff_model_builders
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionModelTestOpts,
+    DiffusionTasks,
+)
+
+# This object defines the (tiny) model configurations for common tests.
+#
+# TO ADD A NEW MODEL:
+# Map the pipeline name (i.e., the name of the pipeline in _DIFFUSION_MODELS)
+# to a DiffusionModelTestOpts and define the supported tasks. A base case with no
+# accelerations is always run. Use extra_test_groups to define additional
+# acceleration configurations to run a basic smoke test for (i.e., simple inference).
+#
+# The model will only be loaded once per test group, and will execute each of its
+# supported tasks as a pytest subtest. The base case also runs additional checks
+# that are independent of accelerations, e.g., for determinism and compatibility with
+# multi-output, unless explicitly disabled in the test settings.
+#
+# TO RUN A SUBSET OF THE TESTS:
+# These tests should be fast, but if you'd like to run a subset of the tests, you can do so
+# with `-k`; the IDs of the tests are the name of the pipeline joined by +.
+#
+# Example:
+# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline]
+#   ^ Runs only the base (no accelerations) case for Flux2KleinPipeline
+#
+# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline
+#   ^ Runs all test groups for Flux2KleinPipeline only
+DIFFUSION_TEST_SETTINGS = {
+    "Flux2KleinPipeline": DiffusionModelTestOpts(
+        model="black-forest-labs/FLUX.2-klein-4B",
+        builder=diff_model_builders.tiny_flux2_klein_builder,
+        supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE],
+        extra_test_groups=[
+            [DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE],
+            [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD],
+            [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD],
+        ],
+    )
+}
diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py
new file mode 100644
index 000000000..6b9c7ff99
--- /dev/null
+++ b/tests/model_tests/diffusion/task_runners.py
@@ -0,0 +1,172 @@
+"""
+Definitions running and validating individual tasks, e.g., text to image,
+image to image, and so on. These are called by the core test runner.
+"""
+
+import base64
+import io
+from dataclasses import replace
+
+import numpy as np
+from PIL import Image
+
+from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+PROMPT = "Dummy prompt"
+IMAGE_DIMS = (512, 512)
+HEIGHT, WIDTH = IMAGE_DIMS
+INPUT_IMAGE = Image.new("RGB", IMAGE_DIMS)
+
+# Offline sampling params
+IMAGE_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams(
+    num_inference_steps=4,
+    height=HEIGHT,
+    width=WIDTH,
+    seed=42,
+)
+
+# Online extra_body for diffusion requests
+IMAGE_GEN_EXTRA_BODY = {
+    "height": HEIGHT,
+    "width": WIDTH,
+    "num_inference_steps": 4,
+    "seed": 42,
+}
+
+
+### Shared validation
+def _validate_images(images: list[Image.Image], expected_n: int = 1):
+    """Given a set of images, ensure we got the expected count of images,
+    and that all provided images match the expected dimensions."""
+    assert len(images) == expected_n
+    for img in images:
+        assert isinstance(img, Image.Image)
+        assert img.size == IMAGE_DIMS
+    return images
+
+
+def _validate_image_gen_determinism(images_a: list[Image.Image], images_b: list[Image.Image]):
+    """Ensure that two image sets are valid and that the results match."""
+    _validate_images(images_a)
+    _validate_images(images_b)
+    assert np.array_equal(np.array(images_a[0]), np.array(images_b[0]))
+
+
+### Output extractor utils for offline / online paths respectively
+def _get_offline_images(outputs: list[OmniRequestOutput]) -> list[Image.Image]:
+    """Extract the images from an Omni .generate() call."""
+    assert len(outputs) == 1
+    return outputs[0].images
+
+
+def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]:
+    """Extract the images from a server response."""
+    assert len(responses) == 1
+    images = responses[0].images
+    assert images is not None
+    return images
+
+
+### Offline helpers
+def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS):
+    return omni.generate({"prompt": PROMPT}, params)
+
+
+def _run_offline_i2i(omni: Omni):
+    return omni.generate(
+        {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}},
+        IMAGE_GEN_SAMPLING_PARAMS,
+    )
+
+
+### Online helpers
+def _build_online_image_data_url() -> str:
+    """Get a valid base 64 encoded data URL corresponding to an image."""
+    buf = io.BytesIO()
+    INPUT_IMAGE.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/png;base64,{b64}"
+
+
+def _run_online_t2i(
+    server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None
+) -> list[DiffusionResponse]:
+    """Run a text to image request through the server."""
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+    request_config = {
+        "model": server.model,
+        "messages": messages,
+        "extra_body": extra_body or IMAGE_GEN_EXTRA_BODY,
+    }
+    return client.send_diffusion_request(request_config)
+
+
+def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]:
+    """Run an image to image request through the server."""
+    image_data_url = _build_online_image_data_url()
+    messages = dummy_messages_from_mix_data(
+        content_text=PROMPT,
+        image_data_url=image_data_url,
+    )
+    request_config = {
+        "model": server.model,
+        "messages": messages,
+        "extra_body": IMAGE_GEN_EXTRA_BODY,
+    }
+    return client.send_diffusion_request(request_config)
+
+
+### Offline task runners
+def run_and_validate_text_to_image_request(omni: Omni):
+    """Run and validate a text to image request."""
+    _validate_images(_get_offline_images(_run_offline_t2i(omni)))
+
+
+def run_and_validate_image_to_image_request(omni: Omni):
+    """Run and validate an image to image request."""
+    _validate_images(_get_offline_images(_run_offline_i2i(omni)))
+
+
+def run_and_validate_text_to_image_determinism(omni: Omni):
+    """Checks for determinism; for now we just keep this for TTI."""
+    _validate_image_gen_determinism(
+        _get_offline_images(_run_offline_t2i(omni)),
+        _get_offline_images(_run_offline_t2i(omni)),
+    )
+
+
+def run_and_validate_text_to_image_multi_output(omni: Omni):
+    """Checks for multi-output; for now we just keep this for TTI."""
+    params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2)
+    _validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2)
+
+
+### Online task runners
+def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler):
+    """Run and validate a text to image request through the server."""
+    _validate_images(_get_online_images(_run_online_t2i(server, client)))
+
+
+def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler):
+    """Run and validate an image to image request through the server."""
+    _validate_images(_get_online_images(_run_online_i2i(server, client)))
+
+
+def run_and_validate_online_text_to_image_determinism(server: OmniServer, client: OpenAIClientHandler):
+    """Checks for determinism through the server; for now we just keep this for TTI."""
+    _validate_image_gen_determinism(
+        _get_online_images(_run_online_t2i(server, client)),
+        _get_online_images(_run_online_t2i(server, client)),
+    )
+
+
+def run_and_validate_online_text_to_image_multi_output(server: OmniServer, client: OpenAIClientHandler):
+    """Checks for multi-output through the server; for now we just keep this for TTI."""
+    extra_body = {**IMAGE_GEN_EXTRA_BODY, "num_outputs_per_prompt": 2}
+    _validate_images(
+        _get_online_images(_run_online_t2i(server, client, extra_body=extra_body)),
+        expected_n=2,
+    )
diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py
new file mode 100644
index 000000000..c95644143
--- /dev/null
+++ b/tests/model_tests/diffusion/test_alignment.py
@@ -0,0 +1,100 @@
+"""
+Tests for alignment between _DIFFUSION_MODELS and DIFFUSION_TEST_SETTINGS; if
+tests in this file are failing, you are probably adding a new model, and need
+to account for it in DIFFUSION_TEST_SETTINGS
+"""
+
+import pytest
+
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+# These diffusion models currently do not have any added (tiny) model tests.
+# In general, pipelines should only be in the list below if they have not been
+# migrated yet or there is a compelling reason that they should not be (which for
+# now may just be CI pressure etc due to the volume of models and the cost of
+# initializing the server).
+#
+# The tests here validate that all keys in the diffusion registry are accounted for;
+# if you are adding a new model and see the tests below fail, please follow the pattern
+# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS.
+EXCLUDED_MODELS = [
+    "QwenImagePipeline",
+    "QwenImageEditPipeline",
+    "QwenImageEditPlusPipeline",
+    "QwenImageLayeredPipeline",
+    "GlmImagePipeline",
+    "ZImagePipeline",
+    "OvisImagePipeline",
+    "WanPipeline",
+    "WanVACEPipeline",
+    "LTX2Pipeline",
+    "LTX2ImageToVideoPipeline",
+    "LTX2TwoStagesPipeline",
+    "LTX2ImageToVideoTwoStagesPipeline",
+    "LTX2T2VDMD2Pipeline",
+    "LTX2I2VDMD2Pipeline",
+    "LTX23Pipeline",
+    "LTX23ImageToVideoPipeline",
+    "StableAudioPipeline",
+    "WanImageToVideoPipeline",
+    "WanS2VPipeline",
+    "WanT2VDMD2Pipeline",
+    "WanI2VDMD2Pipeline",
+    "LongCatImagePipeline",
+    "BagelPipeline",
+    "LancePipeline",
+    "MingImagePipeline",
+    "InternVLAA1Pipeline",
+    "LongCatImageEditPipeline",
+    "StableDiffusion3Pipeline",
+    "FluxKontextPipeline",
+    "HunyuanImage3ForCausalMM",
+    "ErnieImagePipeline",
+    "NextStep11Pipeline",
+    "FluxPipeline",
+    "FluxDMD2Pipeline",
+    "Krea2Pipeline",
+    "QwenImageDMD2Pipeline",
+    "OmniGen2Pipeline",
+    "HeliosPipeline",
+    "HeliosPyramidPipeline",
+    "Flux2Pipeline",
+    "DreamIDOmniPipeline",
+    "SenseNovaU1Pipeline",
+    "AudioXPipeline",
+    "HunyuanVideo15Pipeline",
+    "HunyuanVideo15ImageToVideoPipeline",
+    "MagiHumanPipeline",
+    "OmniVoicePipeline",
+    "OmniVoice",
+    "Cosmos3OmniDiffusersPipeline",
+    "Cosmos3OmniPipeline",
+    "DiffusersAdapterPipeline",
+    "HiDreamImagePipeline",
+    "DreamZeroPipeline",
+    "StableDiffusionXLPipeline",
+    "Gr00tN1d7Pipeline",
+    "SoulXSingerPipeline",
+    "SoulXSingerSVCPipeline",
+]
+
+
+def test_all_non_excluded_pipelines_have_tests():
+    """Ensure that every pipeline in the diffusion registry either has a test
+    configuration in DIFFUSION_TEST_SETTINGS, or is explicitly excluded for a
+    known reason."""
+    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
+    missing_tests = non_excluded_diff_pipes - DIFFUSION_TEST_SETTINGS.keys()
+    assert len(missing_tests) == 0, f"Pipelines missing test settings: {missing_tests}"
+
+
+def test_no_excluded_models_have_test_settings():
+    """Ensure that no models in the exclude list have a test configuration in
+    DIFFUSION_TEST_SETTINGS (i.e., if they have test settings, we should not
+    exclude them)."""
+    non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS)
+    missing_pipes = DIFFUSION_TEST_SETTINGS.keys() - set(non_excluded_diff_pipes)
+    assert len(missing_pipes) == 0
diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py
new file mode 100644
index 000000000..f24ed3780
--- /dev/null
+++ b/tests/model_tests/diffusion/test_common_offline.py
@@ -0,0 +1,69 @@
+import pytest
+
+from tests.model_tests.diffusion.case_filtering import get_parametrized_options
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionTasks,
+    build_omni_from_diff_accelerations,
+)
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from tests.model_tests.diffusion.task_runners import (
+    run_and_validate_image_to_image_request,
+    run_and_validate_text_to_image_determinism,
+    run_and_validate_text_to_image_multi_output,
+    run_and_validate_text_to_image_request,
+)
+
+# NOTE : Hardware marks are added dynamically based on test requirements
+pytestmark = [pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
+    get_parametrized_options(DIFFUSION_TEST_SETTINGS),
+)
+def test_pipeline_on_supported_tasks(
+    model_name,
+    accelerations: list[DiffusionAccs] | None,
+    supported_tasks: list[DiffusionTasks],
+    check_multioutput: bool,
+    check_determinism: bool,
+    tiny_model_paths: dict[str, str],
+    subtests,
+):
+    """Run a smoke test on all of the pipelines supported tasks using a set of enabled accelerations."""
+    assert len(supported_tasks) > 0
+    # We initialize the Omni object before running the tasks, then run each task as a pytest subtask.
+    # This lets us init the model once, but display separate failures in pytest, and avoid halting the
+    # checks on other tasks if one fails.
+    #
+    # This allows us to have some degree of test isolation without the cost of redundant initialization,
+    # since starting the server can take 10+ seconds, even for tiny models.
+    #
+    # NOTE: Be sure to install pytest-subtests if you're running on pytest < 9
+    omni = build_omni_from_diff_accelerations(
+        accelerations=accelerations,
+        model=tiny_model_paths[model_name],
+        enforce_eager=True,
+    )
+    try:
+        for task_type in supported_tasks:
+            with subtests.test(msg=task_type):
+                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+                    run_and_validate_text_to_image_request(omni)
+                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+                    run_and_validate_image_to_image_request(omni)
+                else:
+                    raise ValueError(f"Task type {task_type} is not yet supported")
+
+        # NOTE: For now, we only check determinism + multi output for the base case,
+        # since checking it on every extra acceleration configuration is redundant
+        # (see case_filtering).
+        if check_determinism:
+            with subtests.test(msg="determinism"):
+                run_and_validate_text_to_image_determinism(omni)
+        if check_multioutput:
+            with subtests.test(msg="multi_output"):
+                run_and_validate_text_to_image_multi_output(omni)
+    finally:
+        omni.close()
diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py
new file mode 100644
index 000000000..e4c6be0d4
--- /dev/null
+++ b/tests/model_tests/diffusion/test_common_online.py
@@ -0,0 +1,75 @@
+"""
+Analogous to test_common_offline, but for server tests. Validates the full
+online serving stack (CLI arg parsing, subprocess, API routing, response
+encoding) using tiny models.
+"""
+
+import pytest
+
+from tests.helpers.runtime import OmniServer, OpenAIClientHandler
+from tests.model_tests.diffusion.case_filtering import get_parametrized_options
+from tests.model_tests.diffusion.config_types import (
+    DiffusionAccs,
+    DiffusionTasks,
+    build_server_args_from_diff_accelerations,
+)
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from tests.model_tests.diffusion.task_runners import (
+    run_and_validate_online_image_to_image_request,
+    run_and_validate_online_text_to_image_determinism,
+    run_and_validate_online_text_to_image_multi_output,
+    run_and_validate_online_text_to_image_request,
+)
+
+# NOTE : Hardware marks are added dynamically based on test requirements
+pytestmark = [pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    "model_name,accelerations,supported_tasks,check_multioutput,check_determinism",
+    get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True),
+)
+def test_online_on_supported_tasks(
+    model_name: str,
+    accelerations: list[DiffusionAccs] | None,
+    supported_tasks: list[DiffusionTasks],
+    check_multioutput: bool,
+    check_determinism: bool,
+    tiny_model_paths: dict[str, str],
+    run_level: str,
+    subtests,
+):
+    """Smoke test: start a tiny model server and run each supported task via the API."""
+    model_path = tiny_model_paths[model_name]
+    server_args = build_server_args_from_diff_accelerations(accelerations)
+    server_args.append("--enforce-eager")
+
+    with OmniServer(model_path, server_args) as server:
+        # TODO: We may want to revisit run_level validation here,
+        # because checks for things like image size etc should not
+        # depend on whether or not the weights are real or random
+        client = OpenAIClientHandler(
+            host=server.host,
+            port=server.port,
+            api_key="EMPTY",
+            run_level=run_level,
+            log_stats=server.log_stats,
+        )
+        for task_type in supported_tasks:
+            with subtests.test(msg=task_type):
+                if task_type == DiffusionTasks.TEXT_TO_IMAGE:
+                    run_and_validate_online_text_to_image_request(server, client)
+                elif task_type == DiffusionTasks.IMAGE_TO_IMAGE:
+                    run_and_validate_online_image_to_image_request(server, client)
+                else:
+                    raise ValueError(f"Task type {task_type} is not yet supported")
+
+        # NOTE: For now, we only check determinism + multi output for the base case,
+        # since checking it on every extra acceleration configuration is redundant
+        # (see case_filtering).
+        if check_determinism:
+            with subtests.test(msg="determinism"):
+                run_and_validate_online_text_to_image_determinism(server, client)
+        if check_multioutput:
+            with subtests.test(msg="multi_output"):
+                run_and_validate_online_text_to_image_multi_output(server, client)
diff --git a/tests/model_tests/diffusion/utils.py b/tests/model_tests/diffusion/utils.py
new file mode 100644
index 000000000..35577b68f
--- /dev/null
+++ b/tests/model_tests/diffusion/utils.py
@@ -0,0 +1,33 @@
+"""
+Utilities for resolving real models to their tiny model equivalents.
+"""
+
+import logging
+
+from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS
+from vllm_omni.diffusion.data import resolve_model_class_name
+
+logger = logging.getLogger(__name__)
+
+
+def resolve_tiny_model_path(model: str) -> str:
+    """Given a real model name/path, resolve it to a tiny model path.
+
+    Raises ValueError if the pipeline class cannot be determined (invalid
+    model). Returns the original model path if no tiny builder exists yet."""
+    pipeline_class = resolve_model_class_name(model)
+    if pipeline_class is None:
+        raise ValueError(
+            f"Cannot resolve pipeline class for model: {model}. The model path may be invalid or its config unreadable."
+        )
+
+    test_opts = DIFFUSION_TEST_SETTINGS.get(pipeline_class)
+    if test_opts is None:
+        logger.warning(
+            "No tiny model builder for pipeline %s (model: %s). Using original model.",
+            pipeline_class,
+            model,
+        )
+        return model
+
+    return test_opts.builder()
diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py
index 255b092c8..e138a6d63 100644
--- a/tests/platforms/npu/test_310p_patches.py
+++ b/tests/platforms/npu/test_310p_patches.py
@@ -67,10 +67,37 @@ def _install_qwen3_tts_patch_fakes(monkeypatch: pytest.MonkeyPatch):
     class FakeCodePredictorBaseModel(torch.nn.Module):
         pass
 
-    class FakeMimiEuclideanCodebook(torch.nn.Module):
-        @property
-        def embed(self):
-            return self._embed
+    class FakeProjection(torch.nn.Linear):
+        def __init__(self):
+            super().__init__(4, 4, bias=False)
+            self.call_shapes: list[tuple[int, ...]] = []
+
+        def forward(self, hidden_states):
+            self.call_shapes.append(tuple(hidden_states.shape))
+            return super().forward(hidden_states)
+
+    class FakeCodePredictorWrapper(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            del args, kwargs
+            super().__init__()
+            self.model = torch.nn.Module()
+            self.model.codec_embedding = torch.nn.ModuleList([torch.nn.Embedding(8, 4), torch.nn.Embedding(8, 4)])
+            self.model.linear = torch.nn.Linear(4, 4, bias=False)
+            self.lm_head = torch.nn.ModuleList([torch.nn.Linear(4, 8, bias=False), torch.nn.Linear(4, 8, bias=False)])
+            self.small_to_mtp_projection = FakeProjection()
+            with torch.no_grad():
+                self.small_to_mtp_projection.weight.copy_(torch.diag(torch.tensor([1.0, 2.0, 3.0, 4.0])))
+            self._wrapper_config = SimpleNamespace(use_parallel_embedding=False, sampling_mode="per_call")
+            self._projected_codec_embed_weight = None
+
+        def load_weights(self, weights):
+            del weights
+            return {"loaded"}
+
+    class FakeCode2WavBase(torch.nn.Module):
+        def __init__(self, *args, **kwargs):
+            del args, kwargs
+            super().__init__()
 
     class FakeEncoder:
         def __init__(self):
@@ -136,8 +163,21 @@ def _install_qwen3_tts_patch_fakes(monkeypatch: pytest.MonkeyPatch):
         CodePredictorAttention=FakeCodePredictorAttention,
         CodePredictorDecoderLayer=FakeCodePredictorDecoderLayer,
         CodePredictorBaseModel=FakeCodePredictorBaseModel,
+        CodePredictorWrapper=FakeCodePredictorWrapper,
         _rotate_half=lambda x: x,
     )
+    fake_qwen3_tts_code_predictor_vllm = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm",
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM=FakeCodePredictorWrapper,
+        Qwen3TTSTalkerCodePredictorModelVLLM=FakeCodePredictorBaseModel,
+        CodePredictorWrapper=FakeCodePredictorWrapper,
+    )
+    fake_qwen3_tts_code2wav = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav",
+        Qwen3TTSCode2Wav=FakeCode2WavBase,
+    )
     fake_prompt_builder = _install_fake_module(
         monkeypatch,
         "vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder",
@@ -150,22 +190,40 @@ def _install_qwen3_tts_patch_fakes(monkeypatch: pytest.MonkeyPatch):
         Qwen3TTSTalkerForConditionalGeneration=FakeTalkerBase,
         Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder,
     )
-    fake_modeling_mimi = _install_fake_module(
-        monkeypatch,
-        "transformers.models.mimi.modeling_mimi",
-        MimiEuclideanCodebook=FakeMimiEuclideanCodebook,
-    )
-    fake_mimi = _install_fake_module(
-        monkeypatch,
-        "transformers.models.mimi",
-        modeling_mimi=fake_modeling_mimi,
-    )
 
     _install_fake_module(monkeypatch, "vllm")
     _install_fake_module(monkeypatch, "vllm.multimodal")
     _install_fake_module(monkeypatch, "vllm.multimodal.audio", AudioResampler=FakeAudioResampler)
-    _install_fake_module(monkeypatch, "transformers")
-    _install_fake_module(monkeypatch, "transformers.models", mimi=fake_mimi)
+    _install_fake_module(monkeypatch, "torch_npu", npu_format_cast=lambda weight, _fmt: weight)
+    _install_fake_module(monkeypatch, "vllm_ascend")
+    _install_fake_module(monkeypatch, "vllm_ascend._310p")
+    _install_fake_module(monkeypatch, "vllm_ascend._310p.attention")
+    _install_fake_module(
+        monkeypatch,
+        "vllm_ascend._310p.attention.attention_mask",
+        AttentionMaskBuilder310=SimpleNamespace(
+            gen_causal_additive_mask=lambda max_seq, device: torch.zeros(
+                max_seq,
+                max_seq,
+                device=device,
+            )
+        ),
+    )
+    _install_fake_module(monkeypatch, "vllm_ascend.sample")
+    _install_fake_module(
+        monkeypatch,
+        "vllm_ascend.sample.sampler",
+        apply_top_k_top_p=lambda logits, **_kwargs: logits,
+        random_sample=lambda probs, _generators: probs.argmax(dim=-1, keepdim=True),
+    )
+    _install_fake_module(
+        monkeypatch,
+        "vllm_ascend.utils",
+        ACL_FORMAT_FRACTAL_NZ=29,
+        aligned_16=lambda tensor: tensor,
+        maybe_trans_nz=lambda weight: weight,
+        nd_to_nz_2d=lambda tensor: tensor,
+    )
     _install_fake_module(monkeypatch, "vllm_omni")
     _install_fake_module(monkeypatch, "vllm_omni.model_executor")
     _install_fake_module(monkeypatch, "vllm_omni.model_executor.models")
@@ -178,9 +236,17 @@ def _install_qwen3_tts_patch_fakes(monkeypatch: pytest.MonkeyPatch):
         monkeypatch,
         "vllm_omni.model_executor.models.qwen3_tts",
         prompt_embeds_builder=fake_prompt_builder,
+        qwen3_tts_code2wav=fake_qwen3_tts_code2wav,
+        qwen3_tts_code_predictor_vllm=fake_qwen3_tts_code_predictor_vllm,
         qwen3_tts_talker=fake_talker,
     )
-    return fake_qwen3_code_predictor, fake_prompt_builder, fake_talker
+    return (
+        fake_qwen3_code_predictor,
+        fake_qwen3_tts_code_predictor_vllm,
+        fake_qwen3_tts_code2wav,
+        fake_prompt_builder,
+        fake_talker,
+    )
 
 
 def _load_qwen3_tts_patch(monkeypatch: pytest.MonkeyPatch):
@@ -193,7 +259,7 @@ def _load_qwen3_tts_patch(monkeypatch: pytest.MonkeyPatch):
 def test_registry_applies_worker_once_and_model_patch_lazily(monkeypatch: pytest.MonkeyPatch) -> None:
     registry_path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "__init__.py"
     registry = _load_source_module("vllm_omni_test_310p_patch_registry", registry_path)
-    calls = {"worker": 0, "talker": 0}
+    calls = {"worker": 0, "talker": 0, "code2wav": 0}
 
     _install_fake_module(
         monkeypatch,
@@ -204,52 +270,35 @@ def test_registry_applies_worker_once_and_model_patch_lazily(monkeypatch: pytest
         monkeypatch,
         "vllm_omni.platforms.npu._310p.patch.qwen3_tts",
         apply_talker_patches=lambda: calls.__setitem__("talker", calls["talker"] + 1),
+        apply_code2wav_patches=lambda: calls.__setitem__("code2wav", calls["code2wav"] + 1),
     )
 
     registry.apply_patches()
     registry.apply_patches()
     registry.apply_model_patches(SimpleNamespace(model_arch="OtherModel"))
     registry.apply_model_patches(SimpleNamespace(model_arch="Qwen3TTSTalkerForConditionalGeneration"))
+    registry.apply_model_patches(SimpleNamespace(model_arch="Qwen3TTSCode2Wav"))
 
-    assert calls == {"worker": 1, "talker": 1}
+    assert calls == {"worker": 1, "talker": 1, "code2wav": 1}
 
 
-def test_worker_patch_replaces_base_and_runs_disable_jit(monkeypatch: pytest.MonkeyPatch) -> None:
-    calls: list[str] = []
+def test_qwen3_tts_patch_registers_target_classes(monkeypatch: pytest.MonkeyPatch) -> None:
+    (
+        module,
+        (
+            fake_code_predictor,
+            fake_code_predictor_vllm,
+            fake_code2wav,
+            fake_prompt_builder,
+            fake_talker,
+        ),
+    ) = _load_qwen3_tts_patch(monkeypatch)
 
-    class FakeOmniNPUWorkerBase:
-        def _init_device(self):
-            calls.append("parent")
-            return "npu:0"
-
-    fake_worker_base = _install_fake_module(
-        monkeypatch,
-        "vllm_omni.platforms.npu.worker.base",
-        OmniNPUWorkerBase=FakeOmniNPUWorkerBase,
-    )
-    _install_fake_module(monkeypatch, "vllm_omni")
-    _install_fake_module(monkeypatch, "vllm_omni.platforms")
-    _install_fake_module(monkeypatch, "vllm_omni.platforms.npu")
-    _install_fake_module(
-        monkeypatch,
-        "vllm_omni.platforms.npu._310p",
-        disable_jit_compile=lambda: calls.append("disable_jit"),
-    )
-    _install_fake_module(monkeypatch, "vllm_omni.platforms.npu.worker", base=fake_worker_base)
-
-    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "worker.py"
-    module = _load_source_module("vllm_omni_test_310p_worker_patch", path)
-    module.apply_patch()
-
-    assert fake_worker_base.OmniNPUWorkerBase is module._OmniNPUWorkerBase310P
-    assert fake_worker_base.OmniNPUWorkerBase()._init_device() == "npu:0"
-    assert calls == ["parent", "disable_jit"]
-
-
-def test_qwen3_tts_patch_replaces_target_classes(monkeypatch: pytest.MonkeyPatch) -> None:
-    module, (fake_code_predictor, fake_prompt_builder, fake_talker) = _load_qwen3_tts_patch(monkeypatch)
+    original_common_wrapper = fake_code_predictor.CodePredictorWrapper
+    original_vllm_wrapper = fake_code_predictor_vllm.CodePredictorWrapper
 
     module.apply_talker_patches()
+    module.apply_code2wav_patches()
 
     assert fake_talker.Qwen3TTSTalkerForConditionalGeneration is module._Qwen3TTSTalker310P
     assert fake_talker.Qwen3TTSPromptEmbedsBuilder is module._Qwen3TTSPromptEmbedsBuilder310P
@@ -257,6 +306,112 @@ def test_qwen3_tts_patch_replaces_target_classes(monkeypatch: pytest.MonkeyPatch
     assert fake_code_predictor.CodePredictorAttention is module._Qwen3CodePredictorAttention310P
     assert fake_code_predictor.CodePredictorDecoderLayer is module._Qwen3CodePredictorDecoderLayer310P
     assert fake_code_predictor.CodePredictorBaseModel is module._Qwen3CodePredictorBaseModel310P
+    assert (
+        fake_code_predictor_vllm.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
+        is module._Qwen3TTSTalkerCodePredictor310P
+    )
+    assert fake_code_predictor.CodePredictorWrapper is original_common_wrapper
+    assert fake_code_predictor_vllm.CodePredictorWrapper is original_vllm_wrapper
+    assert fake_code2wav.Qwen3TTSCode2Wav is module._Qwen3TTSCode2Wav310P
+
+    code2wav = module._Qwen3TTSCode2Wav310P(
+        vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu")))
+    )
+
+    assert code2wav._npu_decoder_runtime_dtype(torch.device("cpu")) is torch.float16
+
+
+def test_qwen3_tts_code_predictor_forward_uses_projected_embedding_and_sampling(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    module, _ = _load_qwen3_tts_patch(monkeypatch)
+    predictor = module._Qwen3TTSTalkerCodePredictor310P(
+        vllm_config=object(),
+        config=object(),
+        talker_config=object(),
+    )
+    codec_weight = predictor.model.codec_embedding[0].weight.detach().clone()
+    projection_weight = predictor.small_to_mtp_projection.weight.detach().clone()
+
+    assert predictor.load_weights(iter(())) == {"loaded"}
+    torch.testing.assert_close(
+        predictor._projected_codec_embed_weight[0],
+        torch.nn.functional.linear(codec_weight, projection_weight),
+    )
+    assert predictor.small_to_mtp_projection.call_shapes == [(8, 4), (8, 4)]
+    predictor.small_to_mtp_projection.call_shapes.clear()
+    predictor._num_groups = 3
+    predictor._model_dtype = torch.float32
+    predictor._prefix_graphs_enabled = False
+    predictor._bucket_pos_ids = {}
+    predictor._device_graphs = {}
+    predictor._lm_heads_list = [
+        lambda hidden: torch.nn.functional.one_hot(
+            torch.full((hidden.shape[0],), 2, dtype=torch.long),
+            num_classes=8,
+        ).to(torch.float32),
+        lambda hidden: torch.nn.functional.one_hot(
+            torch.full((hidden.shape[0],), 3, dtype=torch.long),
+            num_classes=8,
+        ).to(torch.float32),
+    ]
+    predictor._setup_compile = lambda: None
+    predictor._padded_bsz = lambda bsz: bsz
+
+    def ensure_buffers(device, dtype, padded_bsz):
+        predictor._proj_buf = torch.zeros(padded_bsz, 4, 4, device=device, dtype=dtype)
+
+    predictor._ensure_buffers = ensure_buffers
+    predictor._compiled_model_fwd = lambda embeds, _positions: torch.zeros_like(embeds)
+
+    codes = predictor.forward(
+        torch.tensor([1]),
+        torch.zeros(1, 4),
+        torch.zeros(1, 4),
+        do_sample=False,
+    )
+
+    assert codes.tolist() == [[1, 2, 3]]
+    assert predictor.small_to_mtp_projection.call_shapes == [(1, 2, 4)]
+    torch.testing.assert_close(
+        predictor._proj_buf[0, 2],
+        predictor._projected_codec_embed_weight[0, 2],
+    )
+
+    filter_calls = []
+    sample_calls = []
+
+    def apply_top_k_top_p(logits, *, p, k, top_k):
+        filter_calls.append((p, k, top_k))
+        return logits
+
+    def random_sample(_probs, generators):
+        sample_calls.append(generators)
+        return torch.tensor([[4]])
+
+    monkeypatch.setattr(module, "apply_top_k_top_p", apply_top_k_top_p)
+    monkeypatch.setattr(module, "random_sample", random_sample)
+    predictor._num_groups = 2
+    predictor._wrapper_config.sampling_mode = "stored"
+    predictor._top_k = 2
+    predictor._top_p = 0.8
+    predictor._lm_heads_list = predictor._lm_heads_list[:1]
+    generator = torch.Generator().manual_seed(1234)
+
+    sampled_codes = predictor.forward(
+        torch.tensor([1]),
+        torch.zeros(1, 4),
+        torch.zeros(1, 4),
+        generator=generator,
+    )
+
+    assert sampled_codes.tolist() == [[1, 4]]
+    assert len(filter_calls) == 1
+    top_p_tensor, top_k_tensor, top_k_hint = filter_calls[0]
+    torch.testing.assert_close(top_p_tensor, torch.tensor([0.8]))
+    assert top_k_tensor.tolist() == [2]
+    assert top_k_hint == 2
+    assert sample_calls == [{0: generator}]
 
 
 def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -265,12 +420,14 @@ def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.Monk
 
     assert talker._embedding_dtype is torch.float16
     assert talker._prompt_builder._embedding_dtype is torch.float16
+    assert talker.talker_mtp_graph_safe is False
+    assert talker.talker_mtp_accepts_per_row_generators is True
     assert talker.load_weights([]) == {"loaded"}
-    assert talker.encoder.to_calls[-1] == {"dtype": torch.float16}
+    assert talker.encoder.to_calls[-1] == {"device": torch.device("cpu"), "dtype": torch.float32}
 
     codes = talker._encode_ref_audio_batch([np.zeros(8, dtype=np.float32)], 24000, device=torch.device("cpu"))
 
-    assert talker.encoder.last_input_dtype is torch.float16
+    assert talker.encoder.last_input_dtype is torch.float32
     assert len(codes) == 1
     assert codes[0].dtype is torch.long
     assert codes[0].shape == (4, 2)
@@ -311,28 +468,154 @@ def test_qwen3_tts_prompt_patch_runs_stft_frontend_on_cpu(monkeypatch: pytest.Mo
     assert speaker.dtype is torch.float16
 
 
-def test_qwen3_tts_mimi_codebook_quantize_uses_cpu_fp32_cdist(monkeypatch: pytest.MonkeyPatch) -> None:
-    module, _ = _load_qwen3_tts_patch(monkeypatch)
-    real_cdist = torch.cdist
-    captured = {}
+def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest.MonkeyPatch) -> None:
+    rotary_calls = []
+    rms_calls = []
 
-    def fake_cdist(x1, x2, p=2):
-        captured["x1_device"] = x1.device
-        captured["x2_device"] = x2.device
-        captured["x1_dtype"] = x1.dtype
-        captured["x2_dtype"] = x2.dtype
-        return real_cdist(x1, x2, p=p)
+    def rotary_mul(hidden_states, cos, sin):
+        rotary_calls.append((hidden_states, cos, sin))
+        return hidden_states + 1
 
-    monkeypatch.setattr(torch, "cdist", fake_cdist)
-    codebook = object.__new__(module._MimiEuclideanCodebook310P)
-    codebook._embed = torch.tensor([[0.0, 0.0], [2.0, 0.0]], dtype=torch.float16)
+    def rms_norm(hidden_states, weight, *, epsilon):
+        rms_calls.append((hidden_states, weight, epsilon))
+        return hidden_states * weight, None
 
-    indices = codebook.quantize(torch.tensor([[0.1, 0.0], [1.8, 0.0]], dtype=torch.float16))
+    _install_fake_module(
+        monkeypatch,
+        "torch_npu",
+        npu_rotary_mul=rotary_mul,
+        npu_rms_norm=rms_norm,
+    )
+    _install_fake_module(monkeypatch, "vllm")
+    _install_fake_module(monkeypatch, "vllm.logger", init_logger=lambda _name: SimpleNamespace(debug=lambda *_: None))
 
-    assert captured == {
-        "x1_device": torch.device("cpu"),
-        "x2_device": torch.device("cpu"),
-        "x1_dtype": torch.float32,
-        "x2_dtype": torch.float32,
+    class FakeRMSNorm(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.tensor([2.0, 3.0]))
+            self.variance_epsilon = 1e-5
+
+        def forward(self, hidden_states):
+            return hidden_states
+
+    def original_rope(q, k, cos, sin):
+        del cos, sin
+        return q, k
+
+    tokenizer = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2",
+        Qwen3TTSTokenizerV2DecoderRMSNorm=FakeRMSNorm,
+        apply_rotary_pos_emb=original_rope,
+    )
+    tokenizer_package = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz",
+        modeling_qwen3_tts_tokenizer_v2=tokenizer,
+    )
+    _install_fake_module(monkeypatch, "vllm_omni")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor.models")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor.models.qwen3_tts")
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz",
+        tokenizer_package,
+    )
+
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts_tokenizer_v2.py"
+    module = _load_source_module("vllm_omni_test_qwen3_tts_tokenizer_npu_patch", path)
+    module.apply_qwen3_tts_tokenizer_v2_patch()
+
+    q = torch.zeros(1, 2, 3, 2)
+    k = torch.ones_like(q)
+    cos = torch.zeros(1, 3, 2)
+    sin = torch.ones_like(cos)
+    q_out, k_out = tokenizer.apply_rotary_pos_emb(q, k, cos, sin)
+    norm = FakeRMSNorm()
+    norm_out = norm(torch.ones(1, 2))
+
+    assert len(rotary_calls) == 2
+    assert rotary_calls[0][1].shape == (1, 1, 3, 2)
+    torch.testing.assert_close(q_out, q + 1)
+    torch.testing.assert_close(k_out, k + 1)
+    assert len(rms_calls) == 1
+    assert rms_calls[0][2] == pytest.approx(1e-5)
+    torch.testing.assert_close(norm_out, torch.tensor([[2.0, 3.0]]))
+
+
+def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytest.MonkeyPatch) -> None:
+    linear_weights = []
+    conv_weights = []
+
+    def maybe_trans_nz(weight):
+        linear_weights.append(weight)
+        return weight
+
+    def format_cast(weight, fmt):
+        conv_weights.append((weight, fmt))
+        return weight
+
+    class FakeDecoder(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(4, 4)
+            self.conv = torch.nn.Conv1d(4, 4, 3)
+            self.deconv = torch.nn.ConvTranspose1d(4, 4, 4)
+            self.grouped_conv = torch.nn.Conv1d(4, 4, 3, groups=2)
+            self.cache_precompute_calls = 0
+
+        def precompute_snake_caches(self):
+            self.cache_precompute_calls += 1
+
+    class FakeCode2Wav:
+        def __init__(self, *, vllm_config, prefix=""):
+            self.vllm_config = vllm_config
+            self.prefix = prefix
+            self.decoder = FakeDecoder()
+
+        def _npu_decoder_runtime_dtype(self, _device):
+            return torch.float16
+
+        def load_weights(self, weights):
+            assert list(weights) == []
+            return {"loaded"}
+
+    logger = SimpleNamespace(info=lambda *_: None, debug=lambda *_: None)
+    current_platform = SimpleNamespace(is_npu=lambda: False)
+    target = _install_fake_module(
+        monkeypatch,
+        "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav",
+        Qwen3TTSCode2Wav=FakeCode2Wav,
+    )
+    _install_fake_module(monkeypatch, "torch_npu", npu_format_cast=format_cast)
+    _install_fake_module(monkeypatch, "vllm")
+    _install_fake_module(monkeypatch, "vllm.config", VllmConfig=object)
+    _install_fake_module(monkeypatch, "vllm.logger", init_logger=lambda _name: logger)
+    _install_fake_module(monkeypatch, "vllm_ascend")
+    _install_fake_module(monkeypatch, "vllm_ascend.utils", maybe_trans_nz=maybe_trans_nz)
+    _install_fake_module(monkeypatch, "vllm_omni")
+    _install_fake_module(monkeypatch, "vllm_omni.platforms", current_omni_platform=current_platform)
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor.models")
+    _install_fake_module(monkeypatch, "vllm_omni.model_executor.models.qwen3_tts")
+
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts_code2wav.py"
+    module = _load_source_module("vllm_omni_test_qwen3_tts_code2wav_npu_patch", path)
+    module.apply_qwen3_tts_code2wav_patch()
+
+    model = target.Qwen3TTSCode2Wav(
+        vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu"))),
+        prefix="stage1",
+    )
+    assert model.load_weights(iter(())) == {"loaded"}
+
+    assert model.prefix == "stage1"
+    assert model.decoder.linear.weight.dtype is torch.float16
+    assert [weight.data_ptr() for weight in linear_weights] == [model.decoder.linear.weight.data_ptr()]
+    assert {weight.data_ptr() for weight, _ in conv_weights} == {
+        model.decoder.conv.weight.data_ptr(),
+        model.decoder.deconv.weight.data_ptr(),
     }
-    assert indices.tolist() == [0, 1]
+    assert all(fmt == module._ACL_FORMAT_FRACTAL_Z for _, fmt in conv_weights)
+    assert model.decoder.cache_precompute_calls == 1
diff --git a/tests/platforms/test_cuda_ir_op_priority.py b/tests/platforms/test_cuda_ir_op_priority.py
new file mode 100644
index 000000000..8542ea63d
--- /dev/null
+++ b/tests/platforms/test_cuda_ir_op_priority.py
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from vllm.config import CompilationConfig, CompilationMode, DeviceConfig, VllmConfig
+from vllm.platforms import current_platform
+
+# Importing CudaOmniPlatform pulls in vllm.platforms.cuda, which imports the
+# CUDA-only ``vllm._C_stable_libtorch`` extension at module top level. That
+# extension is absent on non-CUDA builds (e.g. XPU), so skip the whole module
+# there before the import can crash collection. ``is_cuda()`` resolves the
+# platform without importing cuda.py.
+if not current_platform.is_cuda():
+    pytest.skip("CUDA-only IR op priority tests", allow_module_level=True)
+
+from vllm_omni.platforms.cuda.platform import CudaOmniPlatform  # noqa: E402
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _vllm_config(*, backend: str = "inductor", mode: CompilationMode) -> VllmConfig:
+    return VllmConfig(
+        device_config=DeviceConfig(device="cpu"),
+        compilation_config=CompilationConfig(backend=backend, mode=mode),
+    )
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [CompilationMode.NONE, CompilationMode.VLLM_COMPILE, CompilationMode.STOCK_TORCH_COMPILE],
+)
+def test_cuda_default_ir_op_priority_prefers_vllm_c_when_inductor_backend(mode: CompilationMode) -> None:
+    """Regression for #4964: inductor-active configs must not switch to native-only."""
+    priority = CudaOmniPlatform.get_default_ir_op_priority(_vllm_config(mode=mode))
+    assert priority.rms_norm == ["vllm_c", "native"]
+    assert priority.fused_add_rms_norm == ["vllm_c", "native"]
+
+
+def test_cuda_default_ir_op_priority_with_oink(monkeypatch: pytest.MonkeyPatch) -> None:
+    import vllm.envs as envs
+
+    monkeypatch.setattr(envs, "VLLM_USE_OINK_OPS", True)
+    priority = CudaOmniPlatform.get_default_ir_op_priority(
+        _vllm_config(mode=CompilationMode.VLLM_COMPILE),
+    )
+    assert priority.rms_norm == ["oink", "vllm_c", "native"]
+    assert priority.fused_add_rms_norm == ["oink", "vllm_c", "native"]
diff --git a/tests/plugin/__init__.py b/tests/plugin/__init__.py
new file mode 100644
index 000000000..988131360
--- /dev/null
+++ b/tests/plugin/__init__.py
@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0
diff --git a/tests/plugin/test_gguf_diffusion.py b/tests/plugin/test_gguf_diffusion.py
new file mode 100644
index 000000000..44de4a1fb
--- /dev/null
+++ b/tests/plugin/test_gguf_diffusion.py
@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end tests for GGUF quantization on diffusion models.
+
+Validates that GGUF-quantized diffusion models generate valid images and
+use less peak GPU memory than BF16 baseline.
+
+Requires vllm-omni to be installed alongside the plugin.
+
+Usage:
+    pytest tests/diffusion/test_gguf_memory.py -v
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import numpy as np
+import pytest
+import torch
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+class DiffusionGGUFTestConfig(NamedTuple):
+    artifact_prefix: str
+    gguf_model: str
+    hf_model: str
+    min_cosine_similarity: float = 0.95
+
+
+Z_IMAGE_CONFIG = DiffusionGGUFTestConfig(
+    artifact_prefix="zimage",
+    hf_model="Tongyi-MAI/Z-Image-Turbo",
+    gguf_model="unsloth/Z-Image-Turbo-GGUF:Q4_0",
+)
+
+
+FLUX_CONFIG = DiffusionGGUFTestConfig(
+    artifact_prefix="flux2_klein",
+    hf_model="black-forest-labs/FLUX.2-klein-4B",
+    gguf_model="unsloth/FLUX.2-klein-4B-GGUF:Q8_0",
+)
+
+
+def _image_cosine_similarity(hf_image, gguf_image) -> float:
+    hf_tensor = torch.as_tensor(np.array(hf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
+    gguf_tensor = torch.as_tensor(np.array(gguf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten()
+    return torch.nn.functional.cosine_similarity(hf_tensor, gguf_tensor, dim=0).item()
+
+
+def _generate_single_stage_image(
+    model: str,
+    height: int = 256,
+    width: int = 256,
+    num_inference_steps: int = 20,
+    seed: int = 42,
+    **extra_kwargs,
+) -> tuple[list, float]:
+    """Generate an image with a single-stage diffusion model.
+
+    Returns (images, peak_memory_gib).
+    """
+    omni_kwargs = dict(extra_kwargs)
+
+    memory_monitor = DeviceMemoryMonitor(device_index=0)
+    memory_monitor.start()
+    try:
+        with OmniRunner(model, **omni_kwargs) as runner:
+            generator = torch.Generator(
+                device=current_omni_platform.device_type,
+            ).manual_seed(seed)
+            outputs = runner.omni.generate(
+                "a photo of a cat sitting on a laptop keyboard",
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=0.0,
+                    generator=generator,
+                ),
+            )
+    finally:
+        memory_monitor.stop()
+    peak_mem = memory_monitor.peak_used_mb / 1024
+
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    if hasattr(first_output, "images") and first_output.images:
+        images = first_output.images
+    else:
+        assert hasattr(first_output, "request_output") and first_output.request_output
+        request_output = first_output.request_output
+        if isinstance(request_output, list):
+            req_out = request_output[0]
+        else:
+            req_out = request_output
+        assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
+        images = req_out.images
+    assert len(images) >= 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+    return images, peak_mem
+
+
+@hardware_test(res={"cuda": "L4"})
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.slow
+@pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"])
+def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Z-Image-Turbo GGUF generates valid images and uses less memory than BF16."""
+    monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+    # BF16 baseline
+    hf_images, mem_bf16 = _generate_single_stage_image(
+        model=model.hf_model,
+    )
+
+    # GGUF
+    images, mem_gguf = _generate_single_stage_image(
+        model=model.hf_model,
+        diffusion_quantization_config={
+            "method": "gguf",
+            "gguf_model": model.gguf_model,
+        },
+    )
+
+    hf_image_path = f"test_{model.artifact_prefix}_hf.png"
+    gguf_image_path = f"test_{model.artifact_prefix}_gguf.png"
+    assert len(hf_images) >= 1
+    hf_images[0].save(hf_image_path)
+    assert len(images) >= 1
+    images[0].save(gguf_image_path)
+    print(f"Saved HF image: {hf_image_path}")
+    print(f"Saved GGUF image: {gguf_image_path}")
+
+    image_cosine_similarity = _image_cosine_similarity(hf_images[0], images[0])
+    print(f"{model.artifact_prefix} image cosine similarity: {image_cosine_similarity:.4f}")
+    assert image_cosine_similarity >= model.min_cosine_similarity, (
+        f"GGUF image cosine similarity ({image_cosine_similarity:.4f}) should be >= {model.min_cosine_similarity:.4f}"
+    )
+
+    print(f"{model.artifact_prefix} BF16 peak VRAM delta: {mem_bf16:.2f} GiB")
+    print(f"{model.artifact_prefix} GGUF peak VRAM delta: {mem_gguf:.2f} GiB")
+    reduction = (mem_bf16 - mem_gguf) / mem_bf16 * 100
+    print(f"VRAM reduction: {reduction:.1f}%")
+    assert mem_gguf < mem_bf16, f"GGUF ({mem_gguf:.2f} GiB) should use less VRAM than BF16 ({mem_bf16:.2f} GiB)"
diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py
index 63a67d652..54683e67c 100644
--- a/tests/test_config_factory.py
+++ b/tests/test_config_factory.py
@@ -14,8 +14,9 @@ from unittest.mock import patch
 import pytest
 from transformers import PretrainedConfig, Qwen3OmniMoeConfig
 
-from tests.helpers.stage_config import get_deploy_config_path
+from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
 from vllm_omni.config.config_factory import StageConfigFactory
+from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
 from vllm_omni.config.stage_config import (
     DeployConfig,
@@ -39,6 +40,15 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_black
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def clear_config_factory_caches():
+    """Clear cached classmethods from the StageConfigFactory to prevent test pollution."""
+    yield
+    StageConfigFactory.get_hf_config.cache_clear()
+    StageConfigFactory.try_infer_model_type.cache_clear()
+    StageConfigFactory.get_pipeline_config.cache_clear()
+
+
 Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
 Q3_OMNI_THINKER_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=False)
 
@@ -571,6 +581,7 @@ class TestStagePipelineConfig:
         assert s.final_output is False
         assert s.sampling_constraints == {}
         assert s.engine_output_type is None
+        assert s.scheduler_cls is None
 
 
 class TestPipelineConfigNew:
@@ -665,6 +676,43 @@ class TestPipelineRegistration:
         assert resolved_config is not None
         assert len(resolved_config) > 0
 
+    def test_deploy_override_uses_correct_endpoint_restrictions(self, clean_pipeline_registry, tmp_path):
+        """Ensure endpoint restrictions must come from the final pipeline
+        after deploy config overrides, not the auto-detected pipeline.
+        """
+        # Register two pipeline configs, where one has an endpoint restriction, and one doesn't
+        restriction = EndpointRestriction(
+            OmniServingCapability.COMPLETIONS,
+            "pipeline_a blocks completions",
+        )
+        pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,))
+        pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=())
+        register_pipeline(pipe_a)
+        register_pipeline(pipe_b)
+
+        # Create a config with the autodetected type, and write the
+        # deploy config specifying the override type to a temp path
+        class FakeConfig(PretrainedConfig):
+            model_type = "detect_type"
+
+        deploy_yaml = tmp_path / "override.yaml"
+        deploy_yaml.write_text("pipeline: override_type\n")
+
+        # Get the endpoint restrictions, passing the deploy config with the override
+        # type + patching the config for the detected type. Ensure that the endpoint
+        # restrictions correspond to the type in the deploy config.
+        with patch(
+            "vllm_omni.config.config_factory.get_config",
+            return_value=FakeConfig(),
+        ):
+            restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions(
+                model="fake/model",
+                trust_remote_code=False,
+                deploy_config_path=str(deploy_yaml),
+            )
+
+        assert restrictions == ()
+
 
 class TestResolveScheduler:
     def test_all_execution_types_handled(self):
@@ -739,6 +787,41 @@ stages:
         assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}}
         assert "compilation_config" not in deploy.stages[0].engine_extras
 
+    def test_load_voxcpm2_deploy_config_preserves_engine_extras(self):
+        deploy_path = get_deploy_config_path("voxcpm2.yaml")
+        raw_stage = get_deploy_config_stage("voxcpm2.yaml", 0)
+        expected_runtime_config = raw_stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"]
+
+        deploy = load_deploy_config(deploy_path)
+        runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"]
+        assert runtime_config == expected_runtime_config
+
+    @pytest.mark.parametrize(
+        ("deploy_name", "pipeline_name", "stage_count", "final_output_type"),
+        [
+            ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image"),
+            ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text"),
+            ("omnivoice.yaml", "omnivoice", 1, "audio"),
+        ],
+    )
+    def test_load_new_registry_backed_deploy_configs(
+        self,
+        deploy_name: str,
+        pipeline_name: str,
+        stage_count: int,
+        final_output_type: str,
+    ):
+        deploy_path = Path(get_deploy_config_path(deploy_name))
+        deploy = load_deploy_config(deploy_path)
+        assert deploy.pipeline == pipeline_name
+
+        with patch("vllm_omni.platforms.current_omni_platform") as platform:
+            platform.device_name = "cuda"
+            stages = merge_pipeline_deploy(OMNI_PIPELINES[pipeline_name], deploy)
+        assert len(stages) == stage_count
+        assert stages[-1].final_output is True
+        assert stages[-1].final_output_type == final_output_type
+
     def test_merge_pipeline_deploy(self):
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
@@ -800,6 +883,26 @@ stages:
 
         assert stages[0].yaml_runtime["requires_multimodal_data"] is True
 
+    def test_merge_pipeline_deploy_preserves_pipeline_scheduler_cls(self):
+        scheduler_cls = "tests.fake.CustomScheduler"
+        pipeline = PipelineConfig(
+            model_type="test_scheduler",
+            model_arch="TestModel",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="ar",
+                    execution_type=StageExecutionType.LLM_AR,
+                    scheduler_cls=scheduler_cls,
+                ),
+            ),
+        )
+        deploy = DeployConfig(async_chunk=False, stages=[StageDeployConfig(stage_id=0)])
+
+        stages = merge_pipeline_deploy(pipeline, deploy)
+
+        assert stages[0].scheduler_cls == scheduler_cls
+
     def test_mixed_schema_preserves_flat_fields(self):
         """Ensure flat fields are not dropped when engine_args are present."""
         fake_config = {
@@ -928,6 +1031,54 @@ stages:
         assert "foo" in stage.engine_extras
         assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3}
 
+    def test_explicit_engine_extras_merge_with_flat_and_engine_args(self):
+        """Explicit engine_extras should be preserved with existing pass-through styles."""
+        fake_config = {
+            "stages": [
+                {
+                    "stage_id": 0,
+                    "engine_extras": {
+                        "hf_overrides": {
+                            "runtime_config": {
+                                "explicit_only": True,
+                                "shared": "explicit",
+                                "nested": {"a": 1},
+                            }
+                        }
+                    },
+                    "hf_overrides": {
+                        "runtime_config": {
+                            "flat_only": True,
+                            "shared": "flat",
+                            "nested": {"b": 2},
+                        }
+                    },
+                    "engine_args": {
+                        "hf_overrides": {
+                            "runtime_config": {
+                                "engine_args_only": True,
+                                "shared": "engine_args",
+                                "nested": {"c": 3},
+                            }
+                        }
+                    },
+                },
+            ]
+        }
+
+        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
+            deploy = load_deploy_config("dummy.yaml")
+
+        assert deploy.stages[0].engine_extras["hf_overrides"] == {
+            "runtime_config": {
+                "engine_args_only": True,
+                "explicit_only": True,
+                "flat_only": True,
+                "nested": {"a": 1, "b": 2, "c": 3},
+                "shared": "engine_args",
+            }
+        }
+
     def test_deep_merge_does_not_mutate_inputs(self):
         """Merging engine_args must not mutate the base stage dict."""
         fake_config = {
@@ -1015,7 +1166,10 @@ class TestQwen3OmniPipeline:
         assert isinstance(s, StagePipelineConfig)
         assert s.input_sources == (0,)
         assert s.sampling_constraints["stop_token_ids"] == [2150]
-        assert s.custom_process_input_func is not None
+        # thinker2talker was removed: sync_process_input_func always wins.
+        assert s.custom_process_input_func is None
+        assert s.sync_process_input_func is not None
+        assert s.sync_process_input_func.endswith("thinker2talker_token_only")
         assert s.custom_process_next_stage_input_func is not None
 
     def test_code2wav(self):
@@ -1029,7 +1183,7 @@ class TestQwen3OmniPipeline:
         assert isinstance(s, StagePipelineConfig)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
-        assert s.custom_process_input_func is not None
+        assert s.engine_output_type == "audio"
 
 
 class TestQwen2_5OmniPipeline:
@@ -1567,7 +1721,7 @@ class TestAuraOmniDeploy:
     def test_aura_omni_deploy_resolves_four_native_stages(self):
         pipeline_cfg = StageConfigFactory.resolve_pipeline_config("aura_omni")
 
-        stages = StageConfigFactory._create_from_registry(
+        stages, _ = StageConfigFactory._create_from_registry(
             "qwen3_tts",
             pipeline_cfg,
             cli_overrides={},
@@ -1588,6 +1742,197 @@ class TestAuraOmniDeploy:
         assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav"
 
 
+class TestDeployCliOverrideFlow:
+    """Test deploy-YAML baselines overridden by CLI runtime overrides."""
+
+    def test_diffusion_deploy_fields_can_be_overridden_by_cli(self, tmp_path):
+        deploy_path = tmp_path / "diffusion_stage.yaml"
+        deploy_path.write_text(
+            """
+async_chunk: false
+stages:
+  - stage_id: 0
+    devices: "4,5"
+    parallel_config:
+      pipeline_parallel_size: 1
+      data_parallel_size: 1
+      tensor_parallel_size: 1
+      enable_expert_parallel: false
+      sequence_parallel_size: 1
+      ulysses_degree: 1
+      ring_degree: 1
+      cfg_parallel_size: 1
+      vae_patch_parallel_size: 1
+      use_hsdp: false
+      hsdp_shard_size: -1
+      hsdp_replicate_size: 1
+    engine_args:
+      cache_backend: cache_dit
+      diffusion_attention_backend: FLASH_ATTN
+      diffusion_kv_cache_dtype: auto
+      step_execution: false
+      vae_use_tiling: false
+      enable_cpu_offload: false
+      max_generated_image_size: 1048576
+      tts_max_instructions_length: 1000
+""",
+            encoding="utf-8",
+        )
+
+        pipeline = PipelineConfig(
+            model_type="test_diffusion",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="dit",
+                    execution_type=StageExecutionType.DIFFUSION,
+                    final_output=True,
+                    final_output_type="image",
+                ),
+            ),
+        )
+        deploy = load_deploy_config(deploy_path)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        stage = stages[0]
+
+        assert stage.yaml_engine_args["parallel_config"]["pipeline_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["data_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["enable_expert_parallel"] is False
+        assert stage.yaml_engine_args["parallel_config"]["sequence_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["ulysses_degree"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["ring_degree"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["cfg_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["vae_patch_parallel_size"] == 1
+        assert stage.yaml_engine_args["cache_backend"] == "cache_dit"
+        assert stage.yaml_engine_args["parallel_config"]["use_hsdp"] is False
+        assert stage.yaml_engine_args["parallel_config"]["hsdp_shard_size"] == -1
+        assert stage.yaml_engine_args["parallel_config"]["hsdp_replicate_size"] == 1
+        assert stage.yaml_engine_args["diffusion_attention_backend"] == "FLASH_ATTN"
+        assert stage.yaml_engine_args["diffusion_kv_cache_dtype"] == "auto"
+        assert stage.yaml_engine_args["step_execution"] is False
+        assert stage.yaml_engine_args["vae_use_tiling"] is False
+        assert stage.yaml_engine_args["enable_cpu_offload"] is False
+        assert stage.yaml_engine_args["max_generated_image_size"] == 1048576
+        assert stage.yaml_engine_args["tts_max_instructions_length"] == 1000
+
+        stage.runtime_overrides = StageConfigFactory._merge_cli_overrides(
+            stage,
+            {
+                "pipeline_parallel_size": 2,
+                "data_parallel_size": 3,
+                "tensor_parallel_size": 4,
+                "enable_expert_parallel": True,
+                "sequence_parallel_size": 24,
+                "ulysses_degree": 2,
+                "ring_degree": 4,
+                "cfg_parallel_size": 2,
+                "vae_patch_parallel_size": 2,
+                "cache_backend": "tea_cache",
+                "use_hsdp": True,
+                "hsdp_shard_size": 8,
+                "hsdp_replicate_size": 2,
+                "diffusion_attention_backend": "SAGE_ATTN",
+                "diffusion_kv_cache_dtype": "fp8",
+                "step_execution": True,
+                "vae_use_tiling": True,
+                "enable_cpu_offload": True,
+                "max_generated_image_size": 2097152,
+                "tts_max_instructions_length": 2000,
+            },
+        )
+
+        omega_config = stage.to_omegaconf()
+
+        assert omega_config.engine_args.cache_backend == "tea_cache"
+        assert omega_config.engine_args.diffusion_attention_backend == "SAGE_ATTN"
+        assert omega_config.engine_args.diffusion_kv_cache_dtype == "fp8"
+        assert omega_config.engine_args.step_execution is True
+        assert omega_config.engine_args.vae_use_tiling is True
+        assert omega_config.engine_args.enable_cpu_offload is True
+        assert omega_config.engine_args.max_generated_image_size == 2097152
+        assert omega_config.engine_args.tts_max_instructions_length == 2000
+        assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.data_parallel_size == 3
+        assert omega_config.engine_args.parallel_config.tensor_parallel_size == 4
+        assert omega_config.engine_args.parallel_config.enable_expert_parallel is True
+        assert omega_config.engine_args.parallel_config.sequence_parallel_size == 24
+        assert omega_config.engine_args.parallel_config.ulysses_degree == 2
+        assert omega_config.engine_args.parallel_config.ring_degree == 4
+        assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.use_hsdp is True
+        assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8
+        assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2
+
+    def test_llm_deploy_fields_can_be_overridden_by_cli(self, tmp_path):
+        deploy_path = tmp_path / "llm_stage.yaml"
+        deploy_path.write_text(
+            """
+async_chunk: false
+stages:
+  - stage_id: 0
+    devices: "0,1"
+    engine_args:
+      tensor_parallel_size: 1
+      enable_expert_parallel: false
+      gpu_memory_utilization: 0.5
+      max_num_seqs: 16
+      max_num_batched_tokens: 1024
+      max_model_len: 4096
+      enforce_eager: false
+""",
+            encoding="utf-8",
+        )
+
+        pipeline = PipelineConfig(
+            model_type="test_llm",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="thinker",
+                    execution_type=StageExecutionType.LLM_AR,
+                    final_output=True,
+                    final_output_type="text",
+                ),
+            ),
+        )
+        deploy = load_deploy_config(deploy_path)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        stage = stages[0]
+
+        assert stage.yaml_engine_args["tensor_parallel_size"] == 1
+        assert stage.yaml_engine_args["enable_expert_parallel"] is False
+        assert stage.yaml_engine_args["gpu_memory_utilization"] == 0.5
+        assert stage.yaml_engine_args["max_num_seqs"] == 16
+        assert stage.yaml_engine_args["max_num_batched_tokens"] == 1024
+        assert stage.yaml_engine_args["max_model_len"] == 4096
+        assert stage.yaml_engine_args["enforce_eager"] is False
+
+        stage.runtime_overrides = StageConfigFactory._merge_cli_overrides(
+            stage,
+            {
+                "tensor_parallel_size": 2,
+                "enable_expert_parallel": True,
+                "gpu_memory_utilization": 0.9,
+                "max_num_seqs": 32,
+                "max_num_batched_tokens": 2048,
+                "max_model_len": 8192,
+                "enforce_eager": True,
+            },
+        )
+
+        omega_config = stage.to_omegaconf()
+
+        assert omega_config.engine_args.tensor_parallel_size == 2
+        assert omega_config.engine_args.enable_expert_parallel is True
+        assert omega_config.engine_args.gpu_memory_utilization == 0.9
+        assert omega_config.engine_args.max_num_seqs == 32
+        assert omega_config.engine_args.max_num_batched_tokens == 2048
+        assert omega_config.engine_args.max_model_len == 8192
+        assert omega_config.engine_args.enforce_eager is True
+
+
 class TestSentinelDefaultPrecedence:
     """Caller-typed (non-None) values win over YAML; None values fall through
     to YAML / dataclass defaults (#3035)."""
@@ -1598,11 +1943,12 @@ class TestSentinelDefaultPrecedence:
             model_type,
             Q3_OMNI_ALL_STAGES_HF_CONFIG,
         )
-        return StageConfigFactory._create_from_registry(
+        stages, _ = StageConfigFactory._create_from_registry(
             "qwen3_omni_moe",
             pipeline_cfg,
             cli_overrides=cli_overrides,
         )
+        return stages
 
     def test_typed_kwarg_overrides_yaml(self):
         stages = self._stages({"max_num_seqs": 999})
diff --git a/tests/test_version.py b/tests/test_version.py
index 07e622a7d..911a79611 100644
--- a/tests/test_version.py
+++ b/tests/test_version.py
@@ -56,3 +56,26 @@ def test_warning_contains_version_strings():
     msg = str(record[0].message)
     assert "0.19.0rc2.dev21" in msg
     assert "0.18.0" in msg
+
+
+def test_omni_plugin_loader_invokes_vllm_general_plugins_first(monkeypatch):
+    import vllm.plugins as vllm_plugins
+
+    import vllm_omni.plugins as omni_plugins
+
+    calls = []
+    monkeypatch.setattr(omni_plugins, "omni_plugins_loaded", False)
+    monkeypatch.setattr(vllm_plugins, "load_general_plugins", lambda: calls.append("vllm"))
+    monkeypatch.setattr(
+        omni_plugins,
+        "load_omni_plugins_by_group",
+        lambda group: {"omni": lambda: calls.append(group)},
+    )
+
+    omni_plugins.load_omni_general_plugins()
+
+    assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP]
+
+    omni_plugins.load_omni_general_plugins()
+
+    assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP]
diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py
old mode 100644
new mode 100755
index b701d9f8b..b92b0095b
--- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py
+++ b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py
@@ -7,6 +7,7 @@ without requiring the actual checkpoint or GPU.
 """
 
 import asyncio
+import json
 import time
 from collections import OrderedDict, defaultdict
 from types import SimpleNamespace
@@ -1309,7 +1310,10 @@ class TestVoiceCloneReferenceCache:
         parent = tmp_path / "OmniVoice"
         subdir = parent / "audio_tokenizer"
         subdir.mkdir(parents=True)
-        (subdir / "config.json").write_text("{}", encoding="utf-8")
+        (subdir / "config.json").write_text(
+            json.dumps({"model_type": "higgs_audio_v2_tokenizer"}),
+            encoding="utf-8",
+        )
 
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
@@ -1317,6 +1321,32 @@ class TestVoiceCloneReferenceCache:
         monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir))
         assert tok._resolve_audio_tokenizer_dir() == str(subdir)
 
+    def test_audio_tokenizer_dir_rejects_omnivoice_config(self, tmp_path, monkeypatch):
+        import huggingface_hub
+        import huggingface_hub.constants as hub_constants
+
+        from vllm_omni.model_executor.models.higgs_audio_v2 import higgs_audio_v2_tokenizer as tok
+
+        parent = tmp_path / "OmniVoice"
+        subdir = parent / "audio_tokenizer"
+        subdir.mkdir(parents=True)
+        (subdir / "config.json").write_text(
+            json.dumps({"model_type": "omnivoice"}),
+            encoding="utf-8",
+        )
+
+        monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
+        monkeypatch.setenv("HIGGS_AUDIO_V2_TOKENIZER_PATH", "")
+        monkeypatch.setattr(
+            huggingface_hub,
+            "try_to_load_from_cache",
+            lambda **_: str(subdir / "config.json"),
+        )
+        monkeypatch.setattr(hub_constants, "HF_HUB_CACHE", str(tmp_path / "empty_cache"))
+
+        assert tok._normalize_audio_tokenizer_dir(str(parent)) is None
+        assert tok._resolve_audio_tokenizer_dir() is None
+
     def test_higgs_v3_ref_code_cache_returns_clone(self):
         from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 
diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py
index 437fbd50a..5bee3e3b0 100644
--- a/tests/utils/test_tracking_parser.py
+++ b/tests/utils/test_tracking_parser.py
@@ -16,7 +16,12 @@ from vllm_omni.config.stage_config import (
     StageDeployConfig,
     StagePipelineConfig,
 )
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace
+from vllm_omni.utils.tracking_parser import (
+    UNSET,
+    TrackingArgumentParser,
+    TrackingNamespace,
+    build_shadow_kwargs,
+)
 
 ### Fake pipeline/deploy config for integration tests
 
@@ -326,6 +331,132 @@ def test_omitted_nargs():
     assert ns.items is None
 
 
+### Tests for in-place mutating actions (append / extend / count).
+# These actions mutate their default in place, which would crash on the bare
+# ``UNSET`` sentinel. The shadow parser remaps them to a non-mutating store-style
+# action (see build_shadow_kwargs in tracking_parser), so the real namespace
+# still accumulates while explicit-arg tracking keeps working.
+def test_append_action_omitted():
+    """Omitted append args parse without error and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--tag", action="append")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag is None
+
+
+def test_append_action_explicit():
+    """Repeated append args are collected and marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--tag", action="append")
+    ns = p.parse_args(["--tag", "a", "--tag", "b"])
+    assert ns.explicit_keys == {"tag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag == ["a", "b"]
+
+
+def test_append_const_action_omitted():
+    """Omitted append_const args aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--dbg", action="append_const", const="x", dest="flags")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flags is None
+
+
+def test_append_const_action_explicit():
+    """Repeated append_const args accumulate the const and are explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--dbg", action="append_const", const="x", dest="flags")
+    ns = p.parse_args(["--dbg", "--dbg"])
+    assert ns.explicit_keys == {"flags"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.flags == ["x", "x"]
+
+
+def test_extend_action_omitted():
+    """Omitted extend args parse without error and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", action="extend", nargs="+")
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items is None
+
+
+def test_extend_action_explicit():
+    """Repeated extend args are flattened into one list and marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("--items", action="extend", nargs="+")
+    ns = p.parse_args(["--items", "a", "b", "--items", "c"])
+    assert ns.explicit_keys == {"items"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.items == ["a", "b", "c"]
+
+
+def test_count_action_omitted():
+    """Omitted count args keep their default and aren't marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("-v", "--verbose", action="count", default=0)
+    ns = p.parse_args([])
+    assert ns.explicit_keys == set()
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose == 0
+
+
+def test_count_action_explicit():
+    """Repeated count flags increment and are marked explicit."""
+    p = TrackingArgumentParser()
+    p.add_argument("-v", "--verbose", action="count", default=0)
+    ns = p.parse_args(["-vv"])
+    assert ns.explicit_keys == {"verbose"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.verbose == 2
+
+
+def test_group_append_action_explicit():
+    """append args added via a group go through the same shadow-default path."""
+    p = TrackingArgumentParser()
+    g = p.add_argument_group("TestGroup")
+    g.add_argument("--tag", action="append")
+    ns = p.parse_args(["--tag", "a", "--tag", "b"])
+    assert ns.explicit_keys == {"tag"}
+    assert isinstance(ns, TrackingNamespace)
+    assert ns.tag == ["a", "b"]
+
+
+@pytest.mark.parametrize(
+    ("action", "expected"),
+    [
+        ("append", "store"),
+        ("extend", "store"),
+        ("append_const", "store_const"),
+        ("count", "store_const"),
+    ],
+)
+def test_build_shadow_kwargs_remaps_mutating_actions(action, expected):
+    """Mutating actions are remapped to a store-style action with an UNSET default."""
+    shadow = build_shadow_kwargs({"action": action, "const": "c"})
+    assert shadow["action"] == expected
+    assert shadow["default"] is UNSET
+
+
+def test_build_shadow_kwargs_count_gets_marker_const():
+    """count carries no const, so the remapped store_const gets a non-UNSET marker."""
+    shadow = build_shadow_kwargs({"action": "count"})
+    assert shadow["action"] == "store_const"
+    assert shadow["const"] is not UNSET
+
+
+def test_build_shadow_kwargs_leaves_other_actions_untouched():
+    """Non-mutating actions keep their action and only get the UNSET default."""
+    shadow = build_shadow_kwargs({"action": "store_true"})
+    assert shadow["action"] == "store_true"
+    assert shadow["default"] is UNSET
+
+
 def test_parse_known_args_tracking():
     """Ensure parse_known_args is also trackable"""
     p = TrackingArgumentParser()
@@ -460,7 +591,7 @@ def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages):
     ns = p.parse_args(["--max-num-seqs", "999"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -477,7 +608,7 @@ def test_omitted_default_not_in_runtime_overrides(mock_stages):
     ns = p.parse_args([])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -499,7 +630,7 @@ def test_config_file_args_reach_runtime_overrides(mock_stages):
         ns = p.parse_args(["--config", "fake.yaml"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -517,7 +648,7 @@ def test_per_stage_override_routes_correctly(mock_stages):
     ns = p.parse_args(["--stage-0-gpu-memory-utilization", "0.42"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
@@ -539,7 +670,7 @@ def test_explicit_args_omitted_from_yaml(mock_stages):
     ns = p.parse_args(["--enforce-eager"])
 
     explicit_kwargs = ns.get_explicit_kwargs_dict()
-    stages = StageConfigFactory._create_from_registry(
+    stages, _ = StageConfigFactory._create_from_registry(
         _TEST_MODEL,
         _TEST_PIPELINE,
         explicit_kwargs,
diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py
index a37c9b993..90cf3afe1 100644
--- a/tests/worker/test_gpu_ar_model_runner.py
+++ b/tests/worker/test_gpu_ar_model_runner.py
@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from types import SimpleNamespace
 
 import numpy as np
@@ -10,6 +11,7 @@ from vllm_omni.worker.gpu_ar_model_runner import (
     GPUARModelRunner,
     OmniAsyncGPUModelRunnerOutput,
 )
+from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -58,6 +60,101 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker():
     assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
 
 
+def test_runner_assisted_full_attention_metadata_request_is_opt_in():
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = object()
+    runner.scheduler_config = SimpleNamespace(max_num_seqs=16)
+
+    request = runner._get_runner_assisted_full_attention_metadata_request(
+        req_ids=["r1", "r2"],
+        num_reqs=2,
+        num_reqs_padded=4,
+        num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
+        num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32),
+        max_num_scheduled_tokens=1,
+    )
+
+    assert request is None
+
+
+def test_runner_assisted_full_attention_metadata_request_and_context_hooks():
+    calls = []
+
+    class Model:
+        def get_runner_assisted_full_attention_metadata_request(
+            self,
+            *,
+            req_ids: Sequence[str],
+            num_reqs: int,
+            num_scheduled_tokens: Sequence[int],
+            num_computed_tokens: Sequence[int],
+            max_num_scheduled_tokens: int,
+        ) -> RunnerAssistedFullAttentionMetadataRequest:
+            calls.append(
+                (
+                    "request",
+                    {
+                        "req_ids": list(req_ids),
+                        "num_reqs": num_reqs,
+                        "num_scheduled_tokens": [int(n) for n in num_scheduled_tokens],
+                        "num_computed_tokens": [int(n) for n in num_computed_tokens],
+                        "max_num_scheduled_tokens": max_num_scheduled_tokens,
+                    },
+                )
+            )
+            return RunnerAssistedFullAttentionMetadataRequest(
+                num_reqs_padded=12,
+                for_cudagraph_capture=True,
+            )
+
+        def set_runner_assisted_full_attention_metadata_context(
+            self,
+            *,
+            enabled: bool,
+            num_reqs: int = 0,
+        ) -> None:
+            calls.append(("context", {"enabled": enabled, "num_reqs": num_reqs}))
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.model = Model()
+    runner.scheduler_config = SimpleNamespace(max_num_seqs=8)
+
+    request = runner._get_runner_assisted_full_attention_metadata_request(
+        req_ids=["r1", "r2"],
+        num_reqs=2,
+        num_reqs_padded=4,
+        num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
+        num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32),
+        max_num_scheduled_tokens=1,
+    )
+    context_enabled = runner._set_runner_assisted_full_attention_metadata_context(
+        enabled=True,
+        num_reqs=2,
+    )
+    context_disabled = runner._set_runner_assisted_full_attention_metadata_context(enabled=False)
+
+    assert request == RunnerAssistedFullAttentionMetadataRequest(
+        num_reqs_padded=8,
+        for_cudagraph_capture=True,
+    )
+    assert context_enabled
+    assert context_disabled
+    assert calls == [
+        (
+            "request",
+            {
+                "req_ids": ["r1", "r2"],
+                "num_reqs": 2,
+                "num_scheduled_tokens": [1, 1],
+                "num_computed_tokens": [5, 6],
+                "max_num_scheduled_tokens": 1,
+            },
+        ),
+        ("context", {"enabled": True, "num_reqs": 2}),
+        ("context", {"enabled": False, "num_reqs": 0}),
+    ]
+
+
 def test_omni_async_gpu_model_runner_output_builds_lazily_once():
     async_output = object.__new__(OmniAsyncGPUModelRunnerOutput)
     calls = []
@@ -169,7 +266,6 @@ def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monke
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=["r2"],
-        seq_len=3,
         num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )
@@ -215,7 +311,6 @@ def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatc
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=6,
         num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long),
     )
@@ -269,7 +364,6 @@ def test_process_additional_information_uses_snapshot_request_order(monkeypatch)
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=3,
         num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )
@@ -330,7 +424,6 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=2,
         num_scheduled_tokens_np=np.array([2], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0], dtype=torch.long),
     )
@@ -342,6 +435,97 @@ def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch):
     assert output.multimodal_outputs is None
 
 
+def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_only(monkeypatch):
+    runner = _make_async_output_runner(engine_output_type="latent")
+    runner.model.omni_pooler_payload_include_hidden = False
+    runner._async_chunk = False
+    runner.requests = {"r1": object(), "r2": object(), "r3": object()}
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+
+    codes = torch.arange(48, dtype=torch.long).reshape(3, 16)
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=3,
+            num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
+        ),
+        hidden_states=torch.tensor([[1.0]]),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"codes": {"audio": codes}},
+        req_ids_output_copy=["r1", "r2", "r3"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
+        valid_sampled_token_ids=[[101], [102], [103]],
+        logprobs_lists=None,
+        prompt_logprobs_dict={},
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        ec_connector_output=None,
+        cudagraph_stats=None,
+        kv_extracted_req_ids=None,
+        num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
+    )
+
+    assert output.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1])
+    assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2])
+    assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3])
+
+
+def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monkeypatch):
+    """Thinker mm rows align to hidden_states.shape[0], not padded scheduled count."""
+    runner = _make_async_output_runner(engine_output_type="latent")
+    runner.model.omni_pooler_payload_include_hidden = False
+    runner._async_chunk = False
+    runner.requests = {"r1": object(), "r2": object(), "r3": object()}
+
+    monkeypatch.setattr(
+        GPUARModelRunner,
+        "_resolve_pooler_payload_req_ids",
+        lambda self, req_ids: ("latent", req_ids),
+    )
+    monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False)
+    monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None)
+    monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None)
+
+    layers = torch.arange(96, dtype=torch.long).reshape(3, 32)
+    output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot(
+        runner,
+        scheduler_output=SimpleNamespace(
+            total_num_scheduled_tokens=5,
+            num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1},
+        ),
+        hidden_states=torch.randn(3, 4),
+        staged_hidden_states_cpu=None,
+        multimodal_outputs={"hidden_states": {"layers": {0: layers}}},
+        req_ids_output_copy=["r1", "r2", "r3"],
+        req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2},
+        valid_sampled_token_ids=[[101], [102], [103]],
+        logprobs_lists=None,
+        prompt_logprobs_dict={},
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        ec_connector_output=None,
+        cudagraph_stats=None,
+        kv_extracted_req_ids=None,
+        num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32),
+        query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long),
+    )
+
+    assert output.inter_stage_outputs is not None
+    assert torch.equal(output.inter_stage_outputs[0]["hidden_states.layer_0"], layers[0:1])
+    assert torch.equal(output.inter_stage_outputs[1]["hidden_states.layer_0"], layers[1:2])
+    assert torch.equal(output.inter_stage_outputs[2]["hidden_states.layer_0"], layers[2:3])
+
+
 def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
     runner = _make_async_output_runner()
     runner.model.omni_pooler_payload_include_hidden = False
@@ -357,6 +541,47 @@ def test_async_snapshot_payload_omits_hidden_when_model_opts_out():
     assert payload["multimodal_outputs"]["codes"]["audio"].tolist() == [[1]]
 
 
+def test_runner_assisted_full_attention_metadata_refresh_pads_buffers():
+    class QueryStartLoc:
+        def __init__(self):
+            self.np = np.full(5, -1, dtype=np.int32)
+            self.copied = False
+
+        def copy_to_gpu(self):
+            self.copied = True
+
+    class BlockTable:
+        def __init__(self):
+            self.commits = []
+
+        def commit_block_table(self, num_reqs_padded):
+            self.commits.append(num_reqs_padded)
+
+    runner = object.__new__(GPUARModelRunner)
+    block_table = BlockTable()
+    runner.input_batch = SimpleNamespace(
+        num_computed_tokens_cpu=torch.tensor([10, 20, 99, 99], dtype=torch.int32),
+        block_table=block_table,
+    )
+    runner.optimistic_seq_lens_cpu = torch.zeros(4, dtype=torch.int32)
+    runner.seq_lens = torch.empty(4, dtype=torch.int32)
+    runner.query_pos = SimpleNamespace(np=np.empty(4, dtype=np.int32))
+    runner.query_start_loc = QueryStartLoc()
+    runner._get_cumsum_and_arange = lambda scheduled, _query_pos: np.cumsum(scheduled, dtype=np.int32)
+
+    runner._refresh_runner_assisted_full_attention_metadata_buffers(
+        num_reqs=2,
+        num_reqs_padded=4,
+        num_scheduled_tokens_np=np.array([1, 2], dtype=np.int32),
+    )
+
+    assert runner.optimistic_seq_lens_cpu.tolist() == [11, 22, 0, 0]
+    assert runner.seq_lens.tolist() == [11, 22, 0, 0]
+    assert runner.query_start_loc.np.tolist() == [0, 1, 3, 3, 3]
+    assert runner.query_start_loc.copied
+    assert block_table.commits == [4]
+
+
 @pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
 def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
     runner = object.__new__(GPUARModelRunner)
@@ -489,7 +714,6 @@ def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch
         ec_connector_output=None,
         cudagraph_stats=None,
         kv_extracted_req_ids=None,
-        seq_len=2,
         num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32),
         query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long),
     )
diff --git a/tests/worker/test_min_tokens_stop_ids_guard.py b/tests/worker/test_min_tokens_stop_ids_guard.py
new file mode 100644
index 000000000..5d7591c99
--- /dev/null
+++ b/tests/worker/test_min_tokens_stop_ids_guard.py
@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Guard against out-of-vocabulary stop ids in min-tokens masking (#4962).
+
+vLLM's input processor folds the stage tokenizer's EOS id into
+``SamplingParams.all_stop_token_ids``. For AR stages whose lm_head is
+narrower than the tokenizer vocabulary (codec talkers such as Qwen3-TTS:
+3072 logits vs text EOS 151645), any ``min_tokens >= 1`` makes
+``MinTokensLogitsProcessor`` write ``-inf`` at an out-of-range index and
+the engine dies with a CUDA device-side assert.
+"""
+
+import pytest
+import torch
+from vllm import SamplingParams
+from vllm.v1.sample.logits_processor import (
+    BatchUpdate,
+    LogitsProcessors,
+    MinTokensLogitsProcessor,
+)
+
+from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TALKER_VOCAB = 3072
+CODEC_EOS = 2150
+TEXT_EOS = 151645
+
+
+def _make_min_tokens_proc(stop_token_ids: list[int], eos_token_id: int | None) -> MinTokensLogitsProcessor:
+    """Build a MinTokensLogitsProcessor with one active request, mirroring
+    the production path (engine folds the tokenizer EOS via
+    ``update_from_generation_config``)."""
+    params = SamplingParams(min_tokens=2, stop_token_ids=stop_token_ids)
+    if eos_token_id is not None:
+        params.update_from_generation_config({}, eos_token_id)
+    proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False)
+    proc.update_state(
+        BatchUpdate(
+            batch_size=1,
+            removed=[],
+            added=[(0, params, None, [])],
+            moved=[],
+        )
+    )
+    return proc
+
+
+def test_oob_stop_id_crashes_without_guard():
+    """Document the failure mode: the folded text EOS is out of range for a
+    narrow codec head and index_put_ raises (device-side assert on CUDA)."""
+    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
+    logits = torch.zeros(1, TALKER_VOCAB)
+    with pytest.raises((IndexError, RuntimeError)):
+        proc.apply(logits)
+
+
+def test_guard_filters_oob_and_keeps_in_range_mask():
+    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+
+    logits = torch.zeros(1, TALKER_VOCAB)
+    out = proc.apply(logits)
+
+    assert out[0, CODEC_EOS] == -float("inf")
+    masked = (out == -float("inf")).nonzero()
+    assert masked.tolist() == [[0, CODEC_EOS]]
+
+
+def test_guard_noop_when_all_stop_ids_in_range():
+    """Full-width heads (e.g. CosyVoice3 talker) must be untouched: no
+    tensor rebuild, mask unchanged."""
+    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=None)
+    slice_before = proc.logits_slice
+
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+
+    assert proc.logits_slice is slice_before
+    out = proc.apply(torch.zeros(1, TALKER_VOCAB))
+    assert out[0, CODEC_EOS] == -float("inf")
+
+
+def test_guard_sanitizes_params_only_once():
+    """The stop-id set is shared with the request's SamplingParams, so the
+    second call finds nothing to drop (no repeated rebuilds/warnings)."""
+    proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS)
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+    slice_after_first = proc.logits_slice
+
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+
+    assert proc.logits_slice is slice_after_first
+
+
+def test_guard_handles_multiple_requests():
+    params_oob = SamplingParams(min_tokens=4, stop_token_ids=[CODEC_EOS])
+    params_oob.update_from_generation_config({}, TEXT_EOS)
+    params_ok = SamplingParams(min_tokens=4, stop_token_ids=[7])
+    proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False)
+    proc.update_state(
+        BatchUpdate(
+            batch_size=2,
+            removed=[],
+            added=[(0, params_oob, None, []), (1, params_ok, None, [])],
+            moved=[],
+        )
+    )
+
+    sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB)
+    out = proc.apply(torch.zeros(2, TALKER_VOCAB))
+
+    assert out[0, CODEC_EOS] == -float("inf")
+    assert out[1, 7] == -float("inf")
+    assert int((out == -float("inf")).sum()) == 2
diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py
index f562d18ce..48c0bc6de 100644
--- a/tests/worker/test_omni_gpu_model_runner.py
+++ b/tests/worker/test_omni_gpu_model_runner.py
@@ -81,6 +81,7 @@ class CaptureTalkerMTP(torch.nn.Module):
         top_k=None,
         top_p=None,
         generator=None,
+        generators=None,
     ):
         self.calls.append(
             {
@@ -90,6 +91,7 @@ class CaptureTalkerMTP(torch.nn.Module):
                 "top_k": top_k,
                 "top_p": top_p,
                 "generator": generator,
+                "generators": generators,
             }
         )
         codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64)
@@ -275,7 +277,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
     runner = _make_runner(req_ids=("r1",), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
         seed=42,
-        extra_args={"qwen3_tts_request_seed": 42},
+        extra_args={"tts_local_seed": 42},
     )
     runner.talker_mtp = CaptureTalkerMTP()
     runner.vllm_config = SimpleNamespace(
@@ -306,6 +308,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker
             "top_k": 9,
             "top_p": 0.55,
             "generator": runner.talker_mtp.calls[0]["generator"],
+            "generators": None,
         }
     ]
     assert runner.talker_mtp.calls[0]["generator"] is not None
@@ -319,11 +322,11 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
     runner.requests["r1"].sampling_params = SimpleNamespace(
         seed=11,
-        extra_args={"qwen3_tts_request_seed": 11},
+        extra_args={"tts_local_seed": 11},
     )
     runner.requests["r2"].sampling_params = SimpleNamespace(
         seed=22,
-        extra_args={"qwen3_tts_request_seed": 22},
+        extra_args={"tts_local_seed": 22},
     )
     runner.talker_mtp = CaptureTalkerMTP()
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
@@ -350,6 +353,57 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
     assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds)
 
 
+def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypatch):
+    """Models with talker_mtp_accepts_per_row_generators get one batched call (#4883)."""
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
+    runner.requests["r1"].sampling_params = SimpleNamespace(
+        seed=11,
+        extra_args={"tts_local_seed": 11},
+    )
+    runner.requests["r2"].sampling_params = SimpleNamespace(
+        seed=22,
+        extra_args={"tts_local_seed": 22},
+    )
+    runner.talker_mtp = CaptureTalkerMTP()
+    runner.model = SimpleNamespace(
+        talker_mtp_output_key=("codes", "audio"),
+        talker_mtp_accepts_per_row_generators=True,
+    )
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
+
+    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
+        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
+        return (False, batch_desc, None, None, None)
+
+    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
+
+    inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+
+    # One batched call with distinct per-row generators, not two scalar calls.
+    assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
+    row_generators = runner.talker_mtp.calls[0]["generators"]
+    assert runner.talker_mtp.calls[0]["generator"] is None
+    assert len(row_generators) == 2
+    assert all(generator is not None for generator in row_generators)
+    assert row_generators[0] is not row_generators[1]
+
+    # The per-request generator stream persists across steps...
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
+    assert runner.talker_mtp.calls[1]["generators"][0] is row_generators[0]
+    assert runner.talker_mtp.calls[1]["generators"][1] is row_generators[1]
+
+    # ...and is evicted once its request finishes.
+    del runner.requests["r2"]
+    OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
+    assert set(runner._talker_mtp_generators) == {"r1"}
+    assert runner.talker_mtp.calls[2]["generator"] is row_generators[0]
+
+
 def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
     """Validate that _update_intermediate_buffer writes to model_intermediate_buffer
     (forward path) and mirrors to additional_information_cpu setattr (backward compat)."""
diff --git a/tools/nightly/run_nightly_jobs.sh b/tools/nightly/run_nightly_jobs.sh
index b115a9ee2..7134a36a1 100644
--- a/tools/nightly/run_nightly_jobs.sh
+++ b/tools/nightly/run_nightly_jobs.sh
@@ -9,7 +9,7 @@
 # python tools/nightly/generate_nightly_perf_excel.py runs (even if some perf jobs failed;
 # the script only aggregates JSON already present under tests/dfx/perf/results).
 # Excel and generate_nightly_perf_excel.log are written under $REPO_ROOT/logs/;
-# per-job tee logs use LOG_DIR (see --log-dir). Then remaining jobs run.
+# per-job tee logs use LOG_DIR (see --log-dir; default is timestamped under $REPO_ROOT/logs/).
 #
 # Test kind (--test-type / TEST_TYPE), multiple allowed (OR semantics):
 #   Repeat the flag and/or use comma-separated values, e.g.
@@ -21,6 +21,8 @@
 #   function — label has neither "Perf Test" nor "Accuracy Test" (incl. Doc, Multi-Replica, etc.)
 #   stability— fixed dfx stability scripts under tests/dfx/stability/scripts/ (see below)
 #   local    — pytest -sv -m "<MODEL_TYPE markers> and local_model" from repo root (no YAML step extract)
+#              When LABEL_SUBSTR is set, also runs matching tests/**/test_*.py and perf JSON configs under
+#              tests/dfx/perf/tests/*.json via run_benchmark.py / run_diffusion_benchmark.py.
 #   all      — no test-kind filter for YAML steps (any leaf step with pytest in test-nightly.yml)
 #
 # Model area (--model-type / MODEL_TYPE), multiple allowed (OR semantics):
@@ -36,10 +38,12 @@
 #   local (when included in TEST_TYPE):
 #     From repo root: pytest -sv -m "<markers> and local_model" (markers from MODEL_TYPE: omni, tts,
 #     diffusion; all → "(omni or tts or diffusion) and local_model"). Not filtered by nightly YAML.
+#     LABEL_SUBSTR: if set, restrict to tests/**/test_*.py basenames and tests/dfx/perf/tests/*.json
+#                   whose filename contains the substring (benchmark runner chosen by JSON family).
 #
 #   stability (when included in TEST_TYPE):
-#     From repo root: pytest -s -v tests/dfx/stability/scripts/test_stability_*.py
-#     model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22; all → all four
+#     From repo root: pytest -s -v --run-level full_model tests/dfx/stability/scripts/test_stability_*.py
+#     model_type: omni → qwen3_omni; tts → qwen3_tts + voxcpm2; diffusion → qwen_image + wan22 + hunyuan_image; all → all six
 #     LABEL_SUBSTR: if set, script path / job key / filename must contain it
 #
 # Requirements: bash, python3, PyYAML (pip install pyyaml)
@@ -58,11 +62,20 @@
 # Optional environment:
 #   REPO_ROOT     - vllm-omni root (working directory for pytest); see above
 #   YML           - path to test-nightly.yml (default: $REPO_ROOT/.buildkite/test-nightly.yml)
-#   LOG_DIR       - logs + generated job scripts (default: $REPO_ROOT/logs/nightly_jobs)
+#   LOG_DIR       - logs + generated job scripts; when unset, a timestamped directory under
+#                   $REPO_ROOT/logs/ is created:
+#                     nightly_jobs_YYYYMMDD-HHMMSS       (default / YAML nightly steps)
+#                     nightly_local_jobs_YYYYMMDD-HHMMSS (--test-type local only)
+#                     nightly_stability_jobs_YYYYMMDD-HHMMSS (--test-type stability only)
+#                   per-job *.log plus timing_summary.log after the run;
+#                   perf JSON from tests/dfx/perf/results/ (produced this run) -> $LOG_DIR/perf_results/
 #   TEST_TYPE     - comma-separated and/or repeated flags (default: all); see above
 #   MODEL_TYPE    - comma-separated and/or repeated flags (default: all); see above
-#   LABEL_SUBSTR  - YAML mode: substring of Buildkite step label; stability: substring of path/key/filename
+#   LABEL_SUBSTR  - YAML mode: substring of Buildkite step label; stability: substring of path/key/filename;
+#                   local: substring of test_*.py basename or tests/dfx/perf/tests/*.json filename
 #   DRY_RUN=1     - print extracted commands only; do not write scripts or run pytest
+#   RUN_JOB_TIMEOUT_KILL_AFTER - seconds after SIGTERM before SIGKILL on inline pytest
+#                              timeout (default: 60; baked into generated job scripts)
 #
 set -euo pipefail
 
@@ -212,6 +225,28 @@ done
 TEST_TYPE="$(_finalize_test_type_csv)"
 MODEL_TYPE="$(_finalize_model_type_csv)"
 
+# Default log directory basename: date + local timestamp; prefix depends on TEST_TYPE.
+_nightly_log_dir_basename() {
+  local ts has_local=0 has_stability=0 has_yaml=0 _t
+  ts="$(date +%Y%m%d-%H%M%S)"
+  IFS=',' read -ra _TTOK <<< "${TEST_TYPE}"
+  for _t in "${_TTOK[@]}"; do
+    _t="$(printf '%s' "${_t}" | tr '[:upper:]' '[:lower:]' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')"
+    case "${_t}" in
+      local) has_local=1 ;;
+      stability) has_stability=1 ;;
+      perf | acc | function | all) has_yaml=1 ;;
+    esac
+  done
+  if [[ "${has_local}" -eq 1 && "${has_yaml}" -eq 0 && "${has_stability}" -eq 0 ]]; then
+    printf 'nightly_local_jobs_%s' "${ts}"
+  elif [[ "${has_stability}" -eq 1 && "${has_yaml}" -eq 0 && "${has_local}" -eq 0 ]]; then
+    printf 'nightly_stability_jobs_%s' "${ts}"
+  else
+    printf 'nightly_jobs_%s' "${ts}"
+  fi
+}
+
 BUILDKITE_REL=".buildkite/test-nightly.yml"
 
 _find_repo_containing_nightly() {
@@ -269,7 +304,11 @@ else
   YML="${REPO_ROOT}/${BUILDKITE_REL}"
 fi
 
-LOG_DIR="${LOG_DIR:-${REPO_ROOT}/logs/nightly_jobs}"
+if [[ -z "${LOG_DIR:-}" ]]; then
+  LOG_DIR="${REPO_ROOT}/logs/$(_nightly_log_dir_basename)"
+fi
+
+echo "Log directory: ${LOG_DIR}" >&2
 
 # Require nightly YAML only when at least one non-stability test kind needs it (validated in Python too).
 _needs_nightly_yml=0
@@ -310,6 +349,7 @@ if [[ "${DRY_RUN}" != "1" ]]; then
     rm -f "${_stale_logs[@]}"
   fi
   rm -f "${LOG_DIR}/jobs/.perf_job_keys"
+  rm -f "${LOG_DIR}/jobs/.job_timeouts"
   rm -f "${REPO_ROOT}/logs/nightly_perf_manual.xlsx"
   rm -f "${LOG_DIR}/nightly_perf_manual.xlsx"
 fi
@@ -320,6 +360,7 @@ from __future__ import annotations
 
 import os
 import re
+import shlex
 import stat
 import sys
 from pathlib import Path
@@ -491,12 +532,60 @@ def job_key_from_step(step: dict, label: str) -> str:
     return slug
 
 
+def step_timeout_minutes(step: dict) -> int | None:
+    raw = step.get("timeout_in_minutes")
+    if raw is None:
+        return None
+    try:
+        minutes = int(raw)
+    except (TypeError, ValueError):
+        return None
+    return minutes if minutes > 0 else None
+
+
+_INLINE_TIMEOUT_PREFIX = re.compile(r"^timeout\s+(?:-\S+\s+)*\S+\s+", re.IGNORECASE)
+
+
+def timeout_kill_after_seconds() -> int:
+    raw = os.environ.get("RUN_JOB_TIMEOUT_KILL_AFTER", "60")
+    try:
+        seconds = int(raw)
+    except (TypeError, ValueError):
+        return 60
+    return seconds if seconds > 0 else 60
+
+
+def prepend_timeout_to_pytest(pytest_line: str, timeout_min: int | None) -> str:
+    """Prepend GNU timeout directly before pytest (Buildkite inline style)."""
+    if timeout_min is None:
+        return pytest_line.strip()
+    line = _INLINE_TIMEOUT_PREFIX.sub("", pytest_line.strip())
+    kill_after = timeout_kill_after_seconds()
+    return (
+        f"timeout --foreground --verbose --kill-after={kill_after} "
+        f"{timeout_min}m {line}"
+    )
+
+
+def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) -> None:
+    manifest_path = jobs_dir / ".job_timeouts"
+    if job_timeouts:
+        manifest_path.write_text(
+            "\n".join(f"{k}={v}" for k, v in sorted(job_timeouts.items())) + "\n",
+            encoding="utf-8",
+        )
+    elif manifest_path.is_file():
+        manifest_path.unlink()
+
+
 # Fixed stability scripts (when stability is selected); MODEL_TYPES narrows which run.
 STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [
     ("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)),
     ("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)),
+    ("stability_voxcpm2", "tests/dfx/stability/scripts/test_stability_voxcpm2.py", ("tts",)),
     ("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)),
     ("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)),
+    ("stability_hunyuan_image", "tests/dfx/stability/scripts/test_stability_hunyuan_image.py", ("diffusion",)),
 ]
 
 
@@ -533,14 +622,148 @@ def local_pytest_marker_expr(model_types: list[str]) -> str:
     return f"({' or '.join(families)}) and local_model"
 
 
-def run_local_mode(jobs_dir: Path, model_types: list[str]) -> int:
-    """Single job: pytest -sv -m '<model markers> and local_model' (testpaths from pytest config)."""
+def local_test_files_by_filename(substr: str) -> list[str]:
+    """Return repo-relative posix paths under tests/ whose basename contains substr."""
+    tests_root = REPO_ROOT / "tests"
+    if not tests_root.is_dir():
+        return []
+    matches: list[str] = []
+    for path in sorted(tests_root.rglob("test_*.py")):
+        if substr in path.name:
+            matches.append(path.relative_to(REPO_ROOT).as_posix())
+    return matches
+
+
+PERF_TESTS_REL = Path("tests/dfx/perf/tests")
+RUN_BENCHMARK_REL = Path("tests/dfx/perf/scripts/run_benchmark.py")
+RUN_DIFFUSION_BENCHMARK_REL = Path("tests/dfx/perf/scripts/run_diffusion_benchmark.py")
+TTS_PERF_JSON_HINTS = ("test_tts", "voxcpm", "higgs_audio")
+
+
+def perf_json_model_family(json_basename: str) -> str:
+    """Classify a perf JSON config as omni, tts, or diffusion (mirrors nightly YAML runners)."""
+    name = json_basename.lower()
+    if name.startswith("test_qwen3_omni"):
+        return "omni"
+    if any(hint in name for hint in TTS_PERF_JSON_HINTS):
+        return "tts"
+    return "diffusion"
+
+
+def perf_json_runner(json_basename: str) -> Path:
+    if perf_json_model_family(json_basename) in ("omni", "tts"):
+        return RUN_BENCHMARK_REL
+    return RUN_DIFFUSION_BENCHMARK_REL
+
+
+def perf_json_matches_model_type(json_basename: str, model_types: list[str]) -> bool:
+    if "all" in model_types:
+        return True
+    return perf_json_model_family(json_basename) in model_types
+
+
+def local_perf_json_configs(substr: str, model_types: list[str]) -> list[tuple[str, str]]:
+    """Return (json_rel_posix, runner_rel_posix) for perf configs whose filename contains substr."""
+    perf_dir = REPO_ROOT / PERF_TESTS_REL
+    if not perf_dir.is_dir():
+        return []
+    matches: list[tuple[str, str]] = []
+    for path in sorted(perf_dir.glob("*.json")):
+        if substr not in path.name:
+            continue
+        if not perf_json_matches_model_type(path.name, model_types):
+            continue
+        runner = perf_json_runner(path.name)
+        if not (REPO_ROOT / runner).is_file():
+            print(f"# skip (missing runner): {runner}", file=sys.stderr)
+            continue
+        matches.append(
+            (
+                path.relative_to(REPO_ROOT).as_posix(),
+                runner.as_posix(),
+            )
+        )
+    return matches
+
+
+def _local_job_slug(label: str) -> str:
+    slug = re.sub(r"[^\w\-.]+", "_", label, flags=re.UNICODE)
+    slug = re.sub(r"_+", "_", slug).strip("_")
+    return slug or "filter"
+
+
+def run_local_mode(
+    jobs_dir: Path,
+    model_types: list[str],
+    perf_job_keys: list[str],
+) -> int:
+    """pytest -sv -m '<model markers> and local_model'; optional LABEL_SUBSTR filters."""
     marker_expr = local_pytest_marker_expr(model_types)
-    key = "local_pytest"
+    matched = 0
+
+    if LABEL_SUBSTR:
+        rel_paths = local_test_files_by_filename(LABEL_SUBSTR)
+        perf_pairs = local_perf_json_configs(LABEL_SUBSTR, model_types)
+
+        if rel_paths:
+            paths_arg = " ".join(shlex.quote(p) for p in rel_paths)
+            pytest_line = f'pytest -sv -m "{marker_expr}" {paths_arg}'
+            slug = _local_job_slug(LABEL_SUBSTR)
+            key = f"local_pytest_{slug}"
+            header = (
+                f"# Local marker tests — MODEL_TYPE={','.join(model_types)}, "
+                f"LABEL_SUBSTR={LABEL_SUBSTR!r} (test_*.py), files: {', '.join(rel_paths)}"
+            )
+            script_lines = [
+                "#!/usr/bin/env bash",
+                header,
+                "set -euo pipefail",
+                f'cd "{REPO_ROOT}"',
+                pytest_line,
+            ]
+            _write_job_script(key, script_lines, jobs_dir)
+            matched += 1
+
+        for json_rel, runner_rel in perf_pairs:
+            json_name = Path(json_rel).name
+            pytest_line = (
+                f'pytest -s -v -m "{marker_expr}" '
+                f"{shlex.quote(runner_rel)} "
+                f"--test-config-file {shlex.quote(json_rel)}"
+            )
+            slug = _local_job_slug(Path(json_rel).stem)
+            key = f"local_perf_{slug}"
+            header = (
+                f"# Local perf benchmark — MODEL_TYPE={','.join(model_types)}, "
+                f"config={json_name!r}, runner={Path(runner_rel).name}"
+            )
+            script_lines = [
+                "#!/usr/bin/env bash",
+                header,
+                "set -euo pipefail",
+                f'cd "{REPO_ROOT}"',
+                pytest_line,
+            ]
+            _write_job_script(key, script_lines, jobs_dir)
+            if key not in perf_job_keys:
+                perf_job_keys.append(key)
+            matched += 1
+
+        if matched == 0:
+            print(
+                f"# skip local: no tests/**/test_*.py or "
+                f"tests/dfx/perf/tests/*.json filename matches "
+                f"LABEL_SUBSTR={LABEL_SUBSTR!r}",
+                file=sys.stderr,
+            )
+        return matched
+
     pytest_line = f'pytest -sv -m "{marker_expr}"'
+    key = "local_pytest"
+    header = f"# Local marker tests — MODEL_TYPE={','.join(model_types)}"
     script_lines = [
         "#!/usr/bin/env bash",
-        f"# Local marker tests — MODEL_TYPE={','.join(model_types)}",
+        header,
         "set -euo pipefail",
         f'cd "{REPO_ROOT}"',
         pytest_line,
@@ -566,7 +789,7 @@ def run_stability_mode(jobs_dir: Path, model_types: list[str]) -> int:
         if not rel_path.is_file():
             print(f"# skip (missing file): {rel_path}", file=sys.stderr)
             continue
-        pytest_line = f"pytest -s -v {rel_posix}"
+        pytest_line = f"pytest -s -v --run-level full_model {rel_posix}"
         matched += 1
         script_lines = [
             "#!/usr/bin/env bash",
@@ -594,13 +817,15 @@ def main() -> None:
 
     matched_stability = 0
     matched_yaml = 0
+    matched_local = 0
     perf_job_keys: list[str] = []
+    job_timeouts: dict[str, int] = {}
 
     if want_stability:
         matched_stability = run_stability_mode(jobs_dir, model_types)
 
     if want_local:
-        run_local_mode(jobs_dir, model_types)
+        matched_local = run_local_mode(jobs_dir, model_types, perf_job_keys)
 
     if needs_yaml:
         if not YML.is_file():
@@ -634,8 +859,12 @@ def main() -> None:
                 "set -euo pipefail",
                 f'cd "{REPO_ROOT}"',
             ]
+            timeout_min = step_timeout_minutes(step)
+            if timeout_min is not None:
+                script_lines.append(f"# Buildkite timeout_in_minutes: {timeout_min}")
+                job_timeouts[key] = timeout_min
             script_lines.extend(exports)
-            script_lines.extend(pys)
+            script_lines.extend(prepend_timeout_to_pytest(p, timeout_min) for p in pys)
             _write_job_script(key, script_lines, jobs_dir)
             if "Perf Test" in label and key not in perf_job_keys:
                 perf_job_keys.append(key)
@@ -646,6 +875,7 @@ def main() -> None:
             manifest_path.write_text("\n".join(perf_job_keys) + "\n", encoding="utf-8")
         elif manifest_path.is_file():
             manifest_path.unlink()
+        _write_job_timeouts_manifest(jobs_dir, job_timeouts)
 
     errs: list[str] = []
     if want_stability and matched_stability == 0:
@@ -653,6 +883,17 @@ def main() -> None:
             f"No stability jobs matched TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r} "
             f"LABEL_SUBSTR={LABEL_SUBSTR!r} under {REPO_ROOT}"
         )
+    if want_local and matched_local == 0:
+        if LABEL_SUBSTR:
+            errs.append(
+                f"No local jobs matched LABEL_SUBSTR={LABEL_SUBSTR!r} "
+                f"(tests/**/test_*.py or tests/dfx/perf/tests/*.json) "
+                f"under {REPO_ROOT}"
+            )
+        else:
+            errs.append(
+                f"Local mode produced no jobs TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r}"
+            )
     if needs_yaml and matched_yaml == 0:
         errs.append(
             f"No YAML steps matched TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r} "
@@ -671,29 +912,70 @@ if [[ "${DRY_RUN}" == "1" ]]; then
   exit 0
 fi
 
+# shellcheck source=tools/run_jobs_common.sh
+source "${SCRIPT_DIR}/../run_jobs_common.sh"
+
+PERF_RESULTS_SRC="${REPO_ROOT}/tests/dfx/perf/results"
+
+# Copy perf benchmark JSON written under tests/dfx/perf/results/ during this run into LOG_DIR.
+_collect_perf_result_jsons() {
+  local since_epoch="${1:?}"
+  local dest="${LOG_DIR}/perf_results"
+
+  if [[ ! -d "${PERF_RESULTS_SRC}" ]]; then
+    return 0
+  fi
+
+  COLLECT_SINCE_EPOCH="${since_epoch}" \
+  python3 - <<'PY'
+import os
+import shutil
+import sys
+from pathlib import Path
+
+repo = Path(os.environ["REPO_ROOT"]).resolve()
+src = repo / "tests" / "dfx" / "perf" / "results"
+dest = Path(os.environ["LOG_DIR"]).resolve() / "perf_results"
+since = float(os.environ["COLLECT_SINCE_EPOCH"])
+
+if not src.is_dir():
+    sys.exit(0)
+
+copied = 0
+for path in sorted(src.rglob("*.json")):
+    if not path.is_file():
+        continue
+    try:
+        if path.stat().st_mtime < since:
+            continue
+    except OSError:
+        continue
+    rel = path.relative_to(src)
+    out = dest / rel
+    out.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(path, out)
+    copied += 1
+
+if copied:
+    print(f"Collected {copied} perf result JSON file(s) -> {dest}", file=sys.stderr)
+PY
+}
+
 # Run jobs: YAML perf steps first (if manifest exists), then generate Excel (always when
 # manifest non-empty, even if some perf jobs failed; Excel uses whatever JSON exists on disk),
 # then remaining jobs. Paths align with nightly Buildkite (tests/dfx/perf/results).
 run_generated_jobs_with_tee() {
   set -o pipefail
   local any_fail=0
+  local _run_start
   local perf_list="${LOG_DIR}/jobs/.perf_job_keys"
   local -a perf_keys=()
+  _run_jobs_reset_timing
+  _run_start="$(_run_jobs_epoch_seconds)"
   if [[ -f "${perf_list}" ]] && [[ -s "${perf_list}" ]]; then
     mapfile -t perf_keys < "${perf_list}"
   fi
 
-  _run_one_job() {
-    local _job="$1"
-    local base out job_status
-    base="$(basename "${_job}" .sh)"
-    out="${LOG_DIR}/${base}.log"
-    echo "==> ${_job}  (tee ${out})" >&2
-    (cd "${REPO_ROOT}" && bash "${_job}") 2>&1 | tee "${out}"
-    job_status="${PIPESTATUS[0]}"
-    return "${job_status}"
-  }
-
   local k
   for k in "${perf_keys[@]}"; do
     [[ -n "${k}" ]] || continue
@@ -702,11 +984,11 @@ run_generated_jobs_with_tee() {
       any_fail=1
       continue
     fi
-    _run_one_job "${LOG_DIR}/jobs/${k}.sh" || any_fail=1
+    _run_one_job_with_timing "${LOG_DIR}/jobs/${k}.sh" || any_fail=1
   done
 
   if ((${#perf_keys[@]})); then
-    local excel_out excel_log excel_status excel_repo_logs
+    local excel_out excel_log excel_status excel_repo_logs excel_start excel_end excel_elapsed
     excel_repo_logs="${REPO_ROOT}/logs"
     mkdir -p "${excel_repo_logs}"
     excel_out="${excel_repo_logs}/nightly_perf_$(date -u +%Y%m%d-%H%M%S).xlsx"
@@ -716,6 +998,7 @@ run_generated_jobs_with_tee() {
         "(report reflects JSON already under tests/dfx/perf/results)." >&2
     fi
     echo "==> python3 tools/nightly/generate_nightly_perf_excel.py -> ${excel_out}  (tee ${excel_log})" >&2
+    excel_start="$(_run_jobs_epoch_seconds)"
     (
       cd "${REPO_ROOT}" && python3 tools/nightly/generate_nightly_perf_excel.py \
         --input-dir "${REPO_ROOT}/tests/dfx/perf/results" \
@@ -723,7 +1006,13 @@ run_generated_jobs_with_tee() {
         --output-file "${excel_out}"
     ) 2>&1 | tee "${excel_log}"
     excel_status="${PIPESTATUS[0]}"
-    if [[ "${excel_status}" -ne 0 ]]; then
+    excel_end="$(_run_jobs_epoch_seconds)"
+    excel_elapsed=$((excel_end - excel_start))
+    _run_jobs_record_timing "generate_nightly_perf_excel" "${excel_elapsed}" "${excel_status}"
+    if [[ "${excel_status}" -eq 0 ]]; then
+      echo "    finished in $(_run_jobs_format_duration "${excel_elapsed}")" >&2
+    else
+      echo "    failed after $(_run_jobs_format_duration "${excel_elapsed}") (exit ${excel_status})" >&2
       any_fail=1
       echo "generate_nightly_perf_excel.py failed (exit ${excel_status}). See ${excel_log}" >&2
     fi
@@ -739,15 +1028,16 @@ run_generated_jobs_with_tee() {
   for _job in "${LOG_DIR}/jobs"/*.sh; do
     base="$(basename "${_job}" .sh)"
     [[ ${_perf_done["${base}"]+isset} ]] && continue
-    _run_one_job "${_job}" || any_fail=1
+    _run_one_job_with_timing "${_job}" || any_fail=1
   done
   shopt -u nullglob
 
+  _collect_perf_result_jsons "${_run_start}"
+
+  _run_jobs_print_timing_summary "${_run_start}" "${any_fail}"
   if [[ "${any_fail}" -ne 0 ]]; then
-    echo "One or more selected jobs failed. See logs under ${LOG_DIR}." >&2
     return 1
   fi
-  echo "All selected jobs finished OK. Logs: ${LOG_DIR}/*.log" >&2
   return 0
 }
 
diff --git a/tools/run_jobs_common.sh b/tools/run_jobs_common.sh
new file mode 100644
index 000000000..ddc1794dc
--- /dev/null
+++ b/tools/run_jobs_common.sh
@@ -0,0 +1,653 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Shared helpers for tools/run_*_jobs.sh (source this file; do not execute).
+#
+# Contents:
+#   1) Timing / tee runners used by ready, merge, and nightly.
+#   2) Ready/merge CLI + YAML extract entrypoint: run_yaml_ci_jobs_main
+#      (set BUILDKITE_REL and DEFAULT_LOG_SUBDIR in the entry script, then call it).
+#
+# Job timeouts: generated job scripts embed GNU timeout on the pytest line itself
+# (see prepend_timeout_to_pytest in the Python heredoc below). This runner only
+# executes bash job.sh and records exit 124 as TIMED OUT via jobs/.job_timeouts
+# metadata.
+
+if [[ -n "${_RUN_JOBS_COMMON_LOADED:-}" ]]; then
+  return 0 2>/dev/null || exit 0
+fi
+_RUN_JOBS_COMMON_LOADED=1
+
+if [[ -z "${SCRIPT_DIR:-}" ]]; then
+  SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+fi
+
+# ---------------------------------------------------------------------------
+# Timing / tee helpers (ready, merge, nightly)
+# ---------------------------------------------------------------------------
+
+RUN_JOB_NAMES=()
+RUN_JOB_SECONDS=()
+RUN_JOB_STATUSES=()
+
+_run_jobs_reset_timing() {
+  RUN_JOB_NAMES=()
+  RUN_JOB_SECONDS=()
+  RUN_JOB_STATUSES=()
+}
+
+_run_jobs_epoch_seconds() {
+  date +%s
+}
+
+_run_jobs_format_duration() {
+  local total="${1}"
+  local h m s
+  if [[ "${total}" -lt 60 ]]; then
+    printf '%ss' "${total}"
+    return 0
+  fi
+  m=$((total / 60))
+  s=$((total % 60))
+  if [[ "${m}" -lt 60 ]]; then
+    if [[ "${s}" -eq 0 ]]; then
+      printf '%dm' "${m}"
+    else
+      printf '%dm %ss' "${m}" "${s}"
+    fi
+    return 0
+  fi
+  h=$((m / 60))
+  m=$((m % 60))
+  if [[ "${m}" -eq 0 && "${s}" -eq 0 ]]; then
+    printf '%dh' "${h}"
+  elif [[ "${s}" -eq 0 ]]; then
+    printf '%dh %dm' "${h}" "${m}"
+  else
+    printf '%dh %dm %ss' "${h}" "${m}" "${s}"
+  fi
+}
+
+_run_jobs_record_timing() {
+  local name="${1}"
+  local seconds="${2}"
+  local status="${3}"
+  RUN_JOB_NAMES+=("${name}")
+  RUN_JOB_SECONDS+=("${seconds}")
+  RUN_JOB_STATUSES+=("${status}")
+}
+
+_run_jobs_lookup_timeout_minutes() {
+  local base="${1}"
+  local manifest="${LOG_DIR}/jobs/.job_timeouts"
+  local line key mins
+  [[ -f "${manifest}" ]] || return 1
+  while IFS='=' read -r key mins; do
+    [[ -n "${key}" ]] || continue
+    if [[ "${key}" == "${base}" ]]; then
+      printf '%s' "${mins}"
+      return 0
+    fi
+  done < "${manifest}"
+  return 1
+}
+
+_run_one_job_with_timing() {
+  local _job="$1"
+  local base out job_status start end elapsed timeout_min=""
+  base="$(basename "${_job}" .sh)"
+  out="${LOG_DIR}/${base}.log"
+  if timeout_min="$(_run_jobs_lookup_timeout_minutes "${base}")"; then
+    echo "==> ${_job}  (tee ${out}, pytest inline timeout ${timeout_min}m)" >&2
+  else
+    echo "==> ${_job}  (tee ${out})" >&2
+  fi
+  start="$(_run_jobs_epoch_seconds)"
+  (cd "${REPO_ROOT}" && bash "${_job}") 2>&1 | tee "${out}"
+  job_status="${PIPESTATUS[0]}"
+  end="$(_run_jobs_epoch_seconds)"
+  elapsed=$((end - start))
+  _run_jobs_record_timing "${base}" "${elapsed}" "${job_status}"
+  if [[ "${job_status}" -eq 0 ]]; then
+    echo "    finished in $(_run_jobs_format_duration "${elapsed}")" >&2
+  elif [[ "${job_status}" -eq 124 && -n "${timeout_min}" ]]; then
+    echo "    timed out (inline timeout ${timeout_min}m on pytest, exit 124)" >&2
+  elif [[ "${job_status}" -eq 124 ]]; then
+    echo "    timed out (exit 124)" >&2
+  else
+    echo "    failed after $(_run_jobs_format_duration "${elapsed}") (exit ${job_status})" >&2
+  fi
+  return "${job_status}"
+}
+
+_run_jobs_print_timing_summary() {
+  local wall_start="${1:-}"
+  local any_fail="${2:-0}"
+  local summary_path="${LOG_DIR}/timing_summary.log"
+  local -a lines=()
+  local i name secs status status_str failed_count=0
+  local wall_elapsed job_count total_elapsed=0
+
+  lines+=("=== Job timing summary ===")
+  for i in "${!RUN_JOB_NAMES[@]}"; do
+    name="${RUN_JOB_NAMES[$i]}"
+    secs="${RUN_JOB_SECONDS[$i]}"
+    status="${RUN_JOB_STATUSES[$i]}"
+    if [[ "${status}" -eq 0 ]]; then
+      status_str="OK"
+    elif [[ "${status}" -eq 124 ]]; then
+      status_str="TIMED OUT"
+      failed_count=$((failed_count + 1))
+    else
+      status_str="FAILED (exit ${status})"
+      failed_count=$((failed_count + 1))
+    fi
+    lines+=("  ${name}  $(_run_jobs_format_duration "${secs}")  ${status_str}")
+  done
+
+  job_count="${#RUN_JOB_NAMES[@]}"
+  if [[ -n "${wall_start}" ]]; then
+    wall_elapsed=$(( $(_run_jobs_epoch_seconds) - wall_start ))
+    lines+=("Total wall time: $(_run_jobs_format_duration "${wall_elapsed}") (${job_count} jobs)")
+  elif [[ "${job_count}" -gt 0 ]]; then
+    for secs in "${RUN_JOB_SECONDS[@]}"; do
+      total_elapsed=$((total_elapsed + secs))
+    done
+    lines+=("Total job time: $(_run_jobs_format_duration "${total_elapsed}") (${job_count} jobs)")
+  else
+    lines+=("Total wall time: 0s (0 jobs)")
+  fi
+
+  if [[ "${failed_count}" -gt 0 ]]; then
+    lines+=("Failed jobs: ${failed_count}/${job_count}")
+  fi
+
+  if [[ "${any_fail}" -ne 0 ]]; then
+    lines+=("Result: one or more jobs failed. See logs under ${LOG_DIR}.")
+  else
+    lines+=("Result: all jobs finished OK. Logs: ${LOG_DIR}/*.log")
+  fi
+
+  {
+    for line in "${lines[@]}"; do
+      printf '%s\n' "${line}"
+    done
+  } | tee "${summary_path}" >&2
+}
+
+# ---------------------------------------------------------------------------
+# Ready / merge: CLI + YAML extract + run (tools/run_ready_jobs.sh,
+# tools/run_merge_jobs.sh). Nightly does not use this entrypoint.
+#
+# Entry scripts must set before calling run_yaml_ci_jobs_main:
+#   BUILDKITE_REL      - e.g. .buildkite/test-ready.yml
+#   DEFAULT_LOG_SUBDIR - e.g. ready_jobs
+# ---------------------------------------------------------------------------
+
+_run_yaml_ci_jobs_usage() {
+  # Print the entry script header (caller is run_ready_jobs.sh / run_merge_jobs.sh).
+  sed -n '2,38p' "$0" | sed 's/^# \{0,1\}//'
+}
+
+_run_jobs_split_append_csv_array() {
+  local -n _arr="${1}"
+  local _raw="${2}"
+  local _ifs="${IFS}"
+  local _part
+  IFS=','
+  for _part in ${_raw}; do
+    IFS="${_ifs}"
+    _part="$(printf '%s' "${_part}" | tr '[:upper:]' '[:lower:]' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')"
+    [[ -n "${_part}" ]] || continue
+    _arr+=("${_part}")
+  done
+}
+
+_run_jobs_finalize_model_type_csv() {
+  if [[ "${MODEL_TYPE_FROM_CLI}" -eq 1 ]]; then
+    if ((${#MODEL_TYPE_CLI_PARTS[@]} == 0)); then
+      echo "--model-type requires a non-empty value" >&2
+      exit 2
+    fi
+    local -A _seen=()
+    local -a _out=()
+    local _x
+    for _x in "${MODEL_TYPE_CLI_PARTS[@]}"; do
+      if [[ "${_x}" == all ]]; then
+        printf '%s' "all"
+        return 0
+      fi
+      [[ ${_seen["${_x}"]+isset} ]] && continue
+      _seen["${_x}"]=1
+      _out+=("${_x}")
+    done
+    (IFS=','; printf '%s' "${_out[*]}")
+  else
+    printf '%s' "${MODEL_TYPE_ENV}"
+  fi
+}
+
+_run_jobs_find_repo_containing_yml() {
+  local dir="${1:-}"
+  [[ -n "$dir" ]] || return 1
+  dir="$(cd "$dir" && pwd)" || return 1
+  while true; do
+    if [[ -f "${dir}/${BUILDKITE_REL}" ]]; then
+      printf '%s\n' "$dir"
+      return 0
+    fi
+    [[ "${dir}" == "/" ]] && return 1
+    dir="$(dirname "${dir}")"
+  done
+}
+
+_run_jobs_derive_repo_root_from_yml() {
+  local yml="$1"
+  local d
+  d="$(cd "$(dirname "${yml}")" && pwd)"
+  [[ "$(basename "${d}")" == ".buildkite" ]] || return 1
+  printf '%s\n' "$(dirname "${d}")"
+}
+
+run_yaml_ci_jobs_main() {
+  if [[ -z "${BUILDKITE_REL:-}" || -z "${DEFAULT_LOG_SUBDIR:-}" ]]; then
+    echo "run_yaml_ci_jobs_main: set BUILDKITE_REL and DEFAULT_LOG_SUBDIR before calling." >&2
+    exit 1
+  fi
+
+  LABEL_SUBSTR="${LABEL_SUBSTR:-}"
+  MODEL_TYPE_ENV="${MODEL_TYPE:-all}"
+  MODEL_TYPE_CLI_PARTS=()
+  MODEL_TYPE_FROM_CLI=0
+  DRY_RUN="${DRY_RUN:-0}"
+  SKIP_SIMPLE=0
+
+  while [[ $# -gt 0 ]]; do
+    case "$1" in
+      -h | --help)
+        _run_yaml_ci_jobs_usage
+        exit 0
+        ;;
+      --dry-run)
+        DRY_RUN=1
+        shift
+        ;;
+      --log-dir)
+        LOG_DIR="$2"
+        shift 2
+        ;;
+      --yaml)
+        YML="$2"
+        shift 2
+        ;;
+      --repo-root)
+        REPO_ROOT="$2"
+        shift 2
+        ;;
+      --label-substr)
+        LABEL_SUBSTR="$2"
+        shift 2
+        ;;
+      --model-type)
+        MODEL_TYPE_FROM_CLI=1
+        _run_jobs_split_append_csv_array MODEL_TYPE_CLI_PARTS "$2"
+        shift 2
+        ;;
+      --skip-simple)
+        SKIP_SIMPLE=1
+        shift
+        ;;
+      *)
+        echo "Unknown option: $1" >&2
+        _run_yaml_ci_jobs_usage >&2
+        exit 2
+        ;;
+    esac
+  done
+
+  MODEL_TYPE="$(_run_jobs_finalize_model_type_csv)"
+
+  if [[ -n "${YML:-}" && -n "${REPO_ROOT:-}" ]]; then
+    REPO_ROOT="$(cd "${REPO_ROOT}" && pwd)"
+    YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")"
+  elif [[ -n "${YML:-}" ]]; then
+    YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")"
+    if ! REPO_ROOT="$(_run_jobs_derive_repo_root_from_yml "${YML}")"; then
+      echo "Could not derive REPO_ROOT from YML=${YML} (expected file at <repo>/${BUILDKITE_REL})." >&2
+      echo "Set REPO_ROOT explicitly (or pass --repo-root) for pytest working directory." >&2
+      exit 2
+    fi
+  elif [[ -n "${REPO_ROOT:-}" ]]; then
+    REPO_ROOT="$(cd "${REPO_ROOT}" && pwd)"
+    YML="${REPO_ROOT}/${BUILDKITE_REL}"
+  else
+    REPO_ROOT=""
+    if command -v git >/dev/null 2>&1; then
+      REPO_ROOT="$(git -C "${PWD}" rev-parse --show-toplevel 2>/dev/null || true)"
+    fi
+    if [[ -z "${REPO_ROOT}" ]]; then
+      REPO_ROOT="$(_run_jobs_find_repo_containing_yml "${PWD}" || true)"
+    fi
+    if [[ -z "${REPO_ROOT}" ]]; then
+      REPO_ROOT="$(_run_jobs_find_repo_containing_yml "${SCRIPT_DIR}" || true)"
+    fi
+    if [[ -z "${REPO_ROOT}" ]]; then
+      echo "Could not locate ${BUILDKITE_REL}. Set REPO_ROOT or YML, run from inside the vllm-omni clone," >&2
+      echo "or place this script (or run from a cwd) under the repository tree." >&2
+      exit 2
+    fi
+    YML="${REPO_ROOT}/${BUILDKITE_REL}"
+  fi
+
+  LOG_DIR="${LOG_DIR:-${REPO_ROOT}/logs/${DEFAULT_LOG_SUBDIR}}"
+
+  if [[ ! -f "${YML}" ]]; then
+    echo "YAML not found: ${YML}" >&2
+    exit 2
+  fi
+
+  if ! command -v python3 >/dev/null 2>&1; then
+    echo "python3 is required." >&2
+    exit 1
+  fi
+
+  mkdir -p "${LOG_DIR}/jobs"
+  LOG_DIR="$(cd "${LOG_DIR}" && pwd)"
+  export REPO_ROOT LOG_DIR YML LABEL_SUBSTR MODEL_TYPE DRY_RUN SKIP_SIMPLE
+
+  if [[ "${DRY_RUN}" != "1" ]]; then
+    shopt -s nullglob
+    _stale=( "${LOG_DIR}/jobs"/*.sh )
+    _stale_logs=( "${LOG_DIR}"/*.log )
+    shopt -u nullglob
+    if ((${#_stale[@]})); then
+      rm -f "${_stale[@]}"
+    fi
+    if ((${#_stale_logs[@]})); then
+      rm -f "${_stale_logs[@]}"
+    fi
+    rm -f "${LOG_DIR}/jobs/.job_timeouts"
+  fi
+
+  # shellcheck disable=SC2016,SC1078,SC1079
+  python3 - <<'PY'
+from __future__ import annotations
+
+import os
+import re
+import stat
+import sys
+from pathlib import Path
+
+try:
+    import yaml
+except ImportError:
+    print("Missing PyYAML. Install with: pip install pyyaml", file=sys.stderr)
+    sys.exit(1)
+
+REPO_ROOT = Path(os.environ["REPO_ROOT"]).resolve()
+LOG_DIR = Path(os.environ["LOG_DIR"]).resolve()
+YML = Path(os.environ["YML"]).resolve()
+LABEL_SUBSTR = (os.environ.get("LABEL_SUBSTR") or "").strip()
+DRY_RUN = os.environ.get("DRY_RUN", "0") == "1"
+SKIP_SIMPLE = os.environ.get("SKIP_SIMPLE", "0") == "1"
+
+ALLOWED_MODEL_TYPES = frozenset({"all", "omni", "tts", "diffusion"})
+
+
+def parse_model_types(raw: str) -> list[str]:
+    parts = [p.strip().lower() for p in (raw or "").split(",") if p.strip()]
+    if not parts:
+        parts = ["all"]
+    bad = [p for p in parts if p not in ALLOWED_MODEL_TYPES]
+    if bad:
+        print(
+            f"Invalid MODEL_TYPE / --model-type value(s): {bad!r} "
+            f"(allowed: {', '.join(sorted(ALLOWED_MODEL_TYPES))})",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+    if "all" in parts:
+        return ["all"]
+    out: list[str] = []
+    seen: set[str] = set()
+    for p in parts:
+        if p not in seen:
+            seen.add(p)
+            out.append(p)
+    return out
+
+
+PYTEST_CMD_RE = re.compile(
+    r"(?:timeout\s+\S+\s+)?(?:python3? -m\s+)?pytest\s+[^\n&|;]*"
+)
+
+
+def iter_leaf_steps(steps, group: str | None = None):
+    for raw in steps or []:
+        if not isinstance(raw, dict):
+            continue
+        nested = raw.get("steps")
+        if isinstance(nested, list) and nested:
+            g = raw.get("group")
+            next_group: str | None
+            if isinstance(g, str):
+                next_group = g
+            elif g is not None:
+                next_group = str(g)
+            else:
+                next_group = group
+            yield from iter_leaf_steps(nested, next_group)
+            continue
+        if raw.get("commands"):
+            yield raw, group
+
+
+def raw_command_text(step: dict) -> str:
+    raw = step.get("commands") or []
+    if isinstance(raw, str):
+        raw = [raw]
+    text = "\n".join((c.strip() if isinstance(c, str) else "") for c in raw if c)
+    return text.replace("$$", "$")
+
+
+def export_lines(text: str) -> list[str]:
+    out: list[str] = []
+    for line in text.splitlines():
+        s = line.strip()
+        if not s or s.startswith("#"):
+            continue
+        if s.startswith("export "):
+            out.append(s)
+    return out
+
+
+def pytest_lines(text: str) -> list[str]:
+    out: list[str] = []
+    for m in PYTEST_CMD_RE.finditer(text):
+        line = m.group(0).strip()
+        line_start = text.rfind("\n", 0, m.start()) + 1
+        before = text[line_start : m.start()]
+        if before.lstrip().startswith("#"):
+            continue
+        if line:
+            out.append(line)
+    return out
+
+
+def is_simple_test_step(label: str, group: str | None) -> bool:
+    if group and "Simple Test" in group:
+        return True
+    return label.startswith("Simple ·")
+
+
+def label_matches_model_type(label: str, model_types: list[str]) -> bool:
+    if "all" in model_types:
+        return True
+    if "omni" in model_types and "Omni ·" in label:
+        return True
+    if "tts" in model_types and "TTS ·" in label:
+        return True
+    if "diffusion" in model_types and "Diffusion" in label:
+        return True
+    return False
+
+
+def job_key_from_step(step: dict, label: str) -> str:
+    k = step.get("key")
+    if isinstance(k, str) and k.strip():
+        return k.strip()
+    slug = re.sub(r"[^\w\-.]+", "_", label or "job", flags=re.UNICODE)
+    slug = re.sub(r"_+", "_", slug).strip("_") or "job"
+    return slug
+
+
+def step_timeout_minutes(step: dict) -> int | None:
+    raw = step.get("timeout_in_minutes")
+    if raw is None:
+        return None
+    try:
+        minutes = int(raw)
+    except (TypeError, ValueError):
+        return None
+    return minutes if minutes > 0 else None
+
+
+_INLINE_TIMEOUT_PREFIX = re.compile(r"^timeout\s+(?:-\S+\s+)*\S+\s+", re.IGNORECASE)
+
+
+def timeout_kill_after_seconds() -> int:
+    raw = os.environ.get("RUN_JOB_TIMEOUT_KILL_AFTER", "60")
+    try:
+        seconds = int(raw)
+    except (TypeError, ValueError):
+        return 60
+    return seconds if seconds > 0 else 60
+
+
+def prepend_timeout_to_pytest(pytest_line: str, timeout_min: int | None) -> str:
+    """Prepend GNU timeout directly before pytest (Buildkite inline style)."""
+    if timeout_min is None:
+        return pytest_line.strip()
+    line = _INLINE_TIMEOUT_PREFIX.sub("", pytest_line.strip())
+    kill_after = timeout_kill_after_seconds()
+    return (
+        f"timeout --foreground --verbose --kill-after={kill_after} "
+        f"{timeout_min}m {line}"
+    )
+
+
+def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) -> None:
+    manifest_path = jobs_dir / ".job_timeouts"
+    if job_timeouts:
+        manifest_path.write_text(
+            "\n".join(f"{k}={v}" for k, v in sorted(job_timeouts.items())) + "\n",
+            encoding="utf-8",
+        )
+    elif manifest_path.is_file():
+        manifest_path.unlink()
+
+
+def _write_job_script(key: str, script_lines: list[str], jobs_dir: Path) -> None:
+    body = "\n".join(script_lines) + "\n"
+    if DRY_RUN:
+        print(f"=== {key} ===")
+        print(body)
+        return
+    job_path = jobs_dir / f"{key}.sh"
+    job_path.write_text(body, encoding="utf-8")
+    try:
+        mode = job_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
+        job_path.chmod(mode)
+    except OSError:
+        pass
+    print(f"generated {job_path}", file=sys.stderr)
+
+
+def main() -> None:
+    model_types = parse_model_types(os.environ.get("MODEL_TYPE", "all"))
+
+    jobs_dir = LOG_DIR / "jobs"
+    if not DRY_RUN:
+        jobs_dir.mkdir(parents=True, exist_ok=True)
+
+    if not YML.is_file():
+        print(f"YAML not found: {YML}", file=sys.stderr)
+        sys.exit(1)
+
+    data = yaml.safe_load(YML.read_text(encoding="utf-8"))
+    top_steps = (data or {}).get("steps") or []
+
+    matched_yaml = 0
+    job_timeouts: dict[str, int] = {}
+    for step, grp in iter_leaf_steps(top_steps):
+        label = step.get("label") or ""
+        if SKIP_SIMPLE and is_simple_test_step(label, grp):
+            continue
+        if LABEL_SUBSTR and LABEL_SUBSTR not in label:
+            continue
+        if not label_matches_model_type(label, model_types):
+            continue
+        text = raw_command_text(step)
+        exports = export_lines(text)
+        pys = pytest_lines(text)
+        if not pys:
+            print(f"# skip (no pytest line): {label!r}", file=sys.stderr)
+            continue
+
+        key = job_key_from_step(step, label)
+        matched_yaml += 1
+
+        script_lines = [
+            "#!/usr/bin/env bash",
+            f'# From Buildkite label: {label.replace(chr(10), " ")}',
+            "set -euo pipefail",
+            f'cd "{REPO_ROOT}"',
+        ]
+        timeout_min = step_timeout_minutes(step)
+        if timeout_min is not None:
+            script_lines.append(f"# Buildkite timeout_in_minutes: {timeout_min}")
+            job_timeouts[key] = timeout_min
+        script_lines.extend(exports)
+        script_lines.extend(prepend_timeout_to_pytest(p, timeout_min) for p in pys)
+        _write_job_script(key, script_lines, jobs_dir)
+
+    if not DRY_RUN:
+        _write_job_timeouts_manifest(jobs_dir, job_timeouts)
+
+    if matched_yaml == 0:
+        print(
+            f"No YAML steps matched MODEL_TYPE={model_types!r} "
+            f"LABEL_SUBSTR={LABEL_SUBSTR!r} SKIP_SIMPLE={SKIP_SIMPLE!r} in {YML}",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
+
+main()
+PY
+
+  if [[ "${DRY_RUN}" == "1" ]]; then
+    exit 0
+  fi
+
+  set -o pipefail
+  local any_fail=0
+  local _run_start
+  local _job
+  _run_jobs_reset_timing
+  _run_start="$(_run_jobs_epoch_seconds)"
+
+  shopt -s nullglob
+  for _job in "${LOG_DIR}/jobs"/*.sh; do
+    _run_one_job_with_timing "${_job}" || any_fail=1
+  done
+  shopt -u nullglob
+
+  _run_jobs_print_timing_summary "${_run_start}" "${any_fail}"
+  if [[ "${any_fail}" -ne 0 ]]; then
+    exit 1
+  fi
+  exit 0
+}
diff --git a/tools/run_merge_jobs.sh b/tools/run_merge_jobs.sh
new file mode 100644
index 000000000..35b38a8e0
--- /dev/null
+++ b/tools/run_merge_jobs.sh
@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Extract steps from .buildkite/test-merge.yml that contain pytest, synthesize
+# small bash wrappers (exports + pytest), run them, and tee output to logs named
+# after each step's Buildkite "key" when present (otherwise a slug of the label).
+#
+# Model area (--model-type / MODEL_TYPE), multiple allowed (OR semantics):
+#   e.g. --model-type omni,tts
+#   omni     — label contains "Omni ·"
+#   tts      — label contains "TTS ·"
+#   diffusion— label contains "Diffusion"
+#   all      — no model filter (default)
+#
+#   --skip-simple — skip steps in the "Simple Test" Buildkite group and labels
+#                   starting with "Simple ·" (L1-style unit tests in the same YAML)
+#
+# Requirements: bash, python3, PyYAML (pip install pyyaml)
+#
+# Usage:
+#   bash tools/run_merge_jobs.sh
+#   REPO_ROOT=/path/to/vllm-omni bash tools/run_merge_jobs.sh --model-type diffusion --dry-run
+#   YML=/path/to/vllm-omni/.buildkite/test-merge.yml bash tools/run_merge_jobs.sh
+#
+# Repository / YAML (no dependency on where this script lives):
+#   • Set REPO_ROOT (or pass --repo-root) — default YAML is $REPO_ROOT/.buildkite/test-merge.yml
+#   • Or set YML (or --yaml) — repo root is inferred as parent of the .buildkite directory
+#   • Or run from inside the clone: git rev-parse --show-toplevel, else walk up from $PWD,
+#     then from the script's directory, until .buildkite/test-merge.yml exists
+#
+# Optional environment:
+#   REPO_ROOT     - vllm-omni root (working directory for pytest); see above
+#   YML           - path to test-merge.yml (default: $REPO_ROOT/.buildkite/test-merge.yml)
+#   LOG_DIR       - logs + generated job scripts (default: $REPO_ROOT/logs/merge_jobs);
+#                   per-job *.log plus timing_summary.log after the run
+#   MODEL_TYPE    - comma-separated and/or repeated flags (default: all); see above
+#   LABEL_SUBSTR  - substring of Buildkite step label
+#   DRY_RUN=1     - print extracted commands only; do not write scripts or run pytest
+#
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+BUILDKITE_REL=".buildkite/test-merge.yml"
+DEFAULT_LOG_SUBDIR="merge_jobs"
+
+# shellcheck source=tools/run_jobs_common.sh
+source "${SCRIPT_DIR}/run_jobs_common.sh"
+run_yaml_ci_jobs_main "$@"
diff --git a/tools/run_ready_jobs.sh b/tools/run_ready_jobs.sh
new file mode 100644
index 000000000..89f082436
--- /dev/null
+++ b/tools/run_ready_jobs.sh
@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Extract steps from .buildkite/test-ready.yml that contain pytest, synthesize
+# small bash wrappers (exports + pytest), run them, and tee output to logs named
+# after each step's Buildkite "key" when present (otherwise a slug of the label).
+#
+# Model area (--model-type / MODEL_TYPE), multiple allowed (OR semantics):
+#   e.g. --model-type omni,tts
+#   omni     — label contains "Omni ·"
+#   tts      — label contains "TTS ·"
+#   diffusion— label contains "Diffusion"
+#   all      — no model filter (default)
+#
+#   --skip-simple — skip steps in the "Simple Test" Buildkite group and labels
+#                   starting with "Simple ·" (L1-style unit tests in the same YAML)
+#
+# Requirements: bash, python3, PyYAML (pip install pyyaml)
+#
+# Usage:
+#   bash tools/run_ready_jobs.sh
+#   REPO_ROOT=/path/to/vllm-omni bash tools/run_ready_jobs.sh --model-type omni --dry-run
+#   YML=/path/to/vllm-omni/.buildkite/test-ready.yml bash tools/run_ready_jobs.sh
+#
+# Repository / YAML (no dependency on where this script lives):
+#   • Set REPO_ROOT (or pass --repo-root) — default YAML is $REPO_ROOT/.buildkite/test-ready.yml
+#   • Or set YML (or --yaml) — repo root is inferred as parent of the .buildkite directory
+#   • Or run from inside the clone: git rev-parse --show-toplevel, else walk up from $PWD,
+#     then from the script's directory, until .buildkite/test-ready.yml exists
+#
+# Optional environment:
+#   REPO_ROOT     - vllm-omni root (working directory for pytest); see above
+#   YML           - path to test-ready.yml (default: $REPO_ROOT/.buildkite/test-ready.yml)
+#   LOG_DIR       - logs + generated job scripts (default: $REPO_ROOT/logs/ready_jobs);
+#                   per-job *.log plus timing_summary.log after the run
+#   MODEL_TYPE    - comma-separated and/or repeated flags (default: all); see above
+#   LABEL_SUBSTR  - substring of Buildkite step label
+#   DRY_RUN=1     - print extracted commands only; do not write scripts or run pytest
+#
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+BUILDKITE_REL=".buildkite/test-ready.yml"
+DEFAULT_LOG_SUBDIR="ready_jobs"
+
+# shellcheck source=tools/run_jobs_common.sh
+source "${SCRIPT_DIR}/run_jobs_common.sh"
+run_yaml_ci_jobs_main "$@"
diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py
index eb06dde70..13c1f2c21 100644
--- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py
+++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py
@@ -36,7 +36,11 @@ Enable with ``SEED_TTS_WER_EVAL=1`` or ``--seed-tts-wer-eval``. Install optional
     pip install 'vllm-omni[dev]'
 
 Env: ``SEED_TTS_EVAL_DEVICE`` (e.g. ``cuda:0``, ``cpu``); ``SEED_TTS_HF_WHISPER_MODEL``
-defaults to ``openai/whisper-large-v3`` (override for debugging only).
+defaults to ``openai/whisper-large-v3`` (override for debugging only). Set
+``SEED_TTS_WER_SAVE_AUDIO_DIR`` to save the captured 24 kHz mono WAV used by
+WER evaluation for each synthesized utterance.
+Streaming PCM is decoded using ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` /
+``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` (default: 24 kHz mono).
 """
 
 from __future__ import annotations
@@ -50,12 +54,14 @@ import string
 import tempfile
 import threading
 import wave
+from pathlib import Path
 from typing import Any
 
 import numpy as np
 from vllm.benchmarks.datasets import SampleRequest
 
 from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSSampleRequest
+from vllm_omni.metrics.definitions import stream_pcm_format_from_env
 
 logger = logging.getLogger(__name__)
 
@@ -87,6 +93,30 @@ def pcm_s16le_mono_to_wav_bytes(pcm: bytes, *, sample_rate: int = 24000) -> byte
     return buf.getvalue()
 
 
+def _safe_filename_part(value: Any, *, default: str = "item", max_len: int = 96) -> str:
+    text = str(value or "").strip() or default
+    safe = "".join(ch if ch.isalnum() or ch in ("-", "_", ".") else "_" for ch in text)
+    safe = safe.strip("._") or default
+    return safe[:max_len]
+
+
+def _save_seed_tts_eval_audio(
+    pcm: bytes,
+    *,
+    output_dir: Path | None,
+    index: int,
+    utterance_id: Any,
+    locale: str,
+) -> str | None:
+    if output_dir is None:
+        return None
+    output_dir.mkdir(parents=True, exist_ok=True)
+    stem = f"{index:05d}_{_safe_filename_part(utterance_id)}_{_safe_filename_part(locale)}"
+    path = output_dir / f"{stem}.wav"
+    path.write_bytes(pcm_s16le_mono_to_wav_bytes(pcm, sample_rate=24000))
+    return str(path)
+
+
 def _get_eval_device() -> str:
     explicit = os.environ.get("SEED_TTS_EVAL_DEVICE", "").strip()
     if explicit:
@@ -146,12 +176,27 @@ def process_one_official(hypo: str, truth: str, lang: str) -> tuple[float, str,
     return wer, raw_truth, raw_hypo
 
 
-def _pcm_s16le_to_f32_16k(pcm: bytes, pcm_sample_rate: int = 24000) -> np.ndarray:
+def _pcm_s16le_to_f32_16k(
+    pcm: bytes,
+    pcm_sample_rate: int | None = None,
+    channels: int | None = None,
+) -> np.ndarray:
     import scipy.signal
 
     if not pcm:
         return np.zeros(0, dtype=np.float32)
+    if pcm_sample_rate is None or channels is None:
+        env_sample_rate, env_channels = stream_pcm_format_from_env()
+        if pcm_sample_rate is None:
+            pcm_sample_rate = env_sample_rate
+        if channels is None:
+            channels = env_channels
     raw = np.frombuffer(pcm, dtype=np.int16).astype(np.float32) / 32768.0
+    if channels > 1:
+        usable = (len(raw) // channels) * channels
+        if usable <= 0:
+            return np.zeros(0, dtype=np.float32)
+        raw = raw[:usable].reshape(-1, channels).mean(axis=1)
     target_len = int(len(raw) * 16000 / pcm_sample_rate)
     if target_len <= 0:
         return np.zeros(0, dtype=np.float32)
@@ -512,13 +557,18 @@ def compute_seed_tts_wer_metrics(
     sim_skipped_no_ref = 0
     utmos_failed = 0
     utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=False)
+    save_audio_raw = os.environ.get("SEED_TTS_WER_SAVE_AUDIO_DIR", "").strip()
+    save_audio_dir = Path(save_audio_raw).expanduser() if save_audio_raw else None
+    saved_audio = 0
+    save_audio_failed = 0
 
-    for req, out in zip(input_requests, outputs, strict=True):
+    for index, (req, out) in enumerate(zip(input_requests, outputs, strict=True)):
         assert isinstance(req, SeedTTSSampleRequest)
         ref = req.prompt
         locale = req.seed_tts_locale or "en"
         row_lang = "zh" if locale.lower().startswith("zh") else "en"
         utmos_v: float | None = None
+        audio_path: str | None = None
 
         if not out.success:
             request_failed += 1
@@ -545,8 +595,30 @@ def compute_seed_tts_wer_metrics(
                     }
                 )
             continue
+        try:
+            audio_path = _save_seed_tts_eval_audio(
+                pcm,
+                output_dir=save_audio_dir,
+                index=index,
+                utterance_id=req.seed_tts_utterance_id,
+                locale=locale,
+            )
+            if audio_path:
+                saved_audio += 1
+        except OSError as e:
+            save_audio_failed += 1
+            logger.warning(
+                "Seed-TTS WER audio save failed for utterance=%s: %s",
+                req.seed_tts_utterance_id,
+                e,
+            )
 
-        wav_16k = _pcm_s16le_to_f32_16k(pcm)
+        # Request functions normalize ``tts_output_pcm_bytes`` to Seed-TTS WER
+        # format before it reaches this evaluator: 24 kHz mono int16 PCM.
+        # Do not apply the streamed-response env format here again; for MOSS
+        # local that would reinterpret already-normalized 24 kHz mono audio as
+        # 48 kHz stereo and corrupt ASR/WER.
+        wav_16k = _pcm_s16le_to_f32_16k(pcm, pcm_sample_rate=24000, channels=1)
         if len(wav_16k) == 0:
             asr_failed += 1
             if include_per_item:
@@ -555,6 +627,7 @@ def compute_seed_tts_wer_metrics(
                         "utterance_id": req.seed_tts_utterance_id,
                         "locale": locale,
                         "error": "empty_audio",
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -609,6 +682,7 @@ def compute_seed_tts_wer_metrics(
                         "locale": locale,
                         "error": "asr_exception",
                         "detail": str(e)[:500],
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -621,6 +695,7 @@ def compute_seed_tts_wer_metrics(
                         "utterance_id": req.seed_tts_utterance_id,
                         "locale": locale,
                         "error": "empty_asr",
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -637,6 +712,7 @@ def compute_seed_tts_wer_metrics(
                         "locale": locale,
                         "error": "wer_compute_failed",
                         "detail": str(e)[:500],
+                        "audio_path": audio_path,
                     }
                 )
             continue
@@ -673,6 +749,8 @@ def compute_seed_tts_wer_metrics(
                 "reference_raw": raw_truth,
                 "asr_raw": raw_hypo,
             }
+            if audio_path:
+                row["audio_path"] = audio_path
             if sim_v is not None:
                 row["sim"] = sim_v
             if utmos_v is not None:
@@ -697,7 +775,11 @@ def compute_seed_tts_wer_metrics(
         "seed_tts_utmos_mean": statistics.fmean(utmos_values) if utmos_values else None,
         "seed_tts_utmos_median": statistics.median(utmos_values) if utmos_values else None,
         "seed_tts_utmos_failed": utmos_failed,
+        "seed_tts_saved_audio": saved_audio,
+        "seed_tts_save_audio_failed": save_audio_failed,
     }
+    if save_audio_dir is not None:
+        result["seed_tts_save_audio_dir"] = str(save_audio_dir)
     if include_per_item:
         result["seed_tts_wer_eval_items"] = items
     return result
@@ -729,6 +811,13 @@ def print_seed_tts_wer_summary(metrics: dict[str, Any]) -> None:
     print("{:<40} {:<10}".format("Request failed:", metrics.get("seed_tts_request_failed", 0)))
     print("{:<40} {:<10}".format("No PCM captured:", metrics.get("seed_tts_no_pcm", 0)))
     print("{:<40} {:<10}".format("ASR / WER failed:", metrics.get("seed_tts_asr_failed", 0)))
+    save_dir = metrics.get("seed_tts_save_audio_dir")
+    if save_dir:
+        print("{:<40} {:<10}".format("Saved eval WAVs:", metrics.get("seed_tts_saved_audio", 0)))
+        print("{:<40} {}".format("Saved eval WAV dir:", save_dir))
+        failed = int(metrics.get("seed_tts_save_audio_failed", 0) or 0)
+        if failed:
+            print("{:<40} {:<10}".format("Save eval WAV errors:", failed))
     if sim_ev or metrics.get("seed_tts_sim_skipped_no_ref") or metrics.get("seed_tts_sim_failed"):
         print("{:<40} {:<10}".format("SIM evaluated (higher ~ closer):", sim_ev))
         sm = metrics.get("seed_tts_sim_mean")
diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py
index a756c97f0..e4f11b1e7 100644
--- a/vllm_omni/benchmarks/patch/patch.py
+++ b/vllm_omni/benchmarks/patch/patch.py
@@ -35,8 +35,6 @@ from vllm.benchmarks.lib.endpoint_request_func import (
 from vllm.logger import init_logger
 from vllm.tokenizers import TokenizerLike
 
-logger = init_logger(__name__)
-
 from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats
 from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset, DailyOmniSampleRequest
 from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset
@@ -51,9 +49,9 @@ from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDa
 from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
 from vllm_omni.metrics import definitions as defs
 
+logger = init_logger(__name__)
+
 _AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S"
-_AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE"
-_AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS"
 RETURN_STAGE_METRICS_FIELD = "return_stage_metrics"
 _IMAGE_STAGE_METRICS_BACKENDS = frozenset({"openai-image-edits-omni"})
 _PRINT_STAGE = False
@@ -111,30 +109,28 @@ def _audio_continuity_threshold_s() -> float:
     return max(value, 0.0)
 
 
-def _audio_pcm_format() -> tuple[int, int]:
-    """Return ``(sample_rate, channels)`` of the streamed PCM response.
+def _pcm_s16le_to_seed_tts_wer_bytes(
+    pcm_bytes: bytes,
+    *,
+    sample_rate: int,
+    channels: int,
+) -> bytes:
+    """Normalize streamed raw PCM to the 24 kHz mono PCM used by Seed-TTS WER."""
+    if not pcm_bytes:
+        return b""
+    channels = max(1, int(channels))
+    pcm = np.frombuffer(pcm_bytes, dtype=np.int16)
+    if channels > 1:
+        usable = (pcm.size // channels) * channels
+        pcm = pcm[:usable].reshape(-1, channels).mean(axis=1)
+    pcm_f32 = pcm.astype(np.float32) / 32767.0
+    if int(sample_rate) != 24000 and pcm_f32.size:
+        from vllm.multimodal.audio import AudioResampler
 
-    Defaults to 24 kHz mono (Qwen3-TTS / VoxCPM2 / most MOSS-TTS variants).
-    Override for models with a different codec output format (e.g.
-    MOSS-TTS-Local-Transformer-v1.5, which streams 48 kHz stereo PCM) via
-    ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` / ``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` -
-    getting this wrong silently skews audio_duration/audio_rtf/underrun stats.
-    """
-    sample_rate = 24000
-    channels = 1
-    raw_sr = os.environ.get(_AUDIO_SAMPLE_RATE_ENV)
-    if raw_sr:
-        try:
-            sample_rate = int(raw_sr)
-        except ValueError:
-            logger.warning("Invalid %s=%r; using default %d", _AUDIO_SAMPLE_RATE_ENV, raw_sr, sample_rate)
-    raw_ch = os.environ.get(_AUDIO_CHANNELS_ENV)
-    if raw_ch:
-        try:
-            channels = int(raw_ch)
-        except ValueError:
-            logger.warning("Invalid %s=%r; using default %d", _AUDIO_CHANNELS_ENV, raw_ch, channels)
-    return sample_rate, channels
+        resampler = AudioResampler(target_sr=24000)
+        pcm_f32 = resampler.resample(pcm_f32, orig_sr=int(sample_rate))
+    pcm_f32 = np.clip(pcm_f32, -1.0, 1.0)
+    return (pcm_f32 * 32767).astype(np.int16).tobytes()
 
 
 get_samples_old = datasets.get_samples
@@ -1163,9 +1159,8 @@ async def async_request_openai_audio_speech(
     output = MixRequestFuncOutput()
     output.prompt_len = request_func_input.prompt_len
 
-    # PCM format: 16-bit signed; sample_rate/channels are model-dependent
-    # (see _audio_pcm_format docstring).
-    sample_rate, channels = _audio_pcm_format()
+    # PCM format: 16-bit signed; sample_rate/channels are model-dependent.
+    sample_rate, channels = defs.stream_pcm_format_from_env()
     sample_width = 2  # 16-bit = 2 bytes
 
     st = time.perf_counter()
@@ -1216,7 +1211,15 @@ async def async_request_openai_audio_speech(
                 output.audio_continuity_ok = continuity.is_continuous
                 output.audio_underrun_event_count = continuity.underrun_event_count
                 if pcm_capture is not None and pcm_capture:
-                    output.tts_output_pcm_bytes = bytes(pcm_capture)
+                    try:
+                        output.tts_output_pcm_bytes = _pcm_s16le_to_seed_tts_wer_bytes(
+                            bytes(pcm_capture),
+                            sample_rate=sample_rate,
+                            channels=channels,
+                        )
+                    except Exception as ex:
+                        logger.warning("Seed-TTS WER PCM normalization failed: %s", ex)
+                        output.tts_output_pcm_bytes = bytes(pcm_capture)
                 elif capture_wer_pcm:
                     ct = response.headers.get("Content-Type", "")
                     logger.warning(
@@ -1261,7 +1264,7 @@ if "daily-omni" not in OPENAI_COMPATIBLE_BACKENDS:
 # Prevent import order from causing patch failures
 from vllm.benchmarks import serve
 from vllm.benchmarks.lib.ready_checker import wait_for_endpoint
-from vllm.benchmarks.serve import TaskType, _merge_overrides, calculate_metrics_for_embeddings, get_request
+from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, get_request
 
 from vllm_omni.benchmarks.metrics.metrics import (
     MultiModalsBenchmarkMetrics,
@@ -1273,6 +1276,19 @@ from vllm_omni.benchmarks.metrics.metrics import (
 benchmark_old = serve.benchmark
 
 
+def _merge_overrides(base: dict | None, overrides: dict | None) -> dict | None:
+    """Merge benchmark extra_body with per-request overrides.
+
+    vLLM 0.24 removed the private helper from ``vllm.benchmarks.serve``.
+    Keep the same shallow-merge behavior here, with request overrides winning.
+    """
+    if not base and not overrides:
+        return None
+    merged = dict(base or {})
+    merged.update(overrides or {})
+    return merged
+
+
 async def benchmark(
     task_type: TaskType,
     endpoint_type: str,
diff --git a/vllm_omni/config/composable_parallel/__init__.py b/vllm_omni/config/composable_parallel/__init__.py
new file mode 100644
index 000000000..499c14548
--- /dev/null
+++ b/vllm_omni/config/composable_parallel/__init__.py
@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Composable parallel strategies for vLLM-Omni.
+
+A small, runtime-agnostic vocabulary for declaring a stage's parallel layout
+(:class:`StrategySpec` + mesh/routing/aggregation patterns), plus the logic to
+translate a per-role stack of specs into concrete engine sizing and overlay it
+onto merged stage configs.
+
+The spec/pattern *types* are deliberately pure data with no engine imports; the
+translation and apply logic lives alongside them and consumes vLLM-Omni's own
+config objects.
+"""
+
+from vllm_omni.config.composable_parallel.aggregation import (
+    AGGREGATION_PATTERN_VARIANTS,
+    AggregationConflictError,
+    AggregationPattern,
+    AllGather,
+    Combine,
+    FanInByStage,
+    GatherDim,
+    StitchPipeline,
+    StitchSpatial,
+    TakeRank,
+    Union,
+)
+from vllm_omni.config.composable_parallel.apply import (
+    StrategyApplyError,
+    StrategyApplyResult,
+    apply_strategy_specs,
+    check_device_layout,
+)
+from vllm_omni.config.composable_parallel.routing import (
+    ROUTING_PATTERN_VARIANTS,
+    Broadcast,
+    DuplicateWithCondTag,
+    PartitionByHash,
+    PipelineMicrobatch,
+    RouteByExpert,
+    RouteByStage,
+    RoutingKeyError,
+    RoutingPattern,
+    ShardSequence,
+    ShardSpatial,
+)
+from vllm_omni.config.composable_parallel.spec import (
+    MESH_AXIS_KINDS,
+    KernelSpec,
+    LayerHookSpec,
+    MeshAxisKind,
+    MeshAxisSpec,
+    SpecMergeConflictError,
+    StrategySpec,
+)
+from vllm_omni.config.composable_parallel.translator import (
+    AxisTranslationError,
+    OmniParallelConfig,
+    translate_strategy_stack,
+)
+
+__all__ = [
+    # spec types
+    "MeshAxisKind",
+    "MESH_AXIS_KINDS",
+    "MeshAxisSpec",
+    "LayerHookSpec",
+    "KernelSpec",
+    "StrategySpec",
+    "SpecMergeConflictError",
+    # routing patterns
+    "RoutingPattern",
+    "RoutingKeyError",
+    "Broadcast",
+    "PartitionByHash",
+    "ShardSequence",
+    "ShardSpatial",
+    "RouteByExpert",
+    "PipelineMicrobatch",
+    "DuplicateWithCondTag",
+    "RouteByStage",
+    "ROUTING_PATTERN_VARIANTS",
+    # aggregation patterns
+    "AggregationPattern",
+    "AggregationConflictError",
+    "TakeRank",
+    "Union",
+    "GatherDim",
+    "AllGather",
+    "StitchSpatial",
+    "StitchPipeline",
+    "Combine",
+    "FanInByStage",
+    "AGGREGATION_PATTERN_VARIANTS",
+    # translation
+    "OmniParallelConfig",
+    "translate_strategy_stack",
+    "AxisTranslationError",
+    # apply
+    "apply_strategy_specs",
+    "StrategyApplyResult",
+    "check_device_layout",
+    "StrategyApplyError",
+]
diff --git a/vllm_omni/config/composable_parallel/aggregation.py b/vllm_omni/config/composable_parallel/aggregation.py
new file mode 100644
index 000000000..75ed8fadd
--- /dev/null
+++ b/vllm_omni/config/composable_parallel/aggregation.py
@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Closed aggregation-pattern hierarchy for declarative parallel strategies.
+
+An aggregation pattern describes how the per-worker results of one mesh axis
+are recombined into a single logical result (e.g. take one tensor-parallel
+rank's output, union disjoint data-parallel shards). The hierarchy is *closed*:
+subclasses may only be defined in this module.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class AggregationPattern:
+    """Base class for closed aggregation patterns. Subclass only within this module."""
+
+    def __init_subclass__(cls, **kwargs: object) -> None:
+        super().__init_subclass__(**kwargs)
+        if cls.__module__ != __name__:
+            raise TypeError("AggregationPattern is closed to external subclassing")
+
+    def pattern_kind(self) -> str:
+        raise NotImplementedError
+
+
+class AggregationConflictError(ValueError):
+    """Raised when two aggregation inputs conflict for the same request id."""
+
+
+@dataclass(frozen=True)
+class TakeRank(AggregationPattern):
+    """Select one rank's result (TP — all ranks agree)."""
+
+    rank: int = 0
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.rank, int) or isinstance(self.rank, bool):
+            raise TypeError(f"TakeRank.rank must be int, got {type(self.rank).__name__}")
+
+    def pattern_kind(self) -> str:
+        return "take_rank"
+
+
+@dataclass(frozen=True)
+class Union(AggregationPattern):
+    """Union of disjoint replica results (DP, HSDP)."""
+
+    def pattern_kind(self) -> str:
+        return "union"
+
+
+@dataclass(frozen=True)
+class GatherDim(AggregationPattern):
+    """Gather along a dimension and materialize on one designated rank.
+
+    Default designated rank is axis-rank 0 unless adapter overrides.
+    """
+
+    dim: int = 1
+
+    def pattern_kind(self) -> str:
+        return "gather_dim"
+
+
+@dataclass(frozen=True)
+class AllGather(AggregationPattern):
+    """All ranks materialize the same aggregate result across the axis."""
+
+    dim: int | None = None
+
+    def pattern_kind(self) -> str:
+        return "all_gather"
+
+
+@dataclass(frozen=True)
+class StitchSpatial(AggregationPattern):
+    """Stitch spatial patches (VAE-PP)."""
+
+    grid: tuple[int, int] = (1, 1)
+
+    def pattern_kind(self) -> str:
+        return "stitch_spatial"
+
+
+@dataclass(frozen=True)
+class StitchPipeline(AggregationPattern):
+    """Stitch pipeline stage outputs (PP)."""
+
+    def pattern_kind(self) -> str:
+        return "stitch_pipeline"
+
+
+@dataclass(frozen=True)
+class Combine(AggregationPattern):
+    """Custom reducer combination (CFG). Adapter resolves reducer_id."""
+
+    reducer_id: str = "default"
+
+    def pattern_kind(self) -> str:
+        return "combine"
+
+
+@dataclass(frozen=True)
+class FanInByStage(AggregationPattern):
+    """Fan-in results routed by stage (stage replicas)."""
+
+    def pattern_kind(self) -> str:
+        return "fan_in_by_stage"
+
+
+AGGREGATION_PATTERN_VARIANTS: tuple[type[AggregationPattern], ...] = (
+    TakeRank,
+    Union,
+    GatherDim,
+    AllGather,
+    StitchSpatial,
+    StitchPipeline,
+    Combine,
+    FanInByStage,
+)
+
+ReducerFn = Callable[..., object]
diff --git a/vllm_omni/config/composable_parallel/apply.py b/vllm_omni/config/composable_parallel/apply.py
new file mode 100644
index 000000000..879c7aec3
--- /dev/null
+++ b/vllm_omni/config/composable_parallel/apply.py
@@ -0,0 +1,264 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Apply declarative parallel strategies onto merged stage configs.
+
+This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses
+the pipeline topology + deploy YAML into a list of ``StageConfig`` objects; this
+module then overlays the engine sizing derived from a per-role
+``StrategySpec`` stack onto those stages.
+
+Design rules (one writer per axis):
+
+* A strategy only writes the axes it actually declares. If a stack has no ``tp``
+  axis, this module never touches ``tensor_parallel_size`` — it does not force a
+  default of 1.
+* **Conflict-on-explicit.** If a knob was already set explicitly by the deploy
+  YAML and the strategy derives a different value, we raise rather than silently
+  override. Equal values are a no-op; unset values are filled.
+* **Pre-spawn device check.** When a stage declares ``devices``, the device
+  count must be consistent with the engine world size (``tp * dp * pp``) and the
+  replica count, so misconfigurations fail at config time instead of at spawn.
+
+``omni_lb_policy`` is *not* a per-stage config knob: omni reads it once at
+orchestrator construction (``AsyncOmniEngine``), not from stage configs. So a
+stage_replica axis's derived policy is surfaced on :class:`StrategyApplyResult`
+for the caller to pass to the orchestrator; it is never silently written into a
+stage where it would be ignored.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any, NoReturn
+
+from vllm.logger import init_logger
+
+from vllm_omni.config.composable_parallel.spec import StrategySpec
+from vllm_omni.config.composable_parallel.translator import (
+    OmniParallelConfig,
+    translate_strategy_stack,
+)
+
+logger = init_logger(__name__)
+
+# Role key -> the StageConfig fields each declared axis writes.
+_ENGINE_FIELD_BY_KIND: dict[str, str] = {
+    "tp": "tensor_parallel_size",
+    "dp": "data_parallel_size",
+    "pp": "pipeline_parallel_size",
+}
+
+RoleKey = str
+
+
+class StrategyApplyError(ValueError):
+    """Error for applying a strategy onto stage configs.
+
+    A single error type (rather than per-cause subclasses) keeps the public
+    surface small and consistent with the rest of the codebase; the specific
+    cause — an unmatched role, a conflict with an explicit deploy value, or a
+    device-count mismatch — is in the message and is logged before the raise so
+    it stays visible even when the type is unavailable to a caller (e.g. across a
+    server boundary).
+    """
+
+
+def _fail(msg: str) -> NoReturn:
+    """Log and raise :class:`StrategyApplyError`."""
+    logger.error("[composable_parallel] %s", msg)
+    raise StrategyApplyError(msg)
+
+
+@dataclass
+class StrategyApplyResult:
+    """Outcome of applying strategies to a stage list.
+
+    ``stages`` is the same list passed in (mutated in place and returned for
+    convenience). ``omni_lb_policy`` is the pipeline-wide load-balancer policy
+    derived from any stage_replica axes (``None`` if none declared one); the
+    caller wires it into the orchestrator, since omni does not read it from
+    stage configs.
+
+    ``per_role_config`` is keyed by the role key the caller supplied — a
+    ``model_stage`` name (str). ``per_stage_config`` is the same information
+    keyed by the resolved integer ``stage_id`` — handy for re-validating a
+    stage after later layers (e.g. CLI overrides) have run.
+    """
+
+    stages: list[Any]
+    omni_lb_policy: str | None = None
+    per_role_config: dict[RoleKey, OmniParallelConfig] = field(default_factory=dict)
+    per_stage_config: dict[int, OmniParallelConfig] = field(default_factory=dict)
+
+
+def _resolve_stage(stages: Sequence[Any], key: RoleKey) -> Any:
+    """Find the single stage matching ``key`` by ``model_stage`` name.
+
+    Role keys are ``model_stage`` names (str) — the same human-readable names
+    that strategy.yaml authors use to refer to a stage (e.g. ``"thinker"``).
+    """
+    matches = [s for s in stages if getattr(s, "model_stage", None) == key]
+    descriptor = f"model_stage={key!r}"
+
+    if not matches:
+        available = ", ".join(f"{getattr(s, 'model_stage', '?')!r}(id={getattr(s, 'stage_id', '?')})" for s in stages)
+        _fail(f"strategy role {descriptor} did not match any stage; available stages: {available}")
+    if len(matches) > 1:
+        _fail(f"strategy role {descriptor} is ambiguous; it matched {len(matches)} stages")
+    return matches[0]
+
+
+def _set_explicit(engine_args: dict[str, Any], field_name: str, value: Any, *, role: RoleKey) -> None:
+    """Write ``field_name`` with conflict-on-explicit semantics.
+
+    ``engine_args`` is the raw mapping loaded straight from the deploy YAML
+    (before the engine-args dataclass is constructed), so a key is *present* only
+    when the YAML set it explicitly. We therefore key off membership rather than
+    ``.get``: an explicit ``None`` in the YAML is a real, present value and must
+    not be conflated with a key miss (which ``.get`` would report as ``None``).
+    """
+    if field_name in engine_args and engine_args[field_name] != value:
+        existing = engine_args[field_name]
+        _fail(
+            f"role {role!r}: strategy derives {field_name}={value} but the deploy config already "
+            f"set {field_name}={existing}. Remove one of them (the strategy is the single writer "
+            "for the axes it declares)."
+        )
+    engine_args[field_name] = value
+
+
+def _set_num_replicas(runtime: dict[str, Any], value: int, *, role: RoleKey) -> None:
+    """Write ``num_replicas`` with conflict-on-explicit (treating 1 as unset)."""
+    existing = runtime.get("num_replicas", 1)
+    if existing not in (1, value):
+        _fail(
+            f"role {role!r}: strategy derives num_replicas={value} but the deploy config already "
+            f"set num_replicas={existing}. Remove one of them."
+        )
+    runtime["num_replicas"] = value
+
+
+def _parse_device_count(devices: Any) -> int | None:
+    """Return the number of device ids in a ``devices`` value, or None."""
+    if devices is None:
+        return None
+    if isinstance(devices, (list, tuple)):
+        return len([d for d in devices])
+    text = str(devices).strip()
+    if not text:
+        return None
+    return len([d for d in text.split(",") if d.strip()])
+
+
+def check_device_layout(
+    devices: Any,
+    *,
+    tensor_parallel_size: int,
+    data_parallel_size: int,
+    pipeline_parallel_size: int,
+    num_replicas: int,
+    role: RoleKey,
+) -> None:
+    """Validate a stage's device count against its world size (× replicas in pool mode).
+
+    Raises :class:`StrategyApplyError` when the declared ``devices`` count is
+    neither the per-replica world size nor the full ``num_replicas`` pool.
+    Exposed publicly so the resolved (post-CLI) layout can be re-checked after
+    later override layers, not just the strategy-derived snapshot.
+    """
+    count = _parse_device_count(devices)
+    if count is None:
+        return
+    world = int(tensor_parallel_size) * int(data_parallel_size) * int(pipeline_parallel_size)
+    replicas = int(num_replicas or 1)
+    # Two valid shapes: a single per-replica template (== world) or the full
+    # pool across replicas (== replicas * world). This mirrors omni's own
+    # per-replica device splitter.
+    valid = {world, replicas * world}
+    if count not in valid:
+        _fail(
+            f"role {role!r}: declared {count} device id(s) but the strategy world size is {world} "
+            f"(tp={tensor_parallel_size} * dp={data_parallel_size} * pp={pipeline_parallel_size}). "
+            f"Provide either {world} (per-replica template) or {replicas * world} "
+            f"(num_replicas={replicas} pool)."
+        )
+
+
+def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: RoleKey) -> None:
+    """Pre-spawn check on the strategy-derived snapshot (see :func:`check_device_layout`)."""
+    check_device_layout(
+        runtime.get("devices"),
+        tensor_parallel_size=cfg.tensor_parallel_size,
+        data_parallel_size=cfg.data_parallel_size,
+        pipeline_parallel_size=cfg.pipeline_parallel_size,
+        num_replicas=int(runtime.get("num_replicas", 1) or 1),
+        role=role,
+    )
+
+
+def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None:
+    engine_args = stage.yaml_engine_args
+    runtime = stage.yaml_runtime
+
+    declared = set(cfg.l1_owners.keys())
+    for kind, field_name in _ENGINE_FIELD_BY_KIND.items():
+        if kind in declared:
+            _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role)
+    if "ep" in declared and cfg.enable_expert_parallel:
+        _set_explicit(engine_args, "enable_expert_parallel", True, role=role)
+    if "stage_replica" in declared:
+        _set_num_replicas(runtime, cfg.stage_replica_size, role=role)
+
+    _check_devices(runtime, cfg, role=role)
+
+
+def apply_strategy_specs(
+    stages: list[Any],
+    strategy_specs: Mapping[RoleKey, Sequence[StrategySpec]],
+) -> StrategyApplyResult:
+    """Overlay per-role strategy specs onto a merged stage list.
+
+    Args:
+        stages: the ``list[StageConfig]`` returned by ``merge_pipeline_deploy``.
+        strategy_specs: maps a role (a ``model_stage`` name, e.g. ``"thinker"``)
+            to that stage's stack of ``StrategySpec`` (one per declared mesh
+            axis). Role keys are ``model_stage`` names (str).
+
+    Returns:
+        A :class:`StrategyApplyResult` with the mutated stages and the derived
+        pipeline-wide ``omni_lb_policy`` (if any stage_replica axis set one).
+
+    Raises:
+        StrategyApplyError: a role matched zero or multiple stages, a derived
+            value conflicts with an explicit deploy value (including two roles
+            deriving different ``omni_lb_policy`` values), or a stage's device
+            count is inconsistent with its world size.
+        AxisTranslationError: the spec stack is invalid/unsupported.
+        NotImplementedError: the spec stack requests routing not built yet.
+    """
+    result = StrategyApplyResult(stages=stages)
+    lb_policy: str | None = None
+    lb_owner: RoleKey | None = None
+
+    for key, specs in strategy_specs.items():
+        stage = _resolve_stage(stages, key)
+        cfg = translate_strategy_stack(specs)
+        _apply_to_stage(stage, cfg, role=key)
+        result.per_role_config[key] = cfg
+        stage_id = getattr(stage, "stage_id", None)
+        if stage_id is not None:
+            result.per_stage_config[int(stage_id)] = cfg
+
+        if cfg.stage_replica_size > 1 and cfg.omni_lb_policy is not None:
+            if lb_policy is not None and lb_policy != cfg.omni_lb_policy:
+                _fail(
+                    f"role {key!r} derives omni_lb_policy={cfg.omni_lb_policy!r} but role "
+                    f"{lb_owner!r} already derived {lb_policy!r}. The load-balancer policy is "
+                    "pipeline-wide; only one value is allowed."
+                )
+            lb_policy = cfg.omni_lb_policy
+            lb_owner = key
+
+    result.omni_lb_policy = lb_policy
+    return result
diff --git a/vllm_omni/config/composable_parallel/routing.py b/vllm_omni/config/composable_parallel/routing.py
new file mode 100644
index 000000000..38e7d4a49
--- /dev/null
+++ b/vllm_omni/config/composable_parallel/routing.py
@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Closed routing-pattern hierarchy for declarative parallel strategies.
+
+A routing pattern describes how a request batch is distributed across the
+workers of one mesh axis (e.g. broadcast to every tensor-parallel rank, hash
+partition across data-parallel ranks). The hierarchy is *closed*: subclasses
+may only be defined in this module, so the set of patterns is a fixed,
+exhaustively matchable contract.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass(frozen=True)
+class RoutingPattern:
+    """Base class for closed routing patterns. Subclass only within this module."""
+
+    def __init_subclass__(cls, **kwargs: object) -> None:
+        super().__init_subclass__(**kwargs)
+        if cls.__module__ != __name__:
+            raise TypeError("RoutingPattern is closed to external subclassing")
+
+    def pattern_kind(self) -> str:
+        raise NotImplementedError
+
+
+class RoutingKeyError(KeyError):
+    """Raised when a routing key is missing from a request."""
+
+    def __init__(self, request_id: object, missing_key: str) -> None:
+        self.request_id = request_id
+        self.missing_key = missing_key
+        super().__init__(f"Missing routing key {missing_key!r} for request_id={request_id!r}")
+
+
+@dataclass(frozen=True)
+class Broadcast(RoutingPattern):
+    """Every worker in the group sees the whole batch (TP, dense EP)."""
+
+    def pattern_kind(self) -> str:
+        return "broadcast"
+
+
+@dataclass(frozen=True)
+class PartitionByHash(RoutingPattern):
+    """Stable hash partition by request key (DP, HSDP)."""
+
+    key: str = "request_id"
+
+    def pattern_kind(self) -> str:
+        return "partition_by_hash"
+
+
+@dataclass(frozen=True)
+class ShardSequence(RoutingPattern):
+    """Shard along a sequence dimension (SP-Ulysses, SP-Ring, CP)."""
+
+    dim: int = 1
+
+    def pattern_kind(self) -> str:
+        return "shard_sequence"
+
+
+@dataclass(frozen=True)
+class ShardSpatial(RoutingPattern):
+    """Shard along a spatial grid (VAE-PP)."""
+
+    grid: tuple[int, int] = (1, 1)
+
+    def pattern_kind(self) -> str:
+        return "shard_spatial"
+
+
+@dataclass(frozen=True)
+class RouteByExpert(RoutingPattern):
+    """Sparse MoE expert routing."""
+
+    router_id: str = "default"
+
+    def pattern_kind(self) -> str:
+        return "route_by_expert"
+
+
+@dataclass(frozen=True)
+class PipelineMicrobatch(RoutingPattern):
+    """Pipeline parallelism microbatch routing (PP)."""
+
+    microbatch_size: int = 1
+
+    def pattern_kind(self) -> str:
+        return "pipeline_microbatch"
+
+
+@dataclass(frozen=True)
+class DuplicateWithCondTag(RoutingPattern):
+    """Duplicate batch with a conditional tag (CFG)."""
+
+    flag_name: str = "guidance_branch"
+
+    def pattern_kind(self) -> str:
+        return "duplicate_with_cond_tag"
+
+
+@dataclass(frozen=True)
+class RouteByStage(RoutingPattern):
+    """Route requests by pipeline stage policy (stage replicas).
+
+    The ``random``/``round_robin``/``least_queue`` policies describe stateless
+    load balancing that a runtime (e.g. an engine's built-in DP load balancer,
+    or the omni StagePool balancer) can realize. ``hash`` describes
+    deterministic, key-stable routing that a load balancer cannot honor on its
+    own (see ``PartitionByHash``).
+    """
+
+    routing_policy: Literal["random", "round_robin", "least_queue", "hash"] = "round_robin"
+
+    def pattern_kind(self) -> str:
+        return "route_by_stage"
+
+
+ROUTING_PATTERN_VARIANTS: tuple[type[RoutingPattern], ...] = (
+    Broadcast,
+    PartitionByHash,
+    ShardSequence,
+    ShardSpatial,
+    RouteByExpert,
+    PipelineMicrobatch,
+    DuplicateWithCondTag,
+    RouteByStage,
+)
+
+# Optional callable slot for Combine-style custom routing reducers in adapters.
+RouterFn = Callable[..., object]
diff --git a/vllm_omni/config/composable_parallel/spec.py b/vllm_omni/config/composable_parallel/spec.py
new file mode 100644
index 000000000..299fc3ad3
--- /dev/null
+++ b/vllm_omni/config/composable_parallel/spec.py
@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Core declarative types for parallel-strategy specification.
+
+A :class:`StrategySpec` declares one parallelism scheme: the mesh axis it sizes
+(kind + degree), how a batch is routed across that axis, and how per-worker
+results are aggregated back. A *stack* of specs (one per axis) fully describes a
+stage's parallel layout in a runtime-agnostic, data-only form that can be
+translated into concrete engine sizing.
+
+The mesh-axis kinds below enumerate every parallelism dimension the contract can
+describe. Only a subset is wired end-to-end today (see ``translator.py``):
+
+* **Wired** (translatable now): ``tp``, ``dp``, ``pp``, ``ep``,
+  ``stage_replica``.
+* **Reserved** (declared in the type system but not yet translatable — the
+  translator raises ``AxisTranslationError`` for them): ``sp_ulysses``,
+  ``sp_ring``, ``cfg``, ``vae_pp``, ``hsdp``, ``stage_pp``, ``cp``.
+
+Reserved kinds fail fast at translation time rather than silently doing nothing,
+so declaring one is an explicit "not yet" rather than a no-op.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Literal
+
+from vllm_omni.config.composable_parallel.aggregation import AggregationPattern
+from vllm_omni.config.composable_parallel.routing import RoutingPattern
+
+
+class MeshAxisKind(str, Enum):
+    """Enumerates every parallelism dimension the strategy contract can describe.
+
+    Subclassing ``str`` keeps existing string comparisons and dict keys working
+    (``MeshAxisKind.TP == "tp"`` and both hash the same), while making the type
+    refactor-safe: a typo'd kind is an ``AttributeError`` on the enum rather than
+    a silently-wrong bare string.
+    """
+
+    TP = "tp"
+    DP = "dp"
+    PP = "pp"
+    EP = "ep"
+    SP_ULYSSES = "sp_ulysses"
+    SP_RING = "sp_ring"
+    CFG = "cfg"
+    VAE_PP = "vae_pp"
+    HSDP = "hsdp"
+    STAGE_PP = "stage_pp"
+    STAGE_REPLICA = "stage_replica"
+    CP = "cp"
+
+
+# Bare string values, kept as a plain tuple so membership checks and error
+# messages read as the kind names (``"tp"``) rather than enum reprs. ``str``-enum
+# equality means both ``"tp"`` and ``MeshAxisKind.TP`` satisfy ``in`` checks.
+MESH_AXIS_KINDS: tuple[str, ...] = tuple(k.value for k in MeshAxisKind)
+
+HookCategory = Literal[
+    "linear",
+    "attention_pre",
+    "attention_post",
+    "ffn_pre",
+    "ffn_post",
+    "moe_dispatch",
+    "other",
+]
+
+HOOK_CATEGORY_ORDER: dict[HookCategory, int] = {
+    "linear": 0,
+    "attention_pre": 1,
+    "attention_post": 2,
+    "ffn_pre": 3,
+    "ffn_post": 4,
+    "moe_dispatch": 5,
+    "other": 6,
+}
+
+
+class SpecMergeConflictError(ValueError):
+    """Raised when merged hook/kernel specs have conflicting declarations."""
+
+
+@dataclass(frozen=True)
+class MeshAxisSpec:
+    """Declares one axis of a process mesh (kind + size)."""
+
+    kind: MeshAxisKind
+    size: int
+
+    def __post_init__(self) -> None:
+        if self.kind not in MESH_AXIS_KINDS:
+            raise ValueError(f"MeshAxisSpec.kind must be one of {MESH_AXIS_KINDS}, got {self.kind!r}")
+        if self.size <= 0:
+            raise ValueError(f"MeshAxisSpec.size must be > 0, got {self.size}")
+
+
+@dataclass(frozen=True)
+class LayerHookSpec:
+    """L2 hook slot referenced by StrategySpec; resolved by the model walker."""
+
+    hook_id: str
+    target: str | None = None
+    category: HookCategory = "other"
+    priority: int = 0
+    axis_index: int = 0
+
+
+@dataclass(frozen=True)
+class KernelSpec:
+    """L3 kernel slot referenced by StrategySpec; resolved by the kernel registry."""
+
+    kernel_id: str
+    target: str | None = None
+    category: HookCategory = "other"
+    priority: int = 0
+    axis_index: int = 0
+    group_axis_kind: MeshAxisKind | None = None
+    requires_collective: bool = False
+
+
+@dataclass(frozen=True)
+class StrategySpec:
+    """Declarative contract for one parallelism scheme (data + hook/kernel slots)."""
+
+    name: str
+    mesh_axis: MeshAxisSpec
+    routing: RoutingPattern
+    aggregation: AggregationPattern
+    layer_hook_specs: tuple[LayerHookSpec, ...] = ()
+    kernel_specs: tuple[KernelSpec, ...] = ()
+    shard_extension: Mapping[str, Any] = field(default_factory=dict)
diff --git a/vllm_omni/config/composable_parallel/strategy_loader.py b/vllm_omni/config/composable_parallel/strategy_loader.py
new file mode 100644
index 000000000..665fbe9c2
--- /dev/null
+++ b/vllm_omni/config/composable_parallel/strategy_loader.py
@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Load per-role parallel strategies from a strategy file.
+
+The strategy file is a small, optional companion to the deploy config. It maps a
+logical role (``model_stage``, e.g. ``thinker``) to a list of axis declarations,
+each of which becomes one :class:`StrategySpec`. YAML is the only format wired up
+today, but the module/symbol names stay format-agnostic so a different source can
+be added without renaming the public surface::
+
+    strategies:
+      thinker:
+        - axis: tp
+          size: 2
+      talker:
+        - axis: stage_replica
+          size: 2
+          routing: round_robin   # random | round_robin | least_queue
+      code2wav:
+        - axis: stage_replica
+          size: 2
+          routing: round_robin
+
+Per-axis fields:
+
+* ``axis`` (required): mesh-axis kind (``tp``, ``dp``, ``pp``, ``ep``,
+  ``stage_replica`` are translatable today).
+* ``size`` (required): axis degree (> 0).
+* ``routing`` (optional): stateless policy for ``dp`` / ``stage_replica`` axes
+  (``random`` / ``round_robin`` / ``least_queue``). Ignored for axes whose
+  routing is fixed (``tp`` broadcast, ``pp`` microbatch).
+* ``l1_owner`` (optional): ``engine`` or ``delegated``; overrides the default
+  owner for the axis (recorded in ``shard_extension``).
+* ``name`` (optional): a label for the spec; defaults to the axis kind.
+
+Routing/aggregation patterns are filled in from sensible per-kind defaults so
+the file stays compact; the translator validates the combination.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from vllm_omni.config.composable_parallel.aggregation import (
+    AggregationPattern,
+    FanInByStage,
+    StitchPipeline,
+    TakeRank,
+    Union,
+)
+from vllm_omni.config.composable_parallel.routing import (
+    Broadcast,
+    PipelineMicrobatch,
+    RouteByStage,
+    RoutingPattern,
+)
+from vllm_omni.config.composable_parallel.spec import MeshAxisSpec, StrategySpec
+from vllm_omni.config.yaml_util import load_yaml_config, to_dict
+
+# Per-kind defaults for routing + aggregation, so the file only needs axis+size.
+_ROUTING_POLICY_KINDS = frozenset({"dp", "stage_replica"})
+
+
+class StrategyLoadError(ValueError):
+    """Raised when a strategy file is malformed."""
+
+
+def _default_routing(kind: str, routing_policy: str | None) -> RoutingPattern:
+    if kind in _ROUTING_POLICY_KINDS:
+        return RouteByStage(routing_policy=routing_policy or "round_robin")  # type: ignore[arg-type]
+    if kind == "pp":
+        return PipelineMicrobatch()
+    # tp, ep, and any other kind default to broadcast (the translator rejects
+    # kinds it cannot handle before it ever inspects routing).
+    return Broadcast()
+
+
+def _default_aggregation(kind: str) -> AggregationPattern:
+    if kind == "stage_replica":
+        return FanInByStage()
+    if kind == "pp":
+        return StitchPipeline()
+    if kind in ("dp", "ep"):
+        return Union()
+    return TakeRank()
+
+
+def _build_spec(role: str, entry: Mapping[str, Any]) -> StrategySpec:
+    if "axis" not in entry:
+        raise StrategyLoadError(f"role {role!r}: each strategy entry needs an 'axis' field")
+    if "size" not in entry:
+        raise StrategyLoadError(f"role {role!r}: strategy entry for axis {entry['axis']!r} needs a 'size'")
+
+    kind = str(entry["axis"])
+    try:
+        size = int(entry["size"])
+    except (TypeError, ValueError) as exc:
+        raise StrategyLoadError(f"role {role!r}: axis {kind!r} size must be an integer, got {entry['size']!r}") from exc
+
+    routing_policy = entry.get("routing")
+    if routing_policy is not None and kind not in _ROUTING_POLICY_KINDS:
+        raise StrategyLoadError(
+            f"role {role!r}: axis {kind!r} does not accept a 'routing' policy (only {sorted(_ROUTING_POLICY_KINDS)} do)"
+        )
+
+    shard_extension: dict[str, Any] = {}
+    l1_owner = entry.get("l1_owner")
+    if l1_owner is not None:
+        shard_extension["l1_owner"] = str(l1_owner)
+
+    return StrategySpec(
+        name=str(entry.get("name", kind)),
+        mesh_axis=MeshAxisSpec(kind=kind, size=size),  # type: ignore[arg-type]
+        routing=_default_routing(kind, routing_policy),
+        aggregation=_default_aggregation(kind),
+        shard_extension=shard_extension,
+    )
+
+
+def parse_strategy_specs(data: Mapping[str, Any]) -> dict[str, list[StrategySpec]]:
+    """Parse an already-loaded strategy mapping into per-role spec stacks."""
+    strategies = data.get("strategies", data)
+    if not isinstance(strategies, Mapping):
+        raise StrategyLoadError("strategy file must contain a 'strategies' mapping of role -> list of axes")
+
+    result: dict[str, list[StrategySpec]] = {}
+    for role, entries in strategies.items():
+        if not isinstance(entries, (list, tuple)):
+            raise StrategyLoadError(
+                f"role {role!r}: expected a list of axis declarations, got {type(entries).__name__}"
+            )
+        specs: list[StrategySpec] = []
+        for entry in entries:
+            if not isinstance(entry, Mapping):
+                raise StrategyLoadError(
+                    f"role {role!r}: each axis declaration must be a mapping, got {type(entry).__name__}"
+                )
+            specs.append(_build_spec(str(role), dict(entry)))
+        result[str(role)] = specs
+    return result
+
+
+def load_strategy_specs(path: str) -> dict[str, list[StrategySpec]]:
+    """Load and parse a strategy file into per-role spec stacks."""
+    data = to_dict(load_yaml_config(path))
+    if not isinstance(data, Mapping):
+        raise StrategyLoadError(f"strategy file {path!r} did not parse to a mapping")
+    return parse_strategy_specs(data)
diff --git a/vllm_omni/config/composable_parallel/translator.py b/vllm_omni/config/composable_parallel/translator.py
new file mode 100644
index 000000000..5c394dc9b
--- /dev/null
+++ b/vllm_omni/config/composable_parallel/translator.py
@@ -0,0 +1,378 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Translate a ``StrategySpec`` stack into vLLM-Omni parallel sizing.
+
+Read a stack of strategy specs (one per mesh axis) and work out the
+tensor/data/pipeline parallel sizes, the dense-EP flag, and the number of stage
+replicas, plus who owns each axis's request routing. The result is a plain,
+CPU-computable sizing struct (:class:`OmniParallelConfig`) keyed by the real
+``OmniEngineArgs`` / ``EngineArgs`` field names, so the deploy layer can splat it
+onto a stage's engine args.
+
+The distinction this module enforces — engine data parallelism (a true vLLM
+intra-engine world dimension) vs. omni stage replicas (independent engines fanned
+out by omni's coordinator) — is documented on the axis validators that depend on
+it (see :func:`_validate_dp` and :func:`_stage_replica_lb_policy`). Routing we do
+not support yet (key-stable / affinity routing) raises ``NotImplementedError``;
+any other invalid spec raises :class:`AxisTranslationError`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Literal, NoReturn, cast, get_args
+
+from vllm.logger import init_logger
+
+from vllm_omni.config.composable_parallel.routing import (
+    Broadcast,
+    PartitionByHash,
+    PipelineMicrobatch,
+    RouteByStage,
+    RoutingPattern,
+)
+from vllm_omni.config.composable_parallel.spec import MeshAxisKind, StrategySpec
+
+logger = init_logger(__name__)
+
+# Who owns an axis's request routing. A closed string type so an unexpected raw
+# value is caught statically and at runtime (``_VALID_L1_OWNERS`` is derived from
+# it, keeping one source of truth) rather than silently flowing through.
+L1Owner = Literal["delegated", "engine"]
+
+# DP, TP, PP, (dense) EP and stage_replica translate today. DP/TP/PP are true
+# world dimensions vLLM realizes intra-engine; dense EP is a flag over the
+# existing TP*DP ranks; stage_replica is an omni-coordinator-level fan-out of
+# independent engine replicas (NOT a vLLM world dimension). The remaining kinds
+# (sequence/context parallel, CFG, VAE pipelines, ...) need per-layer hooks or
+# custom collectives and arrive in later stages.
+_SUPPORTED_KINDS: tuple[MeshAxisKind, ...] = ("dp", "tp", "pp", "ep", "stage_replica")
+
+# Which EngineArgs field each axis kind *sizes*. EP and stage_replica are
+# deliberately absent: EP is not an independent world dimension but an
+# ``enable_expert_parallel`` flag whose degree equals tp*dp; stage_replica is a
+# per-stage deploy ``num_replicas`` count, not an intra-engine world dimension.
+_AXIS_TO_ENGINE_FIELD: dict[MeshAxisKind, str] = {
+    "tp": "tensor_parallel_size",
+    "dp": "data_parallel_size",
+    "pp": "pipeline_parallel_size",
+}
+
+# Default L1 owner per axis kind. ``stage_replica`` is the only delegated axis:
+# omni's coordinator/StagePool owns routing across replicas. ``dp`` is engine
+# data parallelism, realized by vLLM's own intra-engine DP load balancer.
+_DEFAULT_L1_OWNER: dict[MeshAxisKind, L1Owner] = {
+    "dp": "engine",
+    "tp": "engine",
+    "pp": "engine",
+    "ep": "engine",
+    "stage_replica": "delegated",
+}
+
+# How a RouteByStage policy maps onto omni's load balancer policy string. These
+# are the stateless options omni can actually do; note there's no "hash" here,
+# because omni has no key-stable balancer.
+_STAGE_POLICY_TO_OMNI_LB: dict[str, str] = {
+    "random": "random",
+    "round_robin": "round-robin",
+    "least_queue": "least-queue-length",
+}
+
+_VALID_L1_OWNERS: frozenset[str] = frozenset(get_args(L1Owner))
+
+
+class AxisTranslationError(ValueError):
+    """Error for an invalid or unsupported strategy spec.
+
+    A single error type (rather than a tree of subclasses) keeps the public
+    surface small and consistent with the rest of the codebase; the specific
+    cause is in the message and is logged before the raise so it is visible even
+    when the type is unavailable to a caller (e.g. across a server boundary).
+
+    Strategies that are *valid but not built yet* — key-stable / affinity routing
+    today — raise ``NotImplementedError`` instead, to distinguish "we haven't
+    implemented this" from "your config is wrong".
+    """
+
+
+def _fail(msg: str) -> NoReturn:
+    """Log and raise :class:`AxisTranslationError` for an invalid/unsupported spec."""
+    logger.error("[composable_parallel] %s", msg)
+    raise AxisTranslationError(msg)
+
+
+def _not_implemented(msg: str) -> NoReturn:
+    """Log and raise ``NotImplementedError`` for a valid-but-unbuilt strategy."""
+    logger.error("[composable_parallel] %s", msg)
+    raise NotImplementedError(msg)
+
+
+@dataclass(frozen=True)
+class OmniParallelConfig:
+    """Result of translating a spec stack into omni parallel sizing."""
+
+    tensor_parallel_size: int = 1
+    data_parallel_size: int = 1
+    pipeline_parallel_size: int = 1
+    # Dense expert parallelism: a flag over the TP*DP ranks, not a world axis.
+    enable_expert_parallel: bool = False
+    # Number of independent omni stage replicas (the ``stage_replica`` axis).
+    # This is NOT a vLLM world dimension; it is the per-stage deploy
+    # ``num_replicas`` count. 1 means a single (un-replicated) engine.
+    stage_replica_size: int = 1
+    # The omni StagePool LB policy string for the (delegated) stage_replica
+    # axis, if any. Only meaningful when ``stage_replica_size > 1``.
+    omni_lb_policy: str | None = None
+    # axis kind -> resolved L1 owner ("delegated" | "engine").
+    l1_owners: Mapping[MeshAxisKind, L1Owner] = field(default_factory=dict)
+
+    @property
+    def world_size(self) -> int:
+        # EP and stage_replica are intentionally excluded. EP reuses the TP*DP
+        # ranks rather than adding a dimension; stage_replica spins up separate
+        # engines (each its own world), not extra ranks in this engine's group.
+        # This product is what vLLM calls ``world_size_across_dp`` (tp * pp * dp).
+        return self.tensor_parallel_size * self.data_parallel_size * self.pipeline_parallel_size
+
+    @property
+    def delegated_axes(self) -> tuple[MeshAxisKind, ...]:
+        return tuple(kind for kind, owner in self.l1_owners.items() if owner == "delegated")
+
+    def as_engine_kwargs(self) -> dict[str, object]:
+        """Return per-stage kwargs keyed by real OmniEngineArgs/EngineArgs field names.
+
+        Two derived values are intentionally *not* emitted here because they are
+        not per-stage engine args:
+
+        * ``stage_replica_size`` — a per-stage deploy ``num_replicas`` knob
+          (``StageDeployConfig``); the deploy layer consumes it separately.
+        * ``omni_lb_policy`` — a pipeline-wide load-balancer policy the engine
+          reads once at construction (see ``StrategyApplyResult.omni_lb_policy``);
+          it is applied at the orchestrator level, not folded into per-stage args.
+        """
+        kwargs: dict[str, object] = {
+            "tensor_parallel_size": self.tensor_parallel_size,
+            "data_parallel_size": self.data_parallel_size,
+            "pipeline_parallel_size": self.pipeline_parallel_size,
+        }
+        if self.enable_expert_parallel:
+            kwargs["enable_expert_parallel"] = True
+        return kwargs
+
+
+def _is_affinity_dp_routing(routing: RoutingPattern) -> bool:
+    """True when DP routing demands key-stable (hash) placement."""
+    if isinstance(routing, PartitionByHash):
+        return True
+    if isinstance(routing, RouteByStage) and routing.routing_policy == "hash":
+        return True
+    return False
+
+
+def _resolve_l1_owner(spec: StrategySpec) -> L1Owner:
+    kind = spec.mesh_axis.kind
+    declared_owner = spec.shard_extension.get("l1_owner")
+    owner = declared_owner
+    if owner is None:
+        owner = _DEFAULT_L1_OWNER.get(kind, "engine")
+        if kind not in _DEFAULT_L1_OWNER:
+            logger.debug(
+                "[composable_parallel] axis %r has unknown kind %r with no default l1_owner; falling back to %r",
+                spec.name,
+                kind,
+                owner,
+            )
+        else:
+            logger.debug(
+                "[composable_parallel] axis %r (kind %r) declared no l1_owner; using default %r",
+                spec.name,
+                kind,
+                owner,
+            )
+    if owner not in _VALID_L1_OWNERS:
+        # Distinguish a bad value supplied via shard_extension from a bad default
+        # indexed out of _DEFAULT_L1_OWNER, so the source of the invalid owner is
+        # debuggable rather than swallowed.
+        source = "shard_extension" if declared_owner is not None else "_DEFAULT_L1_OWNER"
+        logger.debug(
+            "[composable_parallel] axis %r (kind %r) resolved invalid l1_owner %r from %s; valid owners are %s",
+            spec.name,
+            kind,
+            owner,
+            source,
+            sorted(_VALID_L1_OWNERS),
+        )
+        _fail(f"axis {kind!r} has invalid l1_owner {owner!r}; expected one of {sorted(_VALID_L1_OWNERS)}")
+    return cast(L1Owner, owner)
+
+
+def _validate_dp(spec: StrategySpec, owner: L1Owner) -> None:
+    """Validate an engine data-parallel axis.
+
+    DP is realized intra-engine by vLLM's own DP load balancer, so it is
+    engine-owned and emits no ``omni_lb_policy`` (that string configures omni's
+    *replica* balancer, a different layer). Key-stable (hash) request affinity is
+    not something vLLM's DP LB guarantees, so it is rejected rather than silently
+    dropped.
+    """
+    if owner != "engine":
+        _fail(
+            f"dp axis {spec.name!r} is engine data parallelism realized intra-engine by "
+            f"vLLM's DP load balancer; l1_owner must be 'engine', got {owner!r}. For "
+            "omni-coordinator-level request fan-out across independent replicas, use a "
+            "'stage_replica' axis."
+        )
+    if _is_affinity_dp_routing(spec.routing):
+        _not_implemented(
+            f"dp axis {spec.name!r} requests key-stable (hash) routing, which vLLM's "
+            "intra-engine DP load balancer does not guarantee — not supported yet. Use "
+            "RouteByStage(random|round_robin|least_queue) for stateless DP balancing."
+        )
+    if not isinstance(spec.routing, RouteByStage):
+        _fail(
+            f"dp axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) routing, "
+            f"got {type(spec.routing).__name__}"
+        )
+    if spec.routing.routing_policy not in _STAGE_POLICY_TO_OMNI_LB:
+        # Recognized-but-unimplemented routing (key-stable/hash) is already
+        # handled above via _is_affinity_dp_routing -> _not_implemented. Anything
+        # left here is an unknown/invalid policy value (e.g. a typo from YAML),
+        # i.e. invalid input -> AxisTranslationError, not NotImplementedError.
+        _fail(
+            f"dp axis {spec.name!r} has invalid routing_policy "
+            f"{spec.routing.routing_policy!r}; expected one of "
+            f"{sorted(_STAGE_POLICY_TO_OMNI_LB)}."
+        )
+
+
+def _stage_replica_lb_policy(spec: StrategySpec, owner: L1Owner) -> str:
+    """Return the omni StagePool LB policy for a delegated stage_replica axis.
+
+    A ``stage_replica`` axis is *not* a vLLM world dimension: it stands up N
+    independent engine replicas of one pipeline stage, coordinated by the omni
+    coordinator and balanced over a StagePool with stateless policies only
+    (random / round-robin / least-queue-length). It maps to the per-stage
+    ``num_replicas`` count plus the pipeline-level ``omni_lb_policy`` string, so
+    its ``l1_owner`` must be ``"delegated"`` (omni owns the routing). Key-stable
+    (hash) routing is rejected because omni has no key-stable balancer yet.
+    """
+    if owner != "delegated":
+        _fail(
+            f"stage_replica axis {spec.name!r} must be 'delegated' to omni's StagePool load "
+            f"balancer; got owner {owner!r}. Replica routing is owned by omni's coordinator."
+        )
+
+    routing = spec.routing
+    if _is_affinity_dp_routing(routing):
+        _not_implemented(
+            f"stage_replica axis {spec.name!r} requests key-stable (hash) routing, which needs a "
+            "dedicated load balancer — not implemented yet. Use "
+            "RouteByStage(random|round_robin|least_queue) to delegate to omni's load balancer."
+        )
+    if not isinstance(routing, RouteByStage):
+        _fail(
+            f"stage_replica axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) "
+            f"routing, got {type(routing).__name__}"
+        )
+    policy = _STAGE_POLICY_TO_OMNI_LB.get(routing.routing_policy)
+    if policy is None:
+        _not_implemented(
+            f"stage_replica axis {spec.name!r} routing_policy {routing.routing_policy!r} has no omni LB policy"
+        )
+    return policy
+
+
+def _validate_tp(spec: StrategySpec, owner: L1Owner) -> None:
+    if not isinstance(spec.routing, Broadcast):
+        _fail(f"tp axis {spec.name!r} expects Broadcast routing, got {type(spec.routing).__name__}")
+    if owner != "engine":
+        _fail(f"tp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
+
+
+def _validate_pp(spec: StrategySpec, owner: L1Owner) -> None:
+    if not isinstance(spec.routing, PipelineMicrobatch):
+        _fail(f"pp axis {spec.name!r} expects PipelineMicrobatch routing, got {type(spec.routing).__name__}")
+    if owner != "engine":
+        _fail(f"pp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
+
+
+def _validate_ep(spec: StrategySpec, owner: L1Owner) -> None:
+    # Dense EP: every rank still sees the whole batch, experts are sharded
+    # across ranks inside the engine (sparse MoE all-to-all is a later stage).
+    if not isinstance(spec.routing, Broadcast):
+        _fail(
+            f"ep axis {spec.name!r} expects Broadcast routing (dense expert parallel), "
+            f"got {type(spec.routing).__name__}"
+        )
+    if owner != "engine":
+        _fail(f"ep axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}")
+
+
+def translate_strategy_stack(specs: Sequence[StrategySpec]) -> OmniParallelConfig:
+    """Translate a spec stack into an ``OmniParallelConfig``.
+
+    Supported kinds: dp (engine data parallel), tp, pp, (dense) ep, and
+    stage_replica (omni replicas). Raises ``NotImplementedError`` for deferred
+    (affinity / key-stable) routing, and :class:`AxisTranslationError` for any
+    other invalid spec (a kind not yet translatable, a repeated kind, an owner
+    incompatible with the axis kind, or unsupported routing). The EP degree must
+    equal tensor_parallel_size * data_parallel_size.
+    """
+    sizes: dict[str, int] = {"tensor_parallel_size": 1, "data_parallel_size": 1, "pipeline_parallel_size": 1}
+    owners: dict[MeshAxisKind, L1Owner] = {}
+    omni_lb_policy: str | None = None
+    stage_replica_size = 1
+    enable_expert_parallel = False
+    ep_size: int | None = None
+
+    for spec in specs:
+        kind = spec.mesh_axis.kind
+        if kind not in _SUPPORTED_KINDS:
+            _fail(
+                f"axis kind {kind!r} is not translatable yet (supported: {list(_SUPPORTED_KINDS)}); "
+                "it is designed-for and lands in a later stage"
+            )
+        if kind in owners:
+            _fail(f"axis kind {kind!r} appears more than once in the spec stack")
+
+        owner = _resolve_l1_owner(spec)
+        if kind == "dp":
+            _validate_dp(spec, owner)
+        elif kind == "tp":
+            _validate_tp(spec, owner)
+        elif kind == "pp":
+            _validate_pp(spec, owner)
+        elif kind == "ep":
+            _validate_ep(spec, owner)
+            enable_expert_parallel = True
+            ep_size = spec.mesh_axis.size
+        elif kind == "stage_replica":
+            omni_lb_policy = _stage_replica_lb_policy(spec, owner)
+            stage_replica_size = spec.mesh_axis.size
+
+        if kind in _AXIS_TO_ENGINE_FIELD:
+            sizes[_AXIS_TO_ENGINE_FIELD[kind]] = spec.mesh_axis.size
+        owners[kind] = owner
+
+    if enable_expert_parallel:
+        # Dense EP shards experts across exactly the TP*DP ranks, so the declared
+        # EP degree must match that product (it is not its own world dimension,
+        # and PP is excluded). An EP axis of size 1 is a degenerate no-op that
+        # still sets the flag; it is only valid when TP*DP == 1.
+        ep_ranks = sizes["tensor_parallel_size"] * sizes["data_parallel_size"]
+        if ep_size != ep_ranks:
+            _fail(
+                f"ep axis size {ep_size} must equal tensor_parallel_size*data_parallel_size "
+                f"(={ep_ranks}); dense expert parallelism shards experts across exactly those ranks"
+            )
+
+    return OmniParallelConfig(
+        tensor_parallel_size=sizes["tensor_parallel_size"],
+        data_parallel_size=sizes["data_parallel_size"],
+        pipeline_parallel_size=sizes["pipeline_parallel_size"],
+        enable_expert_parallel=enable_expert_parallel,
+        stage_replica_size=stage_replica_size,
+        omni_lb_policy=omni_lb_policy,
+        l1_owners=owners,
+    )
diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py
index 20c32c5bc..1b38a3d34 100644
--- a/vllm_omni/config/config_factory.py
+++ b/vllm_omni/config/config_factory.py
@@ -5,6 +5,8 @@
 from __future__ import annotations
 
 import dataclasses
+import functools
+from collections.abc import Mapping
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any
@@ -14,6 +16,7 @@ from vllm.logger import init_logger
 from vllm.transformers_utils.config import get_config
 from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
 
+from vllm_omni.config.endpoint_policy import EndpointRestriction
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
     _DEPLOY_DIR,
@@ -31,6 +34,16 @@ from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
 logger = init_logger(__name__)
 
 
+# Default degree for any parallel axis / replica count that isn't set anywhere
+# (CLI, deploy YAML, or pipeline default): a single, un-parallelized rank.
+# TODO(composable_parallel): this "1" default and the per-axis device-layout
+# fallbacks are currently re-derived in the merge, apply, and reconcile layers.
+# Centralize them in one schema so the pre-spawn device guard and the engine-args
+# defaults can't drift apart. This is the light slice; the full device-layout
+# centralization is tracked as a follow-up.
+_DEFAULT_PARALLEL_DEGREE = 1
+
+
 class StageConfigFactory:
     """Factory that loads pipeline YAML and merges CLI overrides.
 
@@ -46,83 +59,174 @@ class StageConfigFactory:
     """
 
     @classmethod
-    def create_from_model(
+    def get_pipeline_endpoint_restrictions(
         cls,
         model: str,
-        cli_overrides: dict[str, Any] | None = None,
-        deploy_config_path: str | None = None,
-        **deprecated_kwargs: Any,
-    ) -> list[StageConfig] | None:
-        """Load pipeline + deploy config, merge with CLI overrides.
+        trust_remote_code: bool,
+        deploy_config_path: str | None,
+    ) -> tuple[EndpointRestriction, ...]:
+        """Given a model string, determine the corresponding endpoint restrictions.
 
-        Checks OMNI_PIPELINES first, since supported models should be explicitly
-        registered. If a model is not registered in OMNI_PIPELINES, tries to fall
-        back to using the Transformers config & finding pipelines that have overlapping
-        supported architectures.
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+            deploy_config_path: Optional path to the deploy config for the pipeline.
+
+        Returns:
+            A tuple of model specific endpoint restrictions.
         """
-        if cli_overrides is None:
-            cli_overrides = {}
+        pipeline_cfg = StageConfigFactory.get_pipeline_config(
+            model=model,
+            trust_remote_code=trust_remote_code,
+            deploy_config_path=deploy_config_path,
+        )
+        return pipeline_cfg.endpoint_restrictions if pipeline_cfg else ()
 
-        trust_remote_code = cli_overrides.get("trust_remote_code", True)
-        if trust_remote_code is None:
-            trust_remote_code = False
+    @classmethod
+    @functools.cache
+    def get_hf_config(cls, model: str, trust_remote_code: bool) -> PretrainedConfig | None:
+        """Fetch the HF config (if it exists) from the model directory.
 
-        model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            the model's config or None.
+        """
+        hf_config = None
+        try:
+            return get_config(model, trust_remote_code=trust_remote_code)
+        except Exception as e:
+            logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None")
+        return hf_config
+
+    @classmethod
+    @functools.cache
+    def try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None:
+        """Auto-detect model_type from model directory and apply any model
+        specific patches to get the correct model_type str. If we are unable
+        to infer it from the model directory, we fall back to the PipelineConfig.
+
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            model_type as a string; may be None on failure.
+        """
+        model_type = cls._try_infer_model_type(
+            model=model,
+            trust_remote_code=trust_remote_code,
+        )
         if model_type == "vla":
             if _looks_like_dreamzero(model):
                 model_type = "dreamzero"
+        return model_type
 
-        # --- 1. Explicit deploy-config pipeline override (highest precedence) ---
-        # A deploy YAML may set ``pipeline: <model_type>`` to force routing for
-        # models whose HF config has a generic/colliding ``model_type`` or no
-        # matching architectures. Ming-omni-tts, for example, reports
-        # model_type="dense" with arch BailingMMNative…, which matches no
-        # pipeline — without this it falls through to the diffusion default and
-        # dies with "Model class BailingMMNativeForConditionalGeneration not
-        # found in diffusion model registry". Honor the key before auto-detection.
-        explicit_pipeline = None
-        if deploy_config_path:
-            try:
-                explicit_pipeline = load_deploy_config(deploy_config_path).pipeline
-            except Exception:
-                logger.exception("Failed to read 'pipeline' key from deploy config %s", deploy_config_path)
-        if explicit_pipeline:
-            pipeline_cfg = cls.resolve_pipeline_config(explicit_pipeline, hf_config)
-            if pipeline_cfg is not None:
-                return cls._create_from_registry(
-                    explicit_pipeline,
-                    pipeline_cfg,
-                    cli_overrides,
-                    deploy_config_path,
-                )
-            logger.warning(
-                "Deploy config %s requested pipeline %r which is not in OMNI_PIPELINES; "
-                "falling back to auto-detection.",
-                deploy_config_path,
-                explicit_pipeline,
-            )
+    @classmethod
+    def _try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None:
+        """Auto-detect model_type from model directory.
 
-        # --- 2. Auto-detected model_type registered in OMNI_PIPELINES ---
+        Args:
+            model: Model name or path.
+            trust_remote_code: Whether to trust remote code for HF config loading.
+
+        Returns:
+            model_type as a string; may be None on failure.
+        """
+        hf_config = cls.get_hf_config(
+            model=model,
+            trust_remote_code=trust_remote_code,
+        )
+        if hf_config is not None:
+            return hf_config.model_type
+
+        # Fallback: read config.json directly for custom model types that
+        # are not registered with transformers (e.g. qwen3_tts).
+        try:
+            config_dict = get_hf_file_to_dict("config.json", model, revision=None)
+            if config_dict:
+                if "model_type" in config_dict:
+                    return config_dict["model_type"]
+                # VoxCPM2-style configs use singular ``architecture`` rather
+                # than HF's standard ``model_type`` / ``architectures``. Accept
+                # it as a fallback so the pipeline registry can still match.
+                if "architecture" in config_dict and isinstance(config_dict["architecture"], str):
+                    return config_dict["architecture"]
+        except Exception as e:
+            logger.debug(f"Failed to auto-detect model type for {model}: {e}")
+
+        # Fallback for diffusers-style models: check model_index.json.
+        # Some models (e.g. GLM-Image) have no root config.json but ship a
+        # model_index.json with _class_name that maps to a pipeline key via
+        # PipelineConfig.diffusers_class_name.
+        try:
+            model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
+            if model_index and "_class_name" in model_index:
+                class_name = model_index["_class_name"]
+                for obj in OMNI_PIPELINES.values():
+                    # If we have a resolver, call it with the optional hf_config
+                    # to get the default pipeline config for this key
+                    pipeline_cfg = obj(hf_config) if callable(obj) else obj
+                    if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
+                        logger.info(
+                            "Detected pipeline %r from model_index.json (_class_name=%r)",
+                            pipeline_cfg.model_type,
+                            class_name,
+                        )
+                        return pipeline_cfg.model_type
+        except Exception as e:
+            logger.debug(f"Failed to detect model type for diffusers-style models: {e}")
+
+        # Final fallback: some models (e.g. CosyVoice3) ship an empty
+        # config.json and rely on naming conventions. Match the model path
+        # basename against registered pipeline keys — longest match wins
+        # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
+        model_lower = model.lower().replace("-", "").replace("_", "")
+        best: str | None = None
+        best_len = 0
+        for registered_key in OMNI_PIPELINES.keys():
+            candidate = registered_key.lower().replace("-", "").replace("_", "")
+            if candidate and candidate in model_lower and len(candidate) > best_len:
+                best = registered_key
+                best_len = len(candidate)
+        if best is not None:
+            return best
+
+        return None
+
+    @classmethod
+    @functools.cache
+    def get_pipeline_config(
+        cls,
+        model: str,
+        trust_remote_code: bool,
+        deploy_config_path: str | None = None,
+    ) -> PipelineConfig | None:
+        """Resolve the PipelineConfig for a model path/name."""
+        model_type = cls.try_infer_model_type(model=model, trust_remote_code=trust_remote_code)
+        hf_config = cls.get_hf_config(model=model, trust_remote_code=trust_remote_code)
+
+        # Resolve the deploy config & check if the user set the pipeline;
+        # If the pipeline is explicitly set, it takes highest priority
+        deploy_config_pipe = cls._get_deploy_override_pipe_config(hf_config, deploy_config_path)
+        if deploy_config_pipe is not None:
+            return deploy_config_pipe
+
+        # Pipeline isn't set in the yaml spec, so we need infer it ourselves.
         if model_type and model_type in OMNI_PIPELINES:
             pipeline_cfg = cls.resolve_pipeline_config(model_type, hf_config)
             if pipeline_cfg is not None:
-                return cls._create_from_registry(
-                    model_type,
-                    pipeline_cfg,
-                    cli_overrides,
-                    deploy_config_path,
-                )
+                return pipeline_cfg
 
-        # --- HF architecture fallback: some models report a generic
-        # model_type that collides with another model. Match by the
-        # hf_architectures declared on each registered PipelineConfig.
         if hf_config is not None:
-            logger.warning("Inferred model type %s is not registered to an Omni pipeline", model_type)
+            if model_type is not None:
+                logger.warning("Inferred model type %s is not registered to an Omni pipeline", model_type)
             hf_archs = set(getattr(hf_config, "architectures", []) or [])
             if hf_archs:
                 for registered in OMNI_PIPELINES.values():
                     pipeline_cfg = registered if isinstance(registered, PipelineConfig) else registered(hf_config)
-                    # Resolvers that get configs of the incorrect type should return None
                     if pipeline_cfg is None:
                         continue
                     predicate = pipeline_cfg.hf_config_predicate
@@ -143,38 +247,77 @@ class StageConfigFactory:
                                 pipeline_cfg.model_type,
                             )
                             continue
-
                     if isinstance(pipeline_cfg, PipelineConfig) and hf_archs.intersection(
                         pipeline_cfg.hf_architectures
                     ):
-                        return cls._create_from_registry(
-                            pipeline_cfg.model_type,
-                            pipeline_cfg,
-                            cli_overrides,
-                            deploy_config_path,
-                        )
-
-        # --- Explicit deploy-config pipeline ---
-        # When auto-detection above resolves nothing (generic/missing HF model_type
-        # and no matching architecture), honor an explicit pipeline key in the deploy config
-        if deploy_config_path is not None:
-            deploy_path = Path(deploy_config_path)
-            if deploy_path.exists():
-                deploy_cfg = load_deploy_config(deploy_path)
-                if deploy_cfg.pipeline:
-                    pipeline_cfg = cls.resolve_pipeline_config(deploy_cfg.pipeline, hf_config)
-                    if pipeline_cfg is not None:
-                        return cls._create_from_registry(
-                            pipeline_cfg.model_type,
-                            pipeline_cfg,
-                            cli_overrides,
-                            deploy_config_path,
-                        )
-
-        # Not in the pipeline registry — let the caller fall back to the
-        # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
+                        return pipeline_cfg
         return None
 
+    @classmethod
+    def _get_deploy_override_pipe_config(
+        cls,
+        hf_config: PretrainedConfig | None,
+        deploy_config_path: str | None,
+    ) -> PipelineConfig | None:
+        """Load the deploy config and extract + resolve its pipeline field."""
+        if deploy_config_path is None:
+            return None
+
+        deploy_path = Path(deploy_config_path)
+        if deploy_path.exists():
+            deploy_cfg = load_deploy_config(deploy_path)
+            if deploy_cfg.pipeline is not None:
+                return cls.resolve_pipeline_config(deploy_cfg.pipeline, hf_config)
+        return None
+
+    @classmethod
+    def create_from_model(
+        cls,
+        model: str,
+        *,
+        trust_remote_code: bool = False,
+        cli_overrides: dict[str, Any] | None = None,
+        deploy_config_path: str | None = None,
+        strategy_specs: Mapping[Any, Any] | None = None,
+        **deprecated_kwargs: Any,
+    ) -> tuple[list[StageConfig] | None, str | None]:
+        """Load pipeline + deploy config, merge with CLI overrides.
+
+        Checks OMNI_PIPELINES first, since supported models should be explicitly
+        registered. If a model is not registered in OMNI_PIPELINES, tries to fall
+        back to using the Transformers config & finding pipelines that have overlapping
+        supported architectures.
+
+        When ``strategy_specs`` is provided (a mapping of role -> list of
+        ``StrategySpec``), the derived parallel sizing is overlaid onto the
+        merged stages (see ``vllm_omni.config.composable_parallel``). This is
+        opt-in: omitting it leaves the existing merge path untouched.
+
+        Returns ``(stages, omni_lb_policy)``: the merged stages (``None`` when the
+        model is not in the pipeline registry and the caller should fall back to
+        the legacy YAML path) and the strategy-derived, pipeline-wide
+        ``omni_lb_policy`` (``None`` when no stage_replica axis set one). The
+        policy is returned rather than threaded through a mutable out-param so the
+        engine can apply it without every intermediate call carrying the dict.
+        """
+        if cli_overrides is None:
+            cli_overrides = {}
+
+        pipeline_cfg = cls.get_pipeline_config(
+            model=model,
+            trust_remote_code=trust_remote_code,
+            deploy_config_path=deploy_config_path,
+        )
+        if pipeline_cfg is not None:
+            return cls._create_from_registry(
+                pipeline_cfg.model_type,
+                pipeline_cfg,
+                cli_overrides,
+                deploy_config_path,
+                strategy_specs=strategy_specs,
+            )
+        return None, None
+
     @classmethod
     def _create_from_registry(
         cls,
@@ -182,12 +325,17 @@ class StageConfigFactory:
         pipeline_cfg: PipelineConfig,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None = None,
+        strategy_specs: Mapping[Any, Any] | None = None,
         **deprecated_kwargs: Any,
-    ) -> list[StageConfig]:
+    ) -> tuple[list[StageConfig], str | None]:
         """Create StageConfigs from pipeline registry + deploy YAML.
 
         Precedence: caller-typed (non-None) value > deploy YAML >
         StageDeployConfig dataclass default.
+
+        Returns ``(stages, omni_lb_policy)`` — the strategy-derived pipeline-wide
+        load-balancer policy (``None`` when no strategy set one) travels with the
+        stages instead of through a mutable out-param.
         """
         # Resolve deploy config path
         if deploy_config_path is None:
@@ -220,12 +368,127 @@ class StageConfigFactory:
 
         stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides)
 
+        # Overlay declarative parallel strategies (opt-in) before CLI overrides.
+        applied = cls._apply_strategy_specs(stages, strategy_specs)
+
         explicit_overrides = {k: v for k, v in cli_overrides.items() if v is not None}
 
         for stage in stages:
             stage.runtime_overrides = cls._merge_cli_overrides(stage, explicit_overrides)
 
-        return stages
+        # Re-validate the resolved layout now that CLI overrides are on top.
+        cls._reconcile_strategy_with_cli(stages, applied)
+
+        omni_lb_policy = applied.omni_lb_policy if applied is not None else None
+        return stages, omni_lb_policy
+
+    @staticmethod
+    def _apply_strategy_specs(
+        stages: list[StageConfig],
+        strategy_specs: Mapping[Any, Any] | None,
+    ) -> Any:
+        """Overlay derived parallel sizing onto merged stages (opt-in).
+
+        ``omni_lb_policy`` cannot be set from stage configs (omni reads it once
+        at orchestrator construction), so a derived non-default policy is logged
+        here and carried on the returned ``StrategyApplyResult`` for the caller
+        to hand to the orchestrator.
+
+        Returns the ``StrategyApplyResult`` (or ``None`` when no strategy was
+        supplied) so the caller can re-validate the resolved layout once CLI
+        overrides have been merged on top, and read its ``omni_lb_policy``.
+        """
+        if not strategy_specs:
+            return None
+        from vllm_omni.config.composable_parallel import apply_strategy_specs
+
+        applied = apply_strategy_specs(stages, strategy_specs)
+        if applied.omni_lb_policy is not None:
+            logger.info(
+                "[composable_parallel] strategy derived omni_lb_policy=%r; it will be applied "
+                "to the orchestrator unless an explicit --omni-lb-policy was given.",
+                applied.omni_lb_policy,
+            )
+        return applied
+
+    @staticmethod
+    def _reconcile_strategy_with_cli(
+        stages: list[StageConfig],
+        applied: Any,
+    ) -> None:
+        """Reconcile CLI overrides applied *after* a strategy overlay.
+
+        CLI overrides are applied last (in ``to_omegaconf``), so for any stage a
+        strategy touched we must (a) warn loudly when a CLI arg overrides a
+        strategy-declared axis — the strategy is meant to be the single writer
+        for the axes it declares, so a silent CLI win is surprising — and
+        (b) re-run the device-count check against the *effective* (post-CLI)
+        ``tp``/``dp``/``pp``/``num_replicas``/``devices`` so a CLI
+        ``--devices``/``--tensor-parallel-size``/``--num-replicas`` cannot slip
+        past the pre-spawn guard that exists to prevent silent OOMs.
+        """
+        if applied is None:
+            return
+        from vllm_omni.config.composable_parallel import check_device_layout
+
+        # Axis kind -> (engine-arg field, strategy-derived attribute on the cfg).
+        axis_fields = {
+            "tp": ("tensor_parallel_size", "tensor_parallel_size"),
+            "dp": ("data_parallel_size", "data_parallel_size"),
+            "pp": ("pipeline_parallel_size", "pipeline_parallel_size"),
+        }
+
+        for stage in stages:
+            cfg = applied.per_stage_config.get(stage.stage_id)
+            if cfg is None:
+                continue
+            overrides = stage.runtime_overrides or {}
+            declared = set(cfg.l1_owners.keys())
+
+            for kind, (field_name, attr) in axis_fields.items():
+                if kind in declared and overrides.get(field_name) is not None:
+                    cli_val = overrides[field_name]
+                    derived = getattr(cfg, attr)
+                    if cli_val != derived:
+                        logger.warning(
+                            "[composable_parallel] stage %s: CLI %s=%s overrides the "
+                            "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.",
+                            stage.stage_id,
+                            field_name,
+                            cli_val,
+                            field_name,
+                            derived,
+                        )
+            if "stage_replica" in declared and overrides.get("num_replicas") is not None:
+                cli_val = overrides["num_replicas"]
+                if cli_val != cfg.stage_replica_size:
+                    logger.warning(
+                        "[composable_parallel] stage %s: CLI num_replicas=%s overrides the "
+                        "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.",
+                        stage.stage_id,
+                        cli_val,
+                        cfg.stage_replica_size,
+                    )
+
+            def _eff(field_name: str, fallback: Any) -> Any:
+                val = overrides.get(field_name)
+                return val if val is not None else fallback
+
+            def _eff_degree(field_name: str, source: dict[str, Any]) -> int:
+                # Single place the per-axis "default to 1" fallback is applied,
+                # for both the override and the YAML-default sides (see the
+                # _DEFAULT_PARALLEL_DEGREE TODO).
+                value = _eff(field_name, source.get(field_name, _DEFAULT_PARALLEL_DEGREE))
+                return int(value or _DEFAULT_PARALLEL_DEGREE)
+
+            check_device_layout(
+                _eff("devices", stage.yaml_runtime.get("devices")),
+                tensor_parallel_size=_eff_degree("tensor_parallel_size", stage.yaml_engine_args),
+                data_parallel_size=_eff_degree("data_parallel_size", stage.yaml_engine_args),
+                pipeline_parallel_size=_eff_degree("pipeline_parallel_size", stage.yaml_engine_args),
+                num_replicas=_eff_degree("num_replicas", stage.yaml_runtime),
+                role=stage.model_stage,
+            )
 
     @classmethod
     def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
@@ -288,79 +551,6 @@ class StageConfigFactory:
 
         return [config_dict]
 
-    @classmethod
-    def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
-        """Auto-detect model_type from model directory.
-
-        Args:
-            model: Model name or path.
-            trust_remote_code: Whether to trust remote code for HF config loading.
-
-        Returns:
-            Tuple of (model_type, hf_config). Both may be None on failure.
-        """
-        hf_config = None
-
-        try:
-            hf_config = get_config(model, trust_remote_code=trust_remote_code)
-            return hf_config.model_type, hf_config
-        except Exception as e:
-            logger.debug(f"`get_config` failed for {e}; Falling back to raw config.json path")
-
-        # Fallback: read config.json directly for custom model types that
-        # are not registered with transformers (e.g. qwen3_tts).
-        try:
-            config_dict = get_hf_file_to_dict("config.json", model, revision=None)
-            if config_dict:
-                if "model_type" in config_dict:
-                    return config_dict["model_type"], None
-                # VoxCPM2-style configs use singular ``architecture`` rather
-                # than HF's standard ``model_type`` / ``architectures``. Accept
-                # it as a fallback so the pipeline registry can still match.
-                if "architecture" in config_dict and isinstance(config_dict["architecture"], str):
-                    return config_dict["architecture"], None
-        except Exception as e:
-            logger.debug(f"Failed to auto-detect model type for {model}: {e}")
-
-        # Fallback for diffusers-style models: check model_index.json.
-        # Some models (e.g. GLM-Image) have no root config.json but ship a
-        # model_index.json with _class_name that maps to a pipeline key via
-        # PipelineConfig.diffusers_class_name.
-        try:
-            model_index = get_hf_file_to_dict("model_index.json", model, revision=None)
-            if model_index and "_class_name" in model_index:
-                class_name = model_index["_class_name"]
-                for obj in OMNI_PIPELINES.values():
-                    # If we have a resolver, call it with the optional hf_config
-                    # to get the default pipeline config for this key
-                    pipeline_cfg = obj(hf_config) if callable(obj) else obj
-                    if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name:
-                        logger.info(
-                            "Detected pipeline %r from model_index.json (_class_name=%r)",
-                            pipeline_cfg.model_type,
-                            class_name,
-                        )
-                        return pipeline_cfg.model_type, None
-        except Exception as e:
-            logger.debug(f"Failed to detect model type for diffusers-style models: {e}")
-
-        # Final fallback: some models (e.g. CosyVoice3) ship an empty
-        # config.json and rely on naming conventions. Match the model path
-        # basename against registered pipeline keys — longest match wins
-        # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9).
-        model_lower = model.lower().replace("-", "").replace("_", "")
-        best: str | None = None
-        best_len = 0
-        for registered_key in OMNI_PIPELINES.keys():
-            candidate = registered_key.lower().replace("-", "").replace("_", "")
-            if candidate and candidate in model_lower and len(candidate) > best_len:
-                best = registered_key
-                best_len = len(candidate)
-        if best is not None:
-            return best, None
-
-        return None, None
-
     @classmethod
     def _merge_cli_overrides(
         cls,
@@ -381,6 +571,7 @@ class StageConfigFactory:
         """Given a model type, resolve to the pipeline to be used. If the pipeline
         maps to a callable we resolve based on the HF config."""
         if model_type not in OMNI_PIPELINES:
+            logger.warning("Model type %s is not registered to OMNI_PIPELINES", model_type)
             return None
         obj = OMNI_PIPELINES[model_type]
         return obj(hf_config) if callable(obj) else obj
diff --git a/vllm_omni/config/endpoint_policy.py b/vllm_omni/config/endpoint_policy.py
new file mode 100644
index 000000000..5a1449614
--- /dev/null
+++ b/vllm_omni/config/endpoint_policy.py
@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Endpoint restriction policy for omni pipelines."""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import NamedTuple
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from vllm.entrypoints.serve.utils.error_response import create_error_response
+
+
+class RouteTarget(NamedTuple):
+    """A server path & supported methods."""
+
+    path: str
+    methods: frozenset[str]
+
+
+class OmniServingCapability(Enum):
+    """Serving capabilities that pipelines can shut down."""
+
+    CHAT_COMPLETIONS_BATCH = RouteTarget("/v1/chat/completions/batch", frozenset({"POST"}))
+    COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"}))
+
+    @property
+    def path(self) -> str:
+        return self.value.path
+
+    @property
+    def methods(self) -> frozenset[str]:
+        return self.value.methods
+
+
+@dataclass(frozen=True)
+class EndpointRestriction:
+    capability: OmniServingCapability
+    reason: str
+
+
+# Routes that are not supported for any model, but are supported in vLLM.
+# This is only temporary to avoid 500s for batched chat completions.
+UNSUPPORTED_ROUTES: tuple[EndpointRestriction, ...] = (
+    EndpointRestriction(
+        OmniServingCapability.CHAT_COMPLETIONS_BATCH,
+        "Batched chat completions are not yet supported by vLLM Omni.",
+    ),
+)
+
+
+def build_rejection_handler(reason: str):
+    """Build a rejection handler for a given endpoint for the provided reason."""
+
+    async def rejection_handler(raw_request: Request):
+        error = create_error_response(message=reason)
+        return JSONResponse(
+            content=error.model_dump(),
+            status_code=error.error.code,
+        )
+
+    return rejection_handler
+
+
+def shutdown_unsupported_routes(
+    app: FastAPI,
+    endpoint_restrictions: tuple[EndpointRestriction, ...],
+):
+    """Given an initialized FastAPI server instance and a set of model specific endpoint
+    restrictions, remove the restricted routes and patch a handler that returns 400.
+    """
+    from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app
+
+    # Generally these should not overlap since there is no point. If they do,
+    # we use the reason message in UNSUPPORTED_ROUTES, for consistent error messages.
+    restricted_endpoints = (*endpoint_restrictions, *UNSUPPORTED_ROUTES)
+
+    for end_restrict in restricted_endpoints:
+        capability = end_restrict.capability
+        # Remove the route from the app
+        _remove_route_from_app(app, capability.path, capability.methods)
+
+        # Patch the bad request error with the model specific
+        # reason for shutting down this endpoint
+        rejection_handler = build_rejection_handler(end_restrict.reason)
+
+        app.add_api_route(
+            capability.path,
+            rejection_handler,
+            methods=list(capability.methods),
+        )
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index c44bd3ea9..256dcd5c7 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -469,6 +469,7 @@ class _DiffusionConfigProjection:
     diffusion_kv_cache_skip_layers: str | list[int] | tuple[int, ...] | set[int] | None = None
     diffusion_kv_cache_skip_step_indices: set[int] | None = None
     diffusion_kv_cache_skip_layer_indices: set[int] | None = None
+    moe_backend: str = "auto"
     force_cutlass_fp8: bool = False
     enable_diffusion_pipeline_profiler: bool = False
     step_execution: bool = False
@@ -551,7 +552,7 @@ class _DiffusionConfigProjection:
             elif isinstance(self.quantization_config, str):
                 self.quantization_config = build_quant_config(self.quantization_config)
             elif isinstance(self.quantization_config, Mapping):
-                self.quantization_config = build_quant_config(dict(self.quantization_config))
+                self.quantization_config = dict(self.quantization_config)
             else:
                 raise TypeError(
                     "quantization_config must be str, dict, QuantizationConfig, or None, "
diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py
index 7bf79d44c..934d15d18 100644
--- a/vllm_omni/config/pipeline_registry.py
+++ b/vllm_omni/config/pipeline_registry.py
@@ -56,6 +56,10 @@ from vllm_omni.model_executor.models.hunyuan_image3.pipeline import (
 )
 from vllm_omni.model_executor.models.indextts2.pipeline import INDEXTTS2_PIPELINE
 from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE
+from vllm_omni.model_executor.models.mammoth_moda2.pipeline import (
+    MAMMOTH_MODA2_AR_PIPELINE,
+    MAMMOTH_MODA2_PIPELINE,
+)
 from vllm_omni.model_executor.models.mimo_audio.pipeline import MIMO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.ming_flash_omni.pipeline import (
     MING_FLASH_OMNI_IMAGE_PIPELINE,
@@ -74,6 +78,7 @@ from vllm_omni.model_executor.models.moss_tts.pipeline import (
     MOSS_TTS_REALTIME_PIPELINE,
 )
 from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE
+from vllm_omni.model_executor.models.omnivoice.pipeline import OMNIVOICE_PIPELINE
 from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
     QWEN2_5_OMNI_PIPELINE,
     QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
@@ -118,6 +123,9 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE,
     "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE,
     "moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
+    "omnivoice": OMNIVOICE_PIPELINE,
+    "mammoth_moda2": MAMMOTH_MODA2_PIPELINE,
+    "mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE,
     "moss_tts_delay": MOSS_TTS_PIPELINE,
     "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
     "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,
diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py
index a269216da..6b8a09423 100644
--- a/vllm_omni/config/stage_config.py
+++ b/vllm_omni/config/stage_config.py
@@ -18,6 +18,7 @@ from transformers import PretrainedConfig
 from vllm.logger import init_logger
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 
+from vllm_omni.config.endpoint_policy import EndpointRestriction
 from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -63,7 +64,12 @@ def build_stage_runtime_overrides(
     if internal_keys is None:
         from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
 
-        internal_keys = internal_blacklist_keys() | SHARED_FIELDS
+        # Some fields are modeled as orchestrator-owned for top-level CLI
+        # parsing, but are also legitimate deploy-time stage overrides. Keep
+        # the default blacklist for true orchestrator/shared fields while
+        # allowing any field explicitly represented by the deploy schema to
+        # continue flowing into per-stage overrides.
+        internal_keys = (internal_blacklist_keys() | SHARED_FIELDS) - deploy_runtime_override_keys()
 
     result: dict[str, Any] = {}
 
@@ -223,6 +229,7 @@ class StagePipelineConfig:
     prompt_expand_func: str | None = None
     cfg_kv_collect_func: str | None = None
     omni_kv_config: dict[str, Any] | None = None
+    scheduler_cls: str | None = None
     # Model subdirectory indirections: for multi-component HF repos where the
     # stage's config/tokenizer lives in a subdirectory (e.g. GLM-Image's AR
     # config is in ``vision_language_encoder/``).  Consumed at stage-init time
@@ -262,6 +269,7 @@ class PipelineConfig:
     # this value to auto-detect the pipeline.  Only needed for diffusers-style
     # multi-component repos (e.g. GLM-Image).  ``None`` = not a diffusers model.
     diffusers_class_name: str | None = None
+    endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
 
     def get_stage(self, stage_id: int) -> StagePipelineConfig | None:
         """Look up a stage by its ID."""
@@ -303,7 +311,7 @@ class StageDeployConfig:
     the top level of ``DeployConfig`` and propagated to every stage.
     """
 
-    # === Omni fields ===
+    # === Omni stage wrapper fields ===
     # Stage identity and Omni runtime placement.
     stage_id: int
     devices: str | None = None
@@ -316,22 +324,33 @@ class StageDeployConfig:
     default_sampling_params: dict[str, Any] | None = None
     subtalker_sampling_params: dict[str, Any] | None = None
 
-    # === vLLM EngineArgs fields ===
-    # Parallelism and scheduler/memory capacity.
+    # === Generic stage engine fields ===
+    # Parallelism, scheduler, and memory-capacity controls.
     tensor_parallel_size: int | None = None
+    enable_expert_parallel: bool | None = None
     gpu_memory_utilization: float | None = None
     max_num_seqs: int | None = None
     max_num_batched_tokens: int | None = None
     max_model_len: int | None = None
 
-    # Execution, scheduling, and KV/cache behavior.
+    # Generic execution, scheduling, and KV/cache behavior.
     enforce_eager: bool | None = None
     async_scheduling: bool | None = None
     disable_hybrid_kv_cache_manager: bool | None = None
     mm_processor_cache_gb: float | None = None
 
-    # Diffusion parallel_config deploy override fields.
-    enable_expert_parallel: bool | None = None
+    # Generic compilation, profiling, tokenizer/config parsing, and model
+    # loading controls.
+    compilation_config: dict[str, Any] | None = None
+    profiler_config: dict[str, Any] | None = None
+    skip_mm_profiling: bool | None = None
+    enable_flashinfer_autotune: bool | None = None
+    config_format: str | None = None
+    load_format: str | None = None
+    tokenizer_mode: str | None = None
+
+    # === Diffusion stage runtime fields ===
+    # Diffusion parallel_config deploy/runtime override fields.
     ulysses_degree: int | None = None
     ulysses_mode: str | None = None
     ring_degree: int | None = None
@@ -343,16 +362,44 @@ class StageDeployConfig:
     hsdp_shard_size: int | None = None
     hsdp_replicate_size: int | None = None
 
-    # Compilation, profiling, tokenizer/config parsing, and model loading.
-    compilation_config: dict[str, Any] | None = None
-    profiler_config: dict[str, Any] | None = None
-    skip_mm_profiling: bool | None = None
-    enable_flashinfer_autotune: bool | None = None
-    config_format: str | None = None
-    load_format: str | None = None
-    tokenizer_mode: str | None = None
+    # Diffusion model loading and adapter construction.
+    model_class_name: str | None = None
+    diffusion_load_format: str | None = None
+    diffusers_load_kwargs: dict[str, Any] | None = None
+    diffusers_call_kwargs: dict[str, Any] | None = None
+    diffusion_quantization_config: str | None = None
+    diffusion_attention_backend: str | None = None
+    diffusion_attention_config: dict[str, Any] | None = None
 
-    # Pass-through vLLM EngineArgs fields that are not represented above.
+    # Diffusion execution, cache, and VAE behavior.
+    cache_backend: str | None = None
+    cache_config: dict[str, Any] | None = None
+    enable_cache_dit_summary: bool | None = None
+    step_execution: bool | None = None
+    vae_use_slicing: bool | None = None
+    vae_use_tiling: bool | None = None
+    boundary_ratio: float | None = None
+    flow_shift: float | None = None
+    diffusion_kv_cache_dtype: str | None = None
+    diffusion_kv_cache_skip_steps: str | None = None
+    diffusion_kv_cache_skip_layers: str | None = None
+    auxiliary_text_encoder: str | None = None
+
+    # Runtime optimizations used by diffusion loading/execution.
+    enable_multithread_weight_load: bool | None = None
+    num_weight_load_threads: int | None = None
+    enable_cpu_offload: bool | None = None
+    enable_layerwise_offload: bool | None = None
+
+    # Diffusion-specific debug and observability knobs.
+    enable_diffusion_pipeline_profiler: bool | None = None
+
+    # Modality/service constraints consumed outside the core engine config.
+    max_generated_image_size: int | None = None
+    tts_max_instructions_length: int | None = None
+
+    # === Pass-through stage engine fields ===
+    # Pass-through stage engine args that are not represented above.
     engine_extras: dict[str, Any] = field(default_factory=dict)
 
 
@@ -409,10 +456,22 @@ _STAGE_RESERVED_KEYS = frozenset(
 _STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_RESERVED_KEYS}
 
 
+def deploy_runtime_override_keys() -> frozenset[str]:
+    """Return deploy-schema fields that are valid CLI/runtime overrides.
+
+    These keys form the positive contract for stage override propagation:
+    stage-scoped deploy knobs plus top-level pipeline-wide engine settings.
+    They must remain overridable even if they are also modeled on
+    ``OrchestratorArgs`` for top-level CLI parsing.
+    """
+    return frozenset(_STAGE_DEPLOY_FIELDS) | frozenset(_PIPELINE_WIDE_ENGINE_FIELDS)
+
+
 def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
     """Parse a single stage entry from deploy YAML into StageDeployConfig."""
     # Get the non-reserved keys for this stage
     flat_args = {k: v for k, v in stage_data.items() if k not in _STAGE_RESERVED_KEYS}
+    explicit_engine_extras = dict(stage_data.get("engine_extras") or {})
     runtime_cfg = dict(stage_data.get("runtime", {}))
     devices = runtime_cfg.get("devices", stage_data.get("devices"))
     num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1))
@@ -440,7 +499,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
     kwargs["output_connectors"] = stage_data.get("output_connectors")
     kwargs["input_connectors"] = stage_data.get("input_connectors")
     kwargs["default_sampling_params"] = stage_data.get("default_sampling_params")
-    kwargs["engine_extras"] = flat_args
+    kwargs["engine_extras"] = _get_recursively_merged_dict(explicit_engine_extras, flat_args)
     return StageDeployConfig(**kwargs)
 
 
@@ -840,7 +899,7 @@ def merge_pipeline_deploy(
                 final_output=ps.final_output,
                 final_output_type=ps.final_output_type,
                 worker_type=worker_type,
-                scheduler_cls=_scheduler_path(sched_cls),
+                scheduler_cls=ps.scheduler_cls or _scheduler_path(sched_cls),
                 hf_config_name=ps.hf_config_name,
                 is_comprehension=ps.owns_tokenizer,
                 yaml_engine_args=engine_args,
diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py
index 178fedee3..74d09ab60 100644
--- a/vllm_omni/core/sched/omni_ar_scheduler.py
+++ b/vllm_omni/core/sched/omni_ar_scheduler.py
@@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStat
 from vllm.logger import init_logger
 from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMScheduler
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs, FinishReason
@@ -133,6 +134,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._omits_kv_transfer_cache[rid] = result
         return result
 
+    def _should_defer_waiting_admission(self) -> bool:
+        return False
+
     def _process_kv_transfer_trigger(self, request: Request, new_token_ids: list[int]) -> bool:
         """
         Check triggers and process side effects (marking transfer).
@@ -222,9 +226,19 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 self.waiting, self.running, scheduler_requests=self.requests
             )
 
+        original_waiting = None
+        if self._should_defer_waiting_admission():
+            original_waiting = self.waiting
+            self.waiting = create_request_queue(self.policy)
+
         try:
             scheduler_output = super().schedule(throttle_prefills)
         finally:
+            if original_waiting is not None:
+                deferred_waiting = list(self.waiting)
+                if deferred_waiting:
+                    original_waiting.prepend_requests(deferred_waiting)
+                self.waiting = original_waiting
             if self.chunk_transfer_adapter:
                 # Add request waiting for chunk to the waiting and running queue
                 self.chunk_transfer_adapter.restore_queues(
@@ -300,13 +314,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if kv_connector_stats and self.connector:
-            kv_stats = self.connector.get_kv_connector_stats()
-            if kv_stats:
-                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
 
         failed_kv_load_req_ids = None
         if kv_connector_output and kv_connector_output.invalid_block_ids:
@@ -549,6 +556,20 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
+        # Worker-side KV connector stats from the model runner output.
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        if self.connector:
+            # Scheduler-side KV connector stats collected after connector update.
+            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
+            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
+                kv_connector_stats = (
+                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
+                    if kv_connector_stats is not None
+                    else scheduler_kv_connector_stats
+                )
+
         # collect KV cache events from KV cache manager
         events = self.kv_cache_manager.take_events()
 
diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py
index 42caea252..973b8c55a 100644
--- a/vllm_omni/core/sched/omni_generation_scheduler.py
+++ b/vllm_omni/core/sched/omni_generation_scheduler.py
@@ -431,14 +431,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        # Merge connector-side stats (align with v0.14.0)
-        if kv_connector_stats and self.connector:
-            kv_stats = self.connector.get_kv_connector_stats()
-            if kv_stats:
-                kv_connector_stats = kv_connector_stats.aggregate(kv_stats)
 
         failed_kv_load_req_ids = None
         if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None):
@@ -646,6 +638,20 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
+        # Worker-side KV connector stats from the model runner output.
+        kv_connector_stats: KVConnectorStats | None = (
+            kv_connector_output.kv_connector_stats if kv_connector_output else None
+        )
+        if self.connector:
+            # Scheduler-side KV connector stats collected after connector update.
+            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
+            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
+                kv_connector_stats = (
+                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
+                    if kv_connector_stats is not None
+                    else scheduler_kv_connector_stats
+                )
+
         # Collect and publish KV cache events (align with v0.14.0)
         events = self.kv_cache_manager.take_events()
         if self.connector is not None:
diff --git a/vllm_omni/deploy/dreamzero.yaml b/vllm_omni/deploy/dreamzero.yaml
index 079101ba0..62b3a5eed 100644
--- a/vllm_omni/deploy/dreamzero.yaml
+++ b/vllm_omni/deploy/dreamzero.yaml
@@ -14,6 +14,8 @@ stages:
     max_num_seqs: 1
     enforce_eager: false
     model_class_name: DreamZeroPipeline
+    # DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache).
+    engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine
     model_config:
       default_robot_embodiment: roboarena
       policy_server_config:
diff --git a/vllm_omni/deploy/mammoth_moda2.yaml b/vllm_omni/deploy/mammoth_moda2.yaml
index d1a49d662..afae38ee0 100644
--- a/vllm_omni/deploy/mammoth_moda2.yaml
+++ b/vllm_omni/deploy/mammoth_moda2.yaml
@@ -1,42 +1,27 @@
-# MammothModa2 AR→DiT (text-to-image). Canonical copy lives at
-# ``vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml``; keep in sync.
-#
-# Must remain legacy ``stage_args`` layout (no top-level ``stages``): Omni resolves
-# explicit ``--stage-configs-path`` via ``load_stage_configs_from_yaml``.
+# MammothModa2 AR→DiT (text-to-image).
 
-stage_args:
+async_chunk: false
+pipeline: mammoth_moda2
+
+stages:
   - stage_id: 0
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: ar
-      max_num_seqs: 100
-      model_arch: MammothModa2ForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 8192
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-    final_output: false
+    devices: "0"
+    max_num_seqs: 100
+    max_model_len: 8192
+    gpu_memory_utilization: 0.5
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
 
   - stage_id: 1
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: dit
-      max_num_seqs: 1
-      model_arch: MammothModa2ForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.3
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: image
-      enable_prefix_caching: false
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit
-    final_output: true
-    final_output_type: image
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.3
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    default_sampling_params:
+      extra_args:
+        text_guidance_scale: 9.0
+        cfg_range: [0.0, 1.0]
+        num_inference_steps: 50
diff --git a/vllm_omni/deploy/mammoth_moda2_ar.yaml b/vllm_omni/deploy/mammoth_moda2_ar.yaml
new file mode 100644
index 000000000..45e2580d9
--- /dev/null
+++ b/vllm_omni/deploy/mammoth_moda2_ar.yaml
@@ -0,0 +1,14 @@
+# MammothModa2 AR-only (image/video understanding and summarization).
+
+async_chunk: false
+pipeline: mammoth_moda2_ar
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 16
+    max_model_len: 8192
+    gpu_memory_utilization: 0.5
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
diff --git a/vllm_omni/deploy/ming_tts.yaml b/vllm_omni/deploy/ming_tts.yaml
index 1a2dc99ec..684a47f98 100644
--- a/vllm_omni/deploy/ming_tts.yaml
+++ b/vllm_omni/deploy/ming_tts.yaml
@@ -15,6 +15,7 @@ connectors:
     name: SharedMemoryConnector
     extra:
       latent_chunk_size: 25
+      initial_latent_chunk_size: 4
       latent_left_context: 0
 
 stages:
diff --git a/vllm_omni/deploy/ming_tts_moe.yaml b/vllm_omni/deploy/ming_tts_moe.yaml
index fe7806b1d..76425712a 100644
--- a/vllm_omni/deploy/ming_tts_moe.yaml
+++ b/vllm_omni/deploy/ming_tts_moe.yaml
@@ -23,6 +23,7 @@ connectors:
     name: SharedMemoryConnector
     extra:
       latent_chunk_size: 25
+      initial_latent_chunk_size: 4
       latent_left_context: 0
 
 stages:
@@ -30,7 +31,9 @@ stages:
     max_num_seqs: 1
     tensor_parallel_size: 1
     gpu_memory_utilization: 0.60
-    enforce_eager: true
+    enforce_eager: false
+    compilation_config:
+      cudagraph_mode: PIECEWISE
     enable_prefix_caching: false
     async_scheduling: false
     max_num_batched_tokens: 8192
diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml
index 7e474caa6..1455ccc09 100644
--- a/vllm_omni/deploy/moss_tts_local.yaml
+++ b/vllm_omni/deploy/moss_tts_local.yaml
@@ -14,8 +14,10 @@ connectors:
       connector_get_sleep_s: 0.005
       connector_get_max_wait_first_chunk: 1000
       connector_get_max_wait: 300
-      codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s chunks
+      initial_codec_chunk_frames: 1   # 1 frame * 80 ms/frame = 80 ms first audio chunk
+      codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s steady chunks
       codec_left_context_frames: 0
+      codec_stream_slots: 8
 
 stages:
   - stage_id: 0
@@ -37,7 +39,7 @@ stages:
       repetition_penalty: 1.1
 
   - stage_id: 1
-    max_num_seqs: 1
+    max_num_seqs: 8
     gpu_memory_utilization: 0.15
     enforce_eager: true
     enable_prefix_caching: false
diff --git a/vllm_omni/deploy/omnivoice.yaml b/vllm_omni/deploy/omnivoice.yaml
index 6491be6dc..6769dd5e9 100644
--- a/vllm_omni/deploy/omnivoice.yaml
+++ b/vllm_omni/deploy/omnivoice.yaml
@@ -1,19 +1,13 @@
 # OmniVoice TTS deploy: diffusion engine single stage (dit).
 # Text → iterative unmasking → codebook tokens → DAC decode → 24kHz audio.
-# (Migrated from model_executor/stage_configs/omnivoice.yaml for deploy/ layout.)
 
-stage_args:
+async_chunk: false
+pipeline: omnivoice
+
+stages:
   - stage_id: 0
-    stage_type: diffusion
-    is_comprehension: true
-    runtime:
-      devices: 0
-    engine_args:
-      model_stage: dit
-      model_class_name: "OmniVoicePipeline"
-      enforce_eager: true
-      trust_remote_code: true
-      distributed_executor_backend: "mp"
-      dtype: "float32"
-    final_output: true
-    final_output_type: audio
+    devices: "0"
+    enforce_eager: true
+    trust_remote_code: true
+    distributed_executor_backend: "mp"
+    dtype: "float32"
diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml
index f91eea5a8..5bcbb0fe3 100644
--- a/vllm_omni/deploy/qwen3_omni_moe.yaml
+++ b/vllm_omni/deploy/qwen3_omni_moe.yaml
@@ -47,7 +47,6 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      seed: 42
       repetition_penalty: 1.05
 
   - stage_id: 2
@@ -67,7 +66,6 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
-      seed: 42
       repetition_penalty: 1.1
 
 platforms:
diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml
index 69085fc5a..f27da1b5b 100644
--- a/vllm_omni/deploy/qwen3_tts.yaml
+++ b/vllm_omni/deploy/qwen3_tts.yaml
@@ -35,7 +35,7 @@ connectors:
 
 stages:
   - stage_id: 0
-    max_num_seqs: 10
+    max_num_seqs: 64
     gpu_memory_utilization: 0.3
     trust_remote_code: true
     # Prefix caching on by default. The HS opt-out (talker only consumes the
@@ -43,9 +43,9 @@ stages:
     # PR make per-step bookkeeping negligible, and any text repetition across
     # requests (same voice or shared ref_audio) is a net throughput win.
     # Disable explicitly if memory headroom is tight on a single-GPU deploy.
-    enable_prefix_caching: true
+    enable_prefix_caching: false
     async_scheduling: true
-    max_num_batched_tokens: 512
+    max_num_batched_tokens: 32768
     max_model_len: 4096
     devices: "0"
     output_connectors:
@@ -54,7 +54,10 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      seed: 42
+      # Official qwen-tts generates the talker with min_new_tokens=2; without
+      # it, checkpoints that sample codec EOS on the first frame return empty
+      # audio (HTTP 400). See #4962.
+      min_tokens: 2
       repetition_penalty: 1.05
     subtalker_sampling_params:
       do_sample: true
@@ -64,7 +67,7 @@ stages:
 
   - stage_id: 1
     # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1.
-    max_num_seqs: 10
+    max_num_seqs: 64
     gpu_memory_utilization: 0.3
     enforce_eager: false
     trust_remote_code: true
@@ -85,7 +88,6 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
-      seed: 42
       repetition_penalty: 1.0
 
 platforms:
@@ -93,3 +95,8 @@ platforms:
     stages:
       - stage_id: 0
         enforce_eager: false
+        compilation_config:
+          cudagraph_mode: PIECEWISE
+      - stage_id: 1
+        compilation_config:
+          cudagraph_mode: PIECEWISE
diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
index 9c6ec19fe..939396840 100644
--- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
+++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml
@@ -60,7 +60,10 @@ stages:
       temperature: 0.9
       top_k: 50
       max_tokens: 4096
-      seed: 42
+      # Official qwen-tts generates the talker with min_new_tokens=2; without
+      # it, checkpoints that sample codec EOS on the first frame return empty
+      # audio (HTTP 400). See #4962.
+      min_tokens: 2
       repetition_penalty: 1.05
     subtalker_sampling_params:
       do_sample: true
@@ -85,7 +88,6 @@ stages:
       top_p: 1.0
       top_k: -1
       max_tokens: 65536
-      seed: 42
       repetition_penalty: 1.0
 
 platforms:
diff --git a/vllm_omni/deploy/voxcpm2.yaml b/vllm_omni/deploy/voxcpm2.yaml
index 5d366521b..63996f1b7 100644
--- a/vllm_omni/deploy/voxcpm2.yaml
+++ b/vllm_omni/deploy/voxcpm2.yaml
@@ -8,16 +8,17 @@ dtype: bfloat16
 
 stages:
   - stage_id: 0
-    max_num_seqs: 32
+    max_num_seqs: 8
     gpu_memory_utilization: 0.9
     # Right-size the KV cache instead of letting it fill gpu_memory_utilization.
-    # max_num_seqs(32) * max_model_len(4096) ~= 131k tokens; 6 GiB holds ~175k
-    # (42x concurrency, above the 32-seq scheduler cap, so no throughput loss).
+    # max_num_seqs(8) * max_model_len(4096) ~= 32k tokens; 6 GiB leaves headroom
+    # for longer live requests while keeping VoxCPM2 TTS decode batches small
+    # enough to avoid the high ITL observed with large admission caps.
     # The remaining VRAM stays free for the VoxCPM2 diffusion side-path
     # (LocDiT / AudioVAE / code2wav + batch CUDA-graph pools), which allocates
     # outside vLLM's KV accounting during inference. Keeps peak ~13 GiB on any
     # card and avoids per-rebase OOM tuning on small CI GPUs.
-    # Assumes prefix caching off and max_num_seqs*max_model_len <= ~175k tokens;
+    # Assumes prefix caching off and typical live decode concurrency <= ~8;
     # raise this cap if either grows or enable_prefix_caching is turned on.
     kv_cache_memory_bytes: 6442450944  # 6 GiB
     enforce_eager: true
@@ -39,6 +40,9 @@ stages:
           enable_loc_dit_fused_mlp: true
           enable_loc_dit_zero_dt_cache: true
           enable_loc_dit_skip_qkv_contig: true
+          enable_loc_dit_reduce_overhead_no_cg: true
+          enable_unified_decode_graph: true
+          unified_decode_graph_max_batch_size: 8
     default_sampling_params:
       temperature: 0.0
       top_p: 1.0
diff --git a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py
new file mode 100644
index 000000000..bce607888
--- /dev/null
+++ b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py
@@ -0,0 +1,406 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import threading
+from functools import partial
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
+    piecewise_attn,
+)
+
+logger = init_logger(__name__)
+
+
+_hub_modules: dict[str, object] = {}
+_hub_lock = threading.Lock()
+
+
+def _load_hub_module(repo_id: str):
+    from kernels import get_kernel
+
+    logger.info("Loading %s kernel from HuggingFace Hub...", repo_id)
+    last_error = None
+    for version in (1, 2, None):
+        try:
+            if version is not None:
+                return get_kernel(repo_id, version=version)
+            return get_kernel(repo_id)
+        except Exception as exc:
+            logger.info("Failed to load %s version %s: %s", repo_id, version, exc)
+            last_error = exc
+    raise RuntimeError(f"Failed to load HuggingFace Hub kernel {repo_id!r}") from last_error
+
+
+def _get_hub_module(repo_id: str):
+    with _hub_lock:
+        if repo_id not in _hub_modules:
+            _hub_modules[repo_id] = _load_hub_module(repo_id)
+        return _hub_modules[repo_id]
+
+
+def _run_varlen_dense(
+    varlen_func,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    *,
+    causal: bool,
+    softmax_scale: float,
+) -> torch.Tensor:
+    batch_size, q_len = query.size()[:2]
+    k_len = key.size(1)
+    cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device)
+    cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, step=k_len, dtype=torch.int32, device=query.device)
+
+    out = varlen_func(
+        q=query.flatten(0, 1),
+        k=key.flatten(0, 1),
+        v=value.flatten(0, 1),
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=q_len,
+        max_seqlen_k=k_len,
+        causal=causal,
+        softmax_scale=softmax_scale,
+    )
+    if isinstance(out, tuple):
+        out = out[0]
+    return out.reshape(batch_size, q_len, *out.shape[1:])
+
+
+class FlashAttentionHubBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN_HUB"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionHubImpl"]:
+        return FlashAttentionHubImpl
+
+
+class FlashAttentionHubImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict | None = None,
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.qkv_layout = qkv_layout
+        if backend_kwargs:
+            logger.warning("FlashAttentionHubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
+
+        hub_module = _get_hub_module("kernels-community/flash-attn2")
+        self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
+        self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
+        if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
+            raise RuntimeError("Failed to load flash-attn2 kernel from HuggingFace Hub: no functions found")
+
+    @staticmethod
+    def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
+        # FA3 may return (out, lse), FA2 returns out
+        return out[0] if isinstance(out, tuple) else out
+
+    def _flash_wrapper(self, q, k, v, *, attn_func, **kwargs):
+        if attn_func is not None:
+            return self._unwrap_flash_output(attn_func(q, k, v, **kwargs))
+        return _run_varlen_dense(
+            self.flash_attn_varlen_func,
+            q,
+            k,
+            v,
+            causal=kwargs.get("causal", self.causal),
+            softmax_scale=kwargs.get("softmax_scale", self.softmax_scale),
+        )
+
+    def _forward_varlen_masked(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            _pad_input,
+            _unpad_input,
+            _upad_input,
+        )
+
+        assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
+        query_length = query.size(1)
+        q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
+            query, key, value, attention_mask, query_length, _unpad_input
+        )
+
+        out_unpad = self.flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seq_lens_q,
+            cu_seqlens_k=cu_seq_lens_k,
+            max_seqlen_q=max_length_q,
+            max_seqlen_k=max_length_k,
+            **{
+                "causal": self.causal,
+                "softmax_scale": self.softmax_scale,
+            },
+        )
+        out_unpad = self._unwrap_flash_output(out_unpad)
+        return _pad_input(out_unpad, indices_q, query.size(0), query_length)
+
+    def _forward_varlen_dense(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        return _run_varlen_dense(
+            self.flash_attn_varlen_func,
+            query,
+            key,
+            value,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
+
+        # Try piecewise attention
+        if full_attn_spans is not None:
+            logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
+            attn_func = partial(
+                self._flash_wrapper,
+                attn_func=self.flash_attn_func,
+            )
+
+            return piecewise_attn(
+                query,
+                key,
+                value,
+                full_attn_spans,
+                self.softmax_scale,
+                attn_func,
+            )
+
+        if attention_mask is not None and torch.any(~attention_mask):
+            return self._forward_varlen_masked(
+                query,
+                key,
+                value,
+                attention_mask,
+            )
+
+        if self.flash_attn_func is not None:
+            out = self.flash_attn_func(
+                query,
+                key,
+                value,
+                causal=self.causal,
+                softmax_scale=self.softmax_scale,
+            )
+            return self._unwrap_flash_output(out)
+
+        return self._forward_varlen_dense(
+            query,
+            key,
+            value,
+        )
+
+
+class FlashAttention3HubBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def supports_attention_mask(cls) -> bool:
+        return True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [64, 96, 128, 192, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN_3_HUB"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttention3HubImpl"]:
+        return FlashAttention3HubImpl
+
+
+class FlashAttention3HubImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        qkv_layout: str | None = None,
+        backend_kwargs: dict | None = None,
+        **extra_impl_args,
+    ) -> None:
+        self.num_heads = num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.qkv_layout = qkv_layout
+        if backend_kwargs:
+            logger.warning("FlashAttention3HubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
+
+        hub_module = _get_hub_module("kernels-community/flash-attn3")
+        self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
+        self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
+        if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
+            raise RuntimeError("Failed to load flash-attn3 kernel from HuggingFace Hub: no functions found")
+
+    @staticmethod
+    def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
+        # FA3 returns (out, lse)
+        return out[0] if isinstance(out, tuple) else out
+
+    def _flash_wrapper(self, q, k, v, *, attn_func, **kwargs):
+        if attn_func is not None:
+            return self._unwrap_flash_output(attn_func(q, k, v, **kwargs))
+        return _run_varlen_dense(
+            self.flash_attn_varlen_func,
+            q,
+            k,
+            v,
+            causal=kwargs.get("causal", self.causal),
+            softmax_scale=kwargs.get("softmax_scale", self.softmax_scale),
+        )
+
+    def _forward_varlen_masked(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            _pad_input,
+            _unpad_input,
+            _upad_input,
+        )
+
+        assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
+        query_length = query.size(1)
+        q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
+            query, key, value, attention_mask, query_length, _unpad_input
+        )
+
+        out_unpad = self.flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seq_lens_q,
+            cu_seqlens_k=cu_seq_lens_k,
+            max_seqlen_q=max_length_q,
+            max_seqlen_k=max_length_k,
+            **{
+                "causal": self.causal,
+                "softmax_scale": self.softmax_scale,
+            },
+        )
+        out_unpad = self._unwrap_flash_output(out_unpad)
+        return _pad_input(out_unpad, indices_q, query.size(0), query_length)
+
+    def _forward_varlen_dense(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        return _run_varlen_dense(
+            self.flash_attn_varlen_func,
+            query,
+            key,
+            value,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+
+    def forward_cuda(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
+
+        # Try piecewise attention
+        if full_attn_spans is not None:
+            logger.debug("Using piecewise Flash Attention for mixed causal/full mask")
+            attn_func = partial(
+                self._flash_wrapper,
+                attn_func=self.flash_attn_func,
+            )
+
+            return piecewise_attn(
+                query,
+                key,
+                value,
+                full_attn_spans,
+                self.softmax_scale,
+                attn_func,
+            )
+
+        if attention_mask is not None and torch.any(~attention_mask):
+            return self._forward_varlen_masked(
+                query,
+                key,
+                value,
+                attention_mask,
+            )
+
+        if self.flash_attn_func is not None:
+            out = self.flash_attn_func(
+                query,
+                key,
+                value,
+                causal=self.causal,
+                softmax_scale=self.softmax_scale,
+            )
+            return self._unwrap_flash_output(out)
+
+        return self._forward_varlen_dense(
+            query,
+            key,
+            value,
+        )
diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py
index 8c22cc2cd..a822d78fb 100644
--- a/vllm_omni/diffusion/attention/backends/registry.py
+++ b/vllm_omni/diffusion/attention/backends/registry.py
@@ -62,6 +62,8 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
     SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend"
     CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend"
     FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend"
+    FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend"
+    FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend"
 
     def get_path(self, include_classname: bool = True) -> str:
         """Get the class path for this backend (respects overrides).
diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py
index e3d07bfcd..5d63ce3c0 100644
--- a/vllm_omni/diffusion/cache/cache_dit_backend.py
+++ b/vllm_omni/diffusion/cache/cache_dit_backend.py
@@ -975,6 +975,38 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheCo
     return enable_cache_for_dit(pipeline, cache_config, block_adapter)
 
 
+def enable_cache_for_krea2(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc:
+    """Enable cache-dit for Krea 2.
+
+    Krea 2 is a single-stream MMDiT: each ``Krea2TransformerBlock`` takes and returns only
+    ``hidden_states`` (text is fused into the token stream), so the blocks follow
+    ``ForwardPattern.Pattern_3``.
+
+    ``has_separate_cfg`` is checkpoint-dependent, which is why this needs a custom enabler
+    rather than a static ``_cache_dit_adapter_config``: the distilled Turbo checkpoint runs
+    no-CFG (a single transformer forward per denoise step), while the Raw checkpoint runs CFG
+    as two separate forwards. cache-dit tells cond/uncond apart purely by transformer-forward
+    parity, so the flag must match the actual per-step forward count — ``True`` only for the
+    CFG (Raw) path. The pipeline exposes this via ``is_distilled`` (read from ``model_index.json``).
+
+    Args:
+        pipeline: The Krea2Pipeline instance.
+        cache_config: DiffusionCacheConfig instance with cache configuration.
+
+    Returns:
+        A refresh function that can be called to update cache context with new num_inference_steps.
+    """
+    transformer = default_get_pipeline_transformer(pipeline)
+    block_adapter = BlockAdapter(
+        transformer=transformer,
+        blocks=[transformer.transformer_blocks],
+        forward_pattern=[ForwardPattern.Pattern_3],
+        has_separate_cfg=not pipeline.is_distilled,
+        check_forward_pattern=False,
+    )
+    return enable_cache_for_dit(pipeline, cache_config, block_adapter)
+
+
 # Register custom cache-dit enablers after function definitions
 CUSTOM_DIT_ENABLERS.update(
     {
@@ -984,6 +1016,8 @@ CUSTOM_DIT_ENABLERS.update(
         "Wan22VACEPipeline": enable_cache_for_wan22,
         "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
         "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
+        "Cosmos3OmniPipeline": enable_cache_for_cosmos3,
+        "Krea2Pipeline": enable_cache_for_krea2,
     }
 )
 
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index f5fe4c30a..e131b833b 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -6,6 +6,7 @@ import os
 import random
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field, fields
+from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 import diffusers
@@ -613,6 +614,20 @@ class OmniDiffusionConfig:
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None
 
+    # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class``
+    # (mirrors ``DiffusionExecutor.get_class``). Config files use a string:
+    # "default" -> DiffusionEngine, or an import path (set e.g. by a deploy
+    # config). Programmatic callers may pass a DiffusionEngine subclass
+    # directly — hence ``str | type`` (structured-config mirrors should expose
+    # the string form only).
+    engine_backend: str | type = "default"
+
+    # Optional override for the diffusion model runner class (import path).
+    # Precedence in the worker: this override > the runner declared by the
+    # selected engine class (``default_diffusion_model_runner_cls``) > the
+    # platform default. Never mutated by engines.
+    diffusion_model_runner_cls: str | None = None
+
     # HuggingFace specific parameters
     trust_remote_code: bool = False
     revision: str | None = None
@@ -649,6 +664,9 @@ class OmniDiffusionConfig:
     # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
     skip_time_steps: int = 15
 
+    # MoE kernel backend selection
+    moe_backend: str = "auto"
+
     # Compilation
     enforce_eager: bool = False
 
@@ -1173,7 +1191,8 @@ class OmniDiffusionConfig:
         valid_fields = {f.name for f in fields(cls)}
         filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
 
-        return cls(**filtered_kwargs)
+        instance = cls(**filtered_kwargs)
+        return instance
 
 
 @dataclass
@@ -1182,8 +1201,10 @@ class DiffusionOutput:
     Final output (after pipeline completion)
     """
 
-    # Fields may be replaced with SHM handle dicts by ipc.pack_diffusion_output_shm
     output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None
+
+    # Legacy compatibility fields. New pipeline-specific payloads should be
+    # carried by output["payload"] instead.
     trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None
     trajectory_latents: torch.Tensor | dict[str, Any] | None = None
     trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None
@@ -1196,10 +1217,6 @@ class DiffusionOutput:
 
     post_process_func: Callable[..., Any] | None = None
 
-    # Extra custom output data (e.g. latent trajectories, prompt embeds)
-    # passed through to OmniRequestOutput.custom_output
-    custom_output: dict[str, Any] = field(default_factory=dict)
-
     # logged timings info, directly from Req.timings
     # timings: Optional["RequestTimings"] = None
 
@@ -1214,10 +1231,9 @@ class DiffusionOutput:
     # memory usage info
     peak_memory_mb: float = 0.0
 
-    # When True, move all tensor fields (including tensors inside
-    # ``custom_output``) to CPU at construction time. Useful when the output
-    # is shipped across process boundaries (e.g. step-execution mode) and the
-    # receiving side must not initialise a stray CUDA context.
+    # When True, move tensor fields to CPU at construction time. Useful when
+    # the output is shipped across process boundaries (e.g. step-execution
+    # mode) and the receiving side must not initialise a stray CUDA context.
     to_cpu: bool = False
 
     def __post_init__(self) -> None:
@@ -1227,14 +1243,18 @@ class DiffusionOutput:
         def _maybe_to_cpu(value: Any) -> Any:
             if isinstance(value, torch.Tensor):
                 return value.detach().cpu()
+            if isinstance(value, dict):
+                return {key: _maybe_to_cpu(item) for key, item in value.items()}
+            if isinstance(value, list):
+                return [_maybe_to_cpu(item) for item in value]
+            if isinstance(value, tuple):
+                return tuple(_maybe_to_cpu(item) for item in value)
             return value
 
         self.output = _maybe_to_cpu(self.output)
         self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps)
         self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents)
         self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs)
-        if self.custom_output:
-            self.custom_output = {k: _maybe_to_cpu(v) for k, v in self.custom_output.items()}
 
     @classmethod
     def from_exception(cls, exc: BaseException) -> "DiffusionOutput":
@@ -1473,5 +1493,12 @@ class OmniWakeTask:
     tags: list[str] | None = None
 
 
+class CuMemTag(str, Enum):
+    """Tags representing specific CuMem allocations for sleep/wake state tracking."""
+
+    WEIGHTS = "weights"
+    KV_CACHE = "kv_cache"
+
+
 # Special message broadcast via scheduler queues to signal worker shutdown.
 SHUTDOWN_MESSAGE = {"type": "shutdown"}
diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index 2554099e7..d6e7b35de 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/vllm_omni/diffusion/diffusion_engine.py
@@ -10,7 +10,7 @@ import queue
 import threading
 import time
 from collections.abc import AsyncGenerator, Iterable
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -40,7 +40,6 @@ from vllm_omni.diffusion.output_formatter import (
 )
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
-    get_diffusion_action_post_process_func,
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
 )
@@ -134,6 +133,14 @@ class _RpcTask:
 class DiffusionEngine:
     """The diffusion engine for vLLM-Omni diffusion models."""
 
+    #: Import path of the model runner this engine's workers should build, or
+    #: ``None`` for the platform default. Subclasses declare their runner here
+    #: (e.g. the AR-Diffusion engine), the worker resolves it through
+    #: :meth:`resolve_engine_class` — so engine->runner routing lives on the
+    #: engine class itself, and ``od_config.diffusion_model_runner_cls``
+    #: remains a pure explicit user override (never mutated by engines).
+    default_diffusion_model_runner_cls: str | None = None
+
     def __init__(
         self,
         od_config: OmniDiffusionConfig,
@@ -147,17 +154,10 @@ class DiffusionEngine:
         self.od_config = od_config
 
         self.post_process_func = get_diffusion_post_process_func(od_config)
-        self.action_post_process_func = get_diffusion_action_post_process_func(od_config)
         self.pre_process_func = get_diffusion_pre_process_func(od_config)
         # Cache whether the model-specific postprocess accepts request-level
         # sampling params so step() can support both legacy and extended hooks.
         self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params")
-        self._action_post_process_accepts_sampling_params = _func_accepts_parameter(
-            self.action_post_process_func, "sampling_params"
-        )
-        self._action_post_process_accepts_custom_output = _func_accepts_parameter(
-            self.action_post_process_func, "custom_output"
-        )
 
         self.step_execution = bool(getattr(od_config, "step_execution", False))
         if self.od_config.streaming_output and not self.step_execution:
@@ -299,46 +299,18 @@ class DiffusionEngine:
         if self.od_config.enable_cpu_offload:
             output_data = _move_tensor_tree_to_cpu(output_data)
 
-        custom_output = output.custom_output or {}
-        action_payload = None
-        action_only_output = bool(custom_output.get("action_only_output"))
-
         postprocess_start_time = time.perf_counter()
-        if action_only_output:
-            outputs = []
-        elif self.post_process_func is not None:
+        if self.post_process_func is not None:
             # Some video pipelines need request-level controls during
             # postprocess (for example worker-side frame interpolation).
+            postprocess_kwargs: dict[str, object] = {}
             if self._post_process_accepts_sampling_params:
-                outputs = self.post_process_func(output_data, sampling_params=request.sampling_params)
-            else:
-                outputs = self.post_process_func(output_data)
+                postprocess_kwargs["sampling_params"] = request.sampling_params
+            outputs = self.post_process_func(output_data, **postprocess_kwargs)
         else:
             outputs = output_data
 
-        postprocess_output = normalize_diffusion_postprocess_output(outputs, custom_output)
-        custom_output = postprocess_output.custom_output
-        action_payload = postprocess_output.action_payload
-        if action_payload is None:
-            action_payload = custom_output.get("actions")
-            if action_payload is not None:
-                postprocess_output = replace(postprocess_output, action_payload=action_payload)
-        action_post_process_func = getattr(self, "action_post_process_func", None)
-        if action_payload is None and action_post_process_func is not None:
-            raw_action_payload = custom_output.get("action")
-            if raw_action_payload is not None:
-                action_kwargs: dict[str, Any] = {}
-                if getattr(self, "_action_post_process_accepts_custom_output", False):
-                    action_kwargs["custom_output"] = custom_output
-                if getattr(self, "_action_post_process_accepts_sampling_params", False):
-                    action_kwargs["sampling_params"] = request.sampling_params
-                action_payload = action_post_process_func(raw_action_payload, **action_kwargs)
-                custom_output = {**custom_output, "actions": action_payload}
-                postprocess_output = replace(
-                    postprocess_output,
-                    custom_output=custom_output,
-                    action_payload=action_payload,
-                )
+        postprocess_output = normalize_diffusion_postprocess_output(outputs)
         postprocess_time = time.perf_counter() - postprocess_start_time
         logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
 
@@ -609,19 +581,55 @@ class DiffusionEngine:
             self._put_streaming_output_with_cv(request_id, out)
 
     @staticmethod
-    def make_engine(
-        config: OmniDiffusionConfig,
-        scheduler: SchedulerInterface | None = None,
-    ) -> DiffusionEngine:
-        """Factory method to create a DiffusionEngine instance.
+    def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
+        """Resolve the engine class selected by ``config.engine_backend``.
+
+        Mirrors ``DiffusionExecutor.get_class``: accepts ``"default"``, a
+        ``DiffusionEngine`` subclass, or an import-path string (e.g. a deploy
+        config's ``engine_backend``). Kept separate from :meth:`make_engine` so the
+        selection is testable without constructing an engine (which runs a dummy forward).
 
         Args:
             config: The configuration for the diffusion engine.
 
         Returns:
-            An instance of DiffusionEngine.
+            The ``DiffusionEngine`` (sub)class to instantiate.
         """
-        return DiffusionEngine(config, scheduler=scheduler)
+        backend = getattr(config, "engine_backend", "default") or "default"
+
+        if isinstance(backend, type):
+            if not issubclass(backend, DiffusionEngine):
+                raise TypeError(f"engine_backend must be a DiffusionEngine subclass. Got {backend}.")
+            return backend
+        if backend == "default":
+            return DiffusionEngine
+        if isinstance(backend, str):
+            try:
+                engine_class = resolve_obj_by_qualname(backend)
+            except (ImportError, ValueError) as e:
+                raise ValueError(
+                    f"Failed to load engine_backend '{backend}'. Ensure it is a valid python path. Error: {e}"
+                ) from e
+            if not issubclass(engine_class, DiffusionEngine):
+                raise TypeError(f"engine_backend must resolve to a DiffusionEngine subclass. Got {engine_class}.")
+            return engine_class
+        raise ValueError(f"Unknown engine_backend: {backend!r}")
+
+    @staticmethod
+    def make_engine(
+        config: OmniDiffusionConfig,
+        scheduler: SchedulerInterface | None = None,
+    ) -> DiffusionEngine:
+        """Factory method to create the engine selected by ``config.engine_backend``.
+
+        Args:
+            config: The configuration for the diffusion engine.
+
+        Returns:
+            An instance of the resolved ``DiffusionEngine`` (sub)class.
+        """
+        engine_class = DiffusionEngine.resolve_engine_class(config)
+        return engine_class(config, scheduler=scheduler)
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         with self._cv:
diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py
index 5c5b092ef..0a67fd941 100644
--- a/vllm_omni/diffusion/distributed/group_coordinator.py
+++ b/vllm_omni/diffusion/distributed/group_coordinator.py
@@ -589,11 +589,7 @@ class GroupCoordinator:
         if dst is None:
             dst = self.group_next_rank
 
-        torch.distributed.send(
-            tensor,
-            self.ranks[dst],
-            group=(self.device_groups[self.rank_in_group % 2] if self.world_size == 2 else self.device_group),
-        )
+        torch.distributed.send(tensor, self.ranks[dst], group=self.device_group)
 
     def recv(self, size: torch.Size, dtype: torch.dtype, src: int | None = None) -> torch.Tensor:
         """Receives a tensor from the src rank."""
@@ -602,11 +598,7 @@ class GroupCoordinator:
             src = self.group_prev_rank
 
         tensor = torch.empty(size, dtype=dtype, device=self.device)
-        torch.distributed.recv(
-            tensor,
-            self.ranks[src],
-            (self.device_groups[(self.rank_in_group + 1) % 2] if self.world_size == 2 else self.device_group),
-        )
+        torch.distributed.recv(tensor, self.ranks[src], self.device_group)
         return tensor
 
     def destroy(self):
@@ -710,6 +702,22 @@ class PipelineGroupCoordinator(GroupCoordinator):
                 self.skip_device_group = skip_device_group
         assert self.skip_device_group is not None
 
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+        if dst is None:
+            dst = self.group_next_rank
+
+        group = self.device_groups[self.rank_in_group % 2] if self.world_size == 2 else self.device_group
+        torch.distributed.send(tensor, self.ranks[dst], group=group)
+
+    def recv(self, size: torch.Size, dtype: torch.dtype, src: int | None = None) -> torch.Tensor:
+        if src is None:
+            src = self.group_prev_rank
+
+        tensor = torch.empty(size, dtype=dtype, device=self.device)
+        group = self.device_groups[(self.rank_in_group + 1) % 2] if self.world_size == 2 else self.device_group
+        torch.distributed.recv(tensor, self.ranks[src], group)
+        return tensor
+
     def reset_buffer(self):
         self.recv_tasks_queue = []
         self.receiving_tasks = []
diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py
index 3a056d6d4..cae4356f4 100644
--- a/vllm_omni/diffusion/distributed/hsdp.py
+++ b/vllm_omni/diffusion/distributed/hsdp.py
@@ -26,6 +26,15 @@ from vllm_omni.platforms import current_omni_platform
 logger = init_logger(__name__)
 
 
+def _unshardable_parameters(model: nn.Module) -> set[nn.Parameter]:
+    """Return packed/integer or scalar parameters unsupported by FSDP2."""
+    return {
+        param
+        for param in model.parameters()
+        if param.ndim == 0 or not (param.is_floating_point() or param.is_complex())
+    }
+
+
 @dataclass
 class HSDPInferenceConfig:
     """Configuration for HSDP inference.
@@ -202,6 +211,32 @@ def apply_hsdp_to_model(
             target_device,
         )
 
+    # Serialized low-bit checkpoints may store packed weights in integer
+    # Parameters (for example, ModelOpt NVFP4 uses uint8) and global scales as
+    # scalar Parameters. FSDP2 cannot represent non-floating or zero-dimensional
+    # sharded parameters. Keep those tensors resident and replicated; eligible
+    # scales, biases, and other parameters remain HSDP-sharded.
+    unshardable_params = _unshardable_parameters(model)
+    if unshardable_params:
+        if target_device is None:
+            raise ValueError(
+                f"Model {type(model).__name__} has parameters that HSDP must ignore, "
+                "but apply_hsdp_to_model was called without target_device."
+            )
+        for param in unshardable_params:
+            if param.device != target_device:
+                param.data = param.data.to(target_device)
+        ignored_params.update(unshardable_params)
+        logger.info(
+            "HSDP excluding %d unshardable parameter tensors from sharding "
+            "(non_floating=%d, scalar=%d, dtypes=%s, moved to %s)",
+            len(unshardable_params),
+            sum(not (param.is_floating_point() or param.is_complex()) for param in unshardable_params),
+            sum(param.ndim == 0 for param in unshardable_params),
+            sorted({str(param.dtype) for param in unshardable_params}),
+            target_device,
+        )
+
     # Apply HSDP sharding, this will automatically handle weight distribution
     shard_model(
         model,
@@ -233,9 +268,9 @@ def shard_model(
     ignored_params (if provided) are excluded from the root fully_shard
     wrap, so they are not collected into the root flat-parameter, are not
     subject to MixedPrecisionPolicy, and retain their original dtype.
-    Per-submodule shard wraps do not receive ignored_params because the
-    ignored modules are expected to live at the root level, not inside any
-    block matched by hsdp_shard_conditions.
+    Each per-submodule wrap receives the subset of ignored_params that it owns.
+    This is required for packed integer parameters inside sharded transformer
+    blocks; the root wrap receives the full set for all remaining parameters.
     """
     hsdp_kwargs: dict[str, Any] = {
         "reshard_after_forward": reshard_after_forward,
@@ -246,7 +281,12 @@ def shard_model(
     num_sharded = 0
     for name, module in reversed(list(model.named_modules())):
         if any(cond(name, module) for cond in hsdp_shard_conditions):
-            fully_shard(module, **hsdp_kwargs)
+            module_kwargs = dict(hsdp_kwargs)
+            if ignored_params:
+                module_ignored_params = ignored_params.intersection(module.parameters())
+                if module_ignored_params:
+                    module_kwargs["ignored_params"] = module_ignored_params
+            fully_shard(module, **module_kwargs)
             num_sharded += 1
 
     if num_sharded == 0:
diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py
index 5904ae94f..7b38d3775 100644
--- a/vllm_omni/diffusion/distributed/parallel_state.py
+++ b/vllm_omni/diffusion/distributed/parallel_state.py
@@ -512,6 +512,21 @@ def init_model_parallel_group(
         )
 
 
+def init_vllm_model_parallel_group(
+    group_ranks: list[list[int]],
+    local_rank: int,
+    backend: str,
+    group_name: str,
+) -> vllm_parallel_state.GroupCoordinator:
+    return vllm_parallel_state.init_model_parallel_group(
+        group_ranks=group_ranks,
+        local_rank=local_rank,
+        backend=backend,
+        group_name=group_name,
+        use_device_communicator=True,
+    )
+
+
 def init_dit_group(
     dit_parallel_size: int,
     backend: str,
@@ -838,23 +853,40 @@ def initialize_model_parallel(
     if use_moe_parallel_mapping:
         # Diffusion normally uses its own SP group. Map it to vLLM PCP only for
         # expert-parallel runtimes that rely on vLLM FusedMoE group semantics.
-        vllm_parallel_state._PCP = _SP
+        # vLLM 0.24 MoE kernels require GroupCoordinator.device_communicator
+        # and reduce_scatter(), which the diffusion SP coordinator intentionally
+        # does not own. Keep the rank layout but build a vLLM coordinator.
+        vllm_parallel_state._PCP = init_vllm_model_parallel_group(
+            group_ranks=sp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            group_name="pcp",
+        )
 
     assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized"
-    vllm_parallel_state._TP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("tp"),
-        local_rank=get_world_group().local_rank,
-        backend=backend,
-        parallel_mode="tensor",
-    )
-    if use_moe_parallel_mapping and cfg_parallel_size > 1:
+    tp_group_ranks = rank_generator.get_ranks("tp")
+    if use_moe_parallel_mapping:
+        vllm_parallel_state._TP = init_vllm_model_parallel_group(
+            group_ranks=tp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            group_name="tp",
+        )
+    else:
+        vllm_parallel_state._TP = init_model_parallel_group(
+            group_ranks=tp_group_ranks,
+            local_rank=get_world_group().local_rank,
+            backend=backend,
+            parallel_mode="tensor",
+        )
+    if use_moe_parallel_mapping:
         # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP
         # only when constructing the vLLM EP layout for expert-parallel paths.
-        vllm_parallel_state._DP = init_model_parallel_group(
+        vllm_parallel_state._DP = init_vllm_model_parallel_group(
             group_ranks=rank_generator.get_ranks("cfg-dp"),
             local_rank=get_world_group().local_rank,
             backend=backend,
-            parallel_mode="data",
+            group_name="dp",
         )
 
     global _FS, _EXPERT_PARALLEL_GROUP_RANKS
@@ -869,11 +901,11 @@ def initialize_model_parallel(
     _EXPERT_PARALLEL_GROUP_RANKS = None
     if use_moe_parallel_mapping:
         ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp")
-        vllm_parallel_state._EP = init_model_parallel_group(
+        vllm_parallel_state._EP = init_vllm_model_parallel_group(
             group_ranks=ep_group_ranks,
             local_rank=get_world_group().local_rank,
             backend=backend,
-            parallel_mode="expert",
+            group_name="ep",
         )
         _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks
 
diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py
index 4cb714465..6a045fca2 100644
--- a/vllm_omni/diffusion/layers/rope.py
+++ b/vllm_omni/diffusion/layers/rope.py
@@ -188,6 +188,12 @@ class RotaryEmbedding(CustomOp):
         cos: torch.Tensor,
         sin: torch.Tensor,
     ) -> torch.Tensor:
+        # All batch elements share the same rotary position encoding.
+        # Strip the batch dim so the underlying op broadcasts over the batch,
+        # consistent with forward_cuda / forward_hip / apply_rotary_emb_mindiesd.
+        if cos.dim() == 3:
+            cos = cos[0]
+            sin = sin[0]
         return apply_rotary_emb_torch(
             x,
             cos,
diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py
index 04fc8ea32..68ec98d5a 100644
--- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py
+++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py
@@ -11,7 +11,12 @@ from vllm.model_executor.utils import get_packed_modules_mapping
 
 logger = init_logger(__name__)
 
-MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_2", ".weight_scale_inv")
+MODEL_OPT_SCALE_SUFFIXES = (
+    ".input_scale",
+    ".weight_scale",
+    ".weight_scale_2",
+    ".weight_scale_inv",
+)
 DEFAULT_PACKED_MODULES_MAPPING = {
     "to_qkv": ("to_q", "to_k", "to_v"),
     "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"),
@@ -32,7 +37,7 @@ FP8_DTYPES = tuple(
 @dataclass
 class _AdaptState:
     scale_tensors: dict[str, torch.Tensor] = field(default_factory=dict)
-    pending_weights: dict[str, list[tuple[str, torch.Tensor, torch.dtype]]] = field(default_factory=dict)
+    pending_weights: dict[str, list[tuple[str, str, torch.Tensor, torch.dtype]]] = field(default_factory=dict)
     skipped_scales: int = 0
     dequantized_weights: int = 0
 
@@ -41,6 +46,7 @@ class ModelOptFp8CheckpointAdapter:
     def __init__(self, model: nn.Module, source: object):
         self._loadable_tensors = self._get_model_loadable_tensors(model)
         self._weights_mapper = self._get_weights_mapper(model)
+        self._checkpoint_key_mapper = getattr(model, "remap_checkpoint_key", None)
         self._source_label = getattr(source, "prefix", "") or getattr(source, "subfolder", None) or "model"
 
     @classmethod
@@ -110,14 +116,20 @@ class ModelOptFp8CheckpointAdapter:
             orig_to_new_prefix=orig_to_new_prefix,
         )
 
-    def _resolve_target_name(self, name: str) -> str | None:
+    def _resolve_target_and_output_names(self, name: str) -> tuple[str | None, str]:
         if name in self._loadable_tensors:
-            return name
+            return name, name
+
+        if callable(self._checkpoint_key_mapper):
+            candidate = self._checkpoint_key_mapper(name)
+            if candidate in self._loadable_tensors:
+                return candidate, candidate
 
         for candidate in self._weights_mapper.apply_list([name]):
             if candidate != name and candidate in self._loadable_tensors:
-                return candidate
-        return None
+                # Preserve shard names so downstream packed-weight loaders can route them.
+                return candidate, name
+        return None, name
 
     @staticmethod
     def _reshape_weight_scale(scale: torch.Tensor, weight_shape: torch.Size) -> torch.Tensor:
@@ -160,13 +172,19 @@ class ModelOptFp8CheckpointAdapter:
         scale_name: str,
         state: _AdaptState,
     ) -> Generator[tuple[str, torch.Tensor], None, None]:
-        for weight_name, weight_tensor, target_dtype in state.pending_weights.pop(scale_name, []):
-            yield weight_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype)
+        for (
+            weight_name,
+            output_name,
+            weight_tensor,
+            target_dtype,
+        ) in state.pending_weights.pop(scale_name, []):
+            yield output_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype)
             state.dequantized_weights += 1
 
     def _handle_scale_tensor(
         self,
         name: str,
+        output_name: str,
         tensor: torch.Tensor,
         target_name: str | None,
         state: _AdaptState,
@@ -175,7 +193,7 @@ class ModelOptFp8CheckpointAdapter:
         if target_name is None:
             state.skipped_scales += 1
         else:
-            yield name, tensor
+            yield output_name, tensor
         yield from self._flush_pending_weights(name, state)
 
     def _target_dtype_for_dequantization(
@@ -194,6 +212,7 @@ class ModelOptFp8CheckpointAdapter:
     def _maybe_dequantize_or_defer_weight(
         self,
         name: str,
+        output_name: str,
         tensor: torch.Tensor,
         target_dtype: torch.dtype,
         state: _AdaptState,
@@ -203,7 +222,7 @@ class ModelOptFp8CheckpointAdapter:
             raise ValueError(f"Missing ModelOpt FP8 weight_scale name for weight {name!r}")
 
         if scale_name not in state.scale_tensors:
-            state.pending_weights.setdefault(scale_name, []).append((name, tensor, target_dtype))
+            state.pending_weights.setdefault(scale_name, []).append((name, output_name, tensor, target_dtype))
             return None
 
         state.dequantized_weights += 1
@@ -235,23 +254,25 @@ class ModelOptFp8CheckpointAdapter:
         state = _AdaptState()
 
         for name, tensor in weights:
-            target_name = self._resolve_target_name(name)
+            target_name, output_name = self._resolve_target_and_output_names(name)
             if self._is_scale(name):
-                yield from self._handle_scale_tensor(name, tensor, target_name, state)
+                yield from self._handle_scale_tensor(name, output_name, tensor, target_name, state)
                 continue
 
             target_dtype = self._target_dtype_for_dequantization(tensor, target_name)
             if target_dtype is not None:
-                tensor = self._maybe_dequantize_or_defer_weight(name, tensor, target_dtype, state)
+                tensor = self._maybe_dequantize_or_defer_weight(name, output_name, tensor, target_dtype, state)
                 if tensor is None:
                     continue
-            yield name, tensor
+            yield output_name, tensor
 
         self._check_pending_weights(state)
         self._log_adaptation_summary(state)
 
 
 class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
+    _PRE_QUANT_SCALE_SUFFIX = ".pre_quant_scale"
+
     @staticmethod
     def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
         return (
@@ -261,6 +282,22 @@ class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
             and bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False))
         )
 
+    def adapt(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        def validated_weights() -> Generator[tuple[str, torch.Tensor], None, None]:
+            for name, tensor in weights:
+                if name.endswith(self._PRE_QUANT_SCALE_SUFFIX):
+                    raise ValueError(
+                        f"ModelOpt NVFP4 checkpoint tensor {name!r} is unsupported: "
+                        "vLLM 0.25.0 does not consume pre_quant_scale. Export the checkpoint "
+                        "with pre-quant scales folded into the weights."
+                    )
+                yield name, tensor
+
+        yield from super().adapt(validated_weights())
+
 
 class ModelOptMixedPrecisionCheckpointAdapter(ModelOptFp8CheckpointAdapter):
     @staticmethod
diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py
index 7ee2b4395..8ee172e6b 100644
--- a/vllm_omni/diffusion/model_loader/diffusers_loader.py
+++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py
@@ -11,7 +11,6 @@ from pathlib import Path
 from typing import cast
 
 import torch
-from huggingface_hub import hf_hub_download
 from torch import nn
 from vllm.config import ModelConfig
 from vllm.config.load import LoadConfig
@@ -36,7 +35,6 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
 )
-from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline
 from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 from vllm_omni.diffusion.registry import initialize_model
@@ -381,15 +379,11 @@ class DiffusersPipelineLoader:
                     # DiffusionPipeline.from_pretrained() internally; it does
                     # not use our native customized pipeline classes.
                     cast(DiffusersAdapterPipeline, model).load_weights()
-                elif self._is_gguf_quantization():
-                    self._load_weights_with_gguf(model)
                 else:
-                    # Quantization does not happen in `load_weights` but after it
                     self.load_weights(model)
-
-            # Process weights after loading for quantization (e.g., FP8 online quantization)
-            # This is needed for vLLM's quantization methods that need to transform weights
-            self._process_weights_after_loading(model, target_device)
+                # HSDP processes quantized weights before wrapping parameters as
+                # DTensors. The non-HSDP path can process them here as usual.
+                self._process_weights_after_loading(model, target_device)
 
             if offload_after_quant:
                 model.to("cpu")
@@ -397,16 +391,77 @@ class DiffusersPipelineLoader:
 
         return model.eval()
 
+    @staticmethod
+    def _has_online_quant(model: nn.Module) -> bool:
+        """Whether any layer uses an online-quant method that defers weight
+        materialization onto the ``meta`` device (upstream vLLM
+        ``uses_meta_device=True``, e.g. online FP8)."""
+        for module in model.modules():
+            quant_method = getattr(module, "quant_method", None)
+            if getattr(quant_method, "uses_meta_device", False):
+                return True
+        return False
+
     def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None:
         """Process weights after loading for quantization methods.
 
         This handles vLLM's quantization methods that need to process weights
         after loading (e.g., FP8 online quantization from BF16/FP16 weights).
         """
+        # Newer upstream vLLM online-quant methods (uses_meta_device=True) create
+        # weights on the ``meta`` device and materialize them just-in-time as each
+        # layer's weights finish loading (via the layerwise online-process loader).
+        # Any "straggler" layers whose weights were not fully materialized during
+        # load (padded / partially-loaded layers) remain on ``meta``. Upstream's
+        # base_loader calls finalize_layerwise_processing() to materialize them;
+        # the diffusion loader must mirror that, otherwise the module.to() below
+        # raises "Cannot copy out of meta tensor; no data!". This whole meta-device
+        # handling is gated on online quant actually being in use, so that the
+        # proven code path for everything else (in particular FSDP/HSDP-sharded
+        # params, whose per-parameter .data cannot be cross-device reassigned) is
+        # left untouched. Import lazily so older vLLM (no meta-device quant) is
+        # unaffected.
+        has_online_quant = self._has_online_quant(model)
+        if has_online_quant:
+            from vllm.model_executor.model_loader.reload.layerwise import (
+                finalize_layerwise_processing,
+            )
+
+            # model_config is only dereferenced by finalize for vLLM Attention /
+            # MLAAttention layers; diffusion DiT models use their own attention and
+            # have none, so passing None is safe here.
+            finalize_layerwise_processing(model, model_config=None)
+
         for _, module in model.named_modules():
             quant_method = getattr(module, "quant_method", None)
-            if quant_method is not None and isinstance(quant_method, QuantizeMethodBase):
-                # Move module to target device for processing if needed
+            if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
+                continue
+
+            if has_online_quant:
+                # Online quant may leave straggler params on the ``meta`` device.
+                # Move only real (non-meta) params onto the target device for
+                # processing and restore them afterward, mirroring upstream vLLM's
+                # device_loading_context — a blanket module.to(target_device) would
+                # raise NotImplementedError on meta params. Online quant initializes
+                # on the accelerator, so params are normally already on the target
+                # device and this loop is a no-op move; the point is to skip meta.
+                original_devices: dict[str, torch.device] = {}
+                for name, param in module.named_parameters():
+                    if param.device.type != "meta" and param.device != target_device:
+                        original_devices[name] = param.device
+                        param.data = param.data.to(target_device)
+
+                quant_method.process_weights_after_loading(module)
+
+                # Restore pre-existing params to their original device; leave any
+                # newly created (e.g. quantized) params on the target device.
+                for name, param in module.named_parameters():
+                    if name in original_devices:
+                        param.data = param.data.to(original_devices[name])
+            else:
+                # No meta params possible here. Preserve the original FSDP/HSDP-aware
+                # whole-module move (module.to()), which correctly handles sharded
+                # DTensor params that per-parameter .data reassignment cannot.
                 module_device = next(module.parameters(), None)
                 if module_device is not None:
                     module_device = module_device.device
@@ -454,7 +509,7 @@ class DiffusersPipelineLoader:
     def _is_expected_quantized_weight(name: str) -> bool:
         """Return True if *name* is a quantization-specific parameter.
 
-        Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra
+        Quantization methods (GPTQ, AWQ, FP8, Autoround, etc.) create extra
         parameters that have no counterpart in an unquantized checkpoint.
         These are expected to be absent and should not trigger a load error.
         """
@@ -467,8 +522,6 @@ class DiffusersPipelineLoader:
             ".weight_scale",
             ".weight_scale_inv",
             ".input_scale",
-            # GGUF
-            ".qweight_type",
             # INT8  (weight_scale already covered above)
         )
         return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES)
@@ -500,116 +553,6 @@ class DiffusersPipelineLoader:
         if unexpected_missing:
             raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}")
 
-    def _get_gguf_model_from_config(self) -> str | None:
-        """Extract gguf_model from quant_config (handles both dict and object)."""
-        if self.quant_config is None:
-            return None
-        if isinstance(self.quant_config, dict):
-            return self.quant_config.get("gguf_model")
-        return getattr(self.quant_config, "gguf_model", None)
-
-    def _is_gguf_quantization(self) -> bool:
-        """Check whether or not this pipeline loader is pointing at a GGUF config."""
-        from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig
-
-        maybe_gguf_model = self._get_gguf_model_from_config()
-        is_gguf = (
-            isinstance(self.quant_config, DiffusionGGUFConfig)
-            or (isinstance(self.quant_config, dict) and self.quant_config.get("method") == "gguf")
-            or (hasattr(self.quant_config, "get_name") and self.quant_config.get_name() == "gguf")
-        )
-        if is_gguf:
-            if not maybe_gguf_model:
-                raise ValueError("GGUF quantization requires gguf_model")
-            return True
-
-        if isinstance(self.quant_config, dict):
-            return False
-
-        return bool(maybe_gguf_model)
-
-    def _is_transformer_source(self, source: "ComponentSource") -> bool:
-        if source.subfolder == "transformer":
-            return True
-        return source.prefix.startswith("transformer.")
-
-    def _get_model_loadable_names(self, model: nn.Module) -> set[str]:
-        # Avoid model.state_dict() here because GGUF uses UninitializedParameter
-        # which raises during detach(). Collect names directly.
-        return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()}
-
-    def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str:
-        if os.path.isfile(gguf_model):
-            return gguf_model
-        # repo_id/filename.gguf
-        if "/" in gguf_model and gguf_model.endswith(".gguf"):
-            repo_id, filename = gguf_model.rsplit("/", 1)
-            return hf_hub_download(
-                repo_id=repo_id,
-                filename=filename,
-                revision=revision,
-                cache_dir=self.load_config.download_dir,
-            )
-        # repo_id:quant_type
-        if "/" in gguf_model and ":" in gguf_model:
-            repo_id, quant_type = gguf_model.rsplit(":", 1)
-            return download_gguf(
-                repo_id,
-                quant_type,
-                cache_dir=self.load_config.download_dir,
-                revision=revision,
-                ignore_patterns=self.load_config.ignore_patterns,
-            )
-        raise ValueError(
-            f"Unrecognized GGUF reference: {gguf_model!r} (expected local file, "
-            "<repo_id>/<filename>.gguf, or <repo_id>:<quant_type>)"
-        )
-
-    def _get_gguf_weights_iterator(
-        self,
-        source: "ComponentSource",
-        model: nn.Module,
-    ) -> Generator[tuple[str, torch.Tensor], None, None]:
-        gguf_model = self._get_gguf_model_from_config()
-        if gguf_model is None:
-            raise ValueError("GGUF quantization requires quantization_config.gguf_model")
-        gguf_file = self._resolve_gguf_model_path(gguf_model, self.od_config.revision)
-        adapter = get_gguf_adapter(gguf_file, model, source, self.od_config)
-        weights_iter = adapter.weights_iterator()
-        return ((source.prefix + name, tensor) for (name, tensor) in weights_iter)
-
-    def _load_weights_with_gguf(self, model: nn.Module) -> set[str]:
-        sources = self._get_weight_sources(model)
-        loaded: set[str] = set()
-        loadable_names: set[str] | None = None
-
-        for source in sources:
-            if self._is_transformer_source(source):
-                loaded |= model.load_weights(self._get_gguf_weights_iterator(source, model))
-
-                # GGUF checkpoints can be transformer-only or partially quantized.
-                # Only fall back to HF if this source still has missing loadable weights.
-                loadable_names = loadable_names or self._get_model_loadable_names(model)
-                has_missing_for_source = any(
-                    name.startswith(source.prefix) and name not in loaded for name in loadable_names
-                )
-                if not has_missing_for_source:
-                    continue
-
-                hf_iter = self._get_weights_iterator(source)
-                hf_iter = (
-                    (name, tensor) for (name, tensor) in hf_iter if name in loadable_names and name not in loaded
-                )
-                loaded |= model.load_weights(hf_iter)
-            else:
-                loaded |= model.load_weights(self._get_weights_iterator(source))
-
-        weights_to_load = self._get_expected_parameter_names(model)
-        weights_not_loaded = weights_to_load - loaded
-        if weights_not_loaded:
-            raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}")
-        return loaded
-
     def _init_from_load_format(
         self,
         load_format: str,
@@ -685,6 +628,11 @@ class DiffusersPipelineLoader:
         model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
         self.load_weights(model)
 
+        # Quantization methods must finish while parameters are ordinary local
+        # tensors. Some post-load transforms use operations (for example,
+        # torch.unique in ModelOpt NVFP4) that do not support DTensor inputs.
+        self._process_weights_after_loading(model, target_device)
+
         # Discover pipeline components (DiT, encoders, VAEs) via
         # ModuleDiscovery, which consults SupportsComponentDiscovery
         # when available and falls back to well-known attribute names.
@@ -692,6 +640,14 @@ class DiffusersPipelineLoader:
         # where the transformer lives at "pipe.transformer").
         discovered_modules = ModuleDiscovery.discover(model)
 
+        # Shard only the outermost DiTs. A pipeline may list a DiT and one of its
+        # submodules as separate DiTs (e.g. Cosmos3's transformer and the nested
+        # transformer.language_model) for offload's independent rings; for HSDP an
+        # inner DiT is already covered by its ancestor's _hsdp_shard_conditions, so
+        # sharding it again would double-wrap blocks and require the inner stack to
+        # declare its own conditions.
+        outer_dit_names, outer_dits = discovered_modules.outermost_dits()
+
         # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights
         # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel
         # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params.
@@ -703,14 +659,14 @@ class DiffusersPipelineLoader:
                 prepare_fp8_layers_for_fsdp,
             )
 
-            for trans in discovered_modules.dits:
+            for trans in outer_dits:
                 prepare_fp8_layers_for_fsdp(trans)
 
-        if not discovered_modules.dits:
+        if not outer_dits:
             raise ValueError("No DiT modules discovered for HSDP sharding")
 
-        # Apply HSDP sharding to all discovered DiT transformers
-        for name, trans in zip(discovered_modules.dit_names, discovered_modules.dits):
+        # Apply HSDP sharding to each outermost DiT transformer
+        for name, trans in zip(outer_dit_names, outer_dits):
             logger.debug("Applying HSDP to %s", name)
             apply_hsdp_to_model(trans, hsdp_config, target_device=target_device)
 
diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py b/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py
deleted file mode 100644
index 3266ba06c..000000000
--- a/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-import torch
-
-from .base import GGUFAdapter
-from .flux2_klein import Flux2KleinGGUFAdapter
-from .qwen_image import QwenImageGGUFAdapter
-from .z_image import ZImageGGUFAdapter
-
-if TYPE_CHECKING:
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-    from vllm_omni.diffusion.model_loader.diffusers_loader import (
-        DiffusersPipelineLoader,
-    )
-
-
-def get_gguf_adapter(
-    gguf_file: str,
-    model: torch.nn.Module,
-    source: DiffusersPipelineLoader.ComponentSource,
-    od_config: OmniDiffusionConfig,
-) -> GGUFAdapter:
-    adapter_classes = (QwenImageGGUFAdapter, ZImageGGUFAdapter, Flux2KleinGGUFAdapter)
-    for adapter_cls in adapter_classes:
-        if adapter_cls.is_compatible(od_config, model, source):
-            return adapter_cls(gguf_file, model, source, od_config)
-    model_type = None
-    if od_config.tf_model_config is not None:
-        model_type = od_config.tf_model_config.get("model_type")
-    supported = ", ".join(cls.__name__ for cls in adapter_classes)
-    raise ValueError(
-        "No GGUF adapter matched diffusion model "
-        f"(model_class_name={od_config.model_class_name!r}, model_type={model_type!r}). "
-        f"Supported adapters: {supported}."
-    )
-
-
-__all__ = [
-    "GGUFAdapter",
-    "Flux2KleinGGUFAdapter",
-    "QwenImageGGUFAdapter",
-    "ZImageGGUFAdapter",
-    "get_gguf_adapter",
-]
diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/base.py b/vllm_omni/diffusion/model_loader/gguf_adapters/base.py
deleted file mode 100644
index f3f4d28a0..000000000
--- a/vllm_omni/diffusion/model_loader/gguf_adapters/base.py
+++ /dev/null
@@ -1,109 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from collections.abc import Generator
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
-
-import numpy as np
-import torch
-
-try:
-    import gguf
-except ImportError:
-    # Optional dependency: only required for GGUF-quantized diffusion models.
-    # Keep this import lazy so a missing ``gguf`` does not break importing the
-    # diffusion model loader (and thus inspection/loading of every other
-    # diffusion model that imports through this package).
-    gguf = None
-
-if TYPE_CHECKING:
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-    from vllm_omni.diffusion.model_loader.diffusers_loader import (
-        DiffusersPipelineLoader,
-    )
-
-
-@dataclass
-class MappedTensor:
-    name: str
-    tensor: Any
-    tensor_type: Any
-    row_slice: slice | None = None
-    swap_scale_shift: bool = False
-
-
-class GGUFAdapter(ABC):
-    """Base class for model-specific GGUF adapters."""
-
-    def __init__(
-        self,
-        gguf_file: str,
-        model: torch.nn.Module,
-        source: DiffusersPipelineLoader.ComponentSource,
-        od_config: OmniDiffusionConfig,
-    ) -> None:
-        self.gguf_file = gguf_file
-        self.model = model
-        self.source = source
-        self.od_config = od_config
-
-    @staticmethod
-    def is_compatible(
-        od_config: OmniDiffusionConfig,
-        model: torch.nn.Module,
-        source: DiffusersPipelineLoader.ComponentSource,
-    ) -> bool:
-        return False
-
-    @abstractmethod
-    def weights_iterator(self) -> Generator[tuple[str, torch.Tensor], None, None]:
-        raise NotImplementedError
-
-
-# FIXME(Isotr0py): Sync implemnentation with upstream vLLM?
-def gguf_quant_weights_iterator(gguf_file: str) -> Generator[tuple[str, torch.Tensor]]:
-    """
-    Iterate over the quant weights in the model gguf files and convert
-    them to torch tensors.
-    Be careful of the order of yielding weight types and weights data,
-    we have to yield all weight types first before yielding any weights.
-    Otherwise it would cause issue when loading weights with for packed
-    layer with different quant types.
-    """
-
-    if gguf is None:
-        raise ImportError(
-            "The 'gguf' package is required to load GGUF-quantized diffusion models. "
-            "Install it with `pip install gguf`."
-        )
-    reader = gguf.GGUFReader(gguf_file)
-
-    for tensor in reader.tensors:
-        weight_type = tensor.tensor_type
-        name = tensor.name
-
-        if weight_type.name not in ("F32", "F16"):
-            weight_type_name = name.replace("weight", "qweight_type")
-            weight_type = torch.tensor(weight_type)
-            yield weight_type_name, weight_type
-
-    for tensor in reader.tensors:
-        weight = tensor.data
-        weight_type = tensor.tensor_type
-        name = tensor.name
-        if weight_type.name not in ("F32", "F16"):
-            name = name.replace("weight", "qweight")
-        if weight_type.name == "BF16" and tensor.data.dtype == np.uint8:
-            # BF16 is currently the only "quantization" type that isn't
-            # actually quantized but is read as a raw byte tensor.
-            # Reinterpret as `torch.bfloat16` tensor.
-            weight = weight.view(np.uint16)
-            if reader.byte_order == "S":
-                # GGUF endianness != system endianness
-                weight = weight.byteswap()
-            param = torch.tensor(weight).view(torch.bfloat16)
-        else:
-            param = torch.tensor(weight)
-        yield name, param
diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py b/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py
deleted file mode 100644
index 4ee6439b2..000000000
--- a/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py
+++ /dev/null
@@ -1,98 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from collections.abc import Iterable
-
-import torch
-from vllm.model_executor.models.utils import WeightsMapper
-
-from .base import GGUFAdapter, gguf_quant_weights_iterator
-
-FLUX2_TRANSFORMER_KEYS_RENAME_DICT = {
-    "single_blocks.": "single_transformer_blocks.",
-    # Image and text input projections
-    "img_in": "x_embedder",
-    "txt_in": "context_embedder",
-    # Timestep and guidance embeddings
-    "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1",
-    "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2",
-    "guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1",
-    "guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2",
-    # Modulation parameters
-    "double_stream_modulation_img.lin": "double_stream_modulation_img.linear",
-    "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear",
-    "single_stream_modulation.lin": "single_stream_modulation.linear",
-    # Final output layer
-    # "final_layer.adaLN_modulation.1": "norm_out.linear",  # Handle separately since we need to swap mod params
-    "final_layer.linear": "proj_out",
-}
-
-FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = {
-    "final_layer.adaLN_modulation.1": "norm_out.linear",
-}
-
-FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = {
-    "double_blocks.": "transformer_blocks.",
-    # Handle fused QKV projections separately as we need to break into Q, K, V projections
-    "img_attn.norm.query_norm": "attn.norm_q",
-    "img_attn.norm.key_norm": "attn.norm_k",
-    "img_attn.proj": "attn.to_out.0",
-    "img_mlp.0": "ff.linear_in",
-    "img_mlp.2": "ff.linear_out",
-    "txt_attn.norm.query_norm": "attn.norm_added_q",
-    "txt_attn.norm.key_norm": "attn.norm_added_k",
-    "txt_attn.proj": "attn.to_add_out",
-    "txt_mlp.0": "ff_context.linear_in",
-    "txt_mlp.2": "ff_context.linear_out",
-    # Additional for fuse qkv
-    "img_attn.qkv": "attn.to_qkv",
-    "txt_attn.qkv": "attn.add_kv_proj",
-}
-
-FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = {
-    "linear1": "attn.to_qkv_mlp_proj",
-    "linear2": "attn.to_out",
-    "norm.query_norm": "attn.norm_q",
-    "norm.key_norm": "attn.norm_k",
-}
-
-
-class Flux2KleinGGUFAdapter(GGUFAdapter):
-    """GGUF adapter for Flux2-Klein models with qkv splitting and adaLN swap."""
-
-    @staticmethod
-    def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
-        model_class = od_config.model_class_name or ""
-        if model_class.startswith("Flux2"):
-            return True
-        cfg = od_config.tf_model_config
-        if cfg is not None:
-            model_type = str(cfg.get("model_type", "")).lower()
-            if model_type.startswith("flux"):
-                return True
-        return False
-
-    gguf_to_hf_mapper = WeightsMapper(
-        # double_stream_modulation
-        orig_to_new_prefix=FLUX2_TRANSFORMER_KEYS_RENAME_DICT | FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP,
-        orig_to_new_substr=FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP | FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP,
-    )
-
-    def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
-        def custom_weights_adapter(
-            weights: Iterable[tuple[str, torch.Tensor]],
-        ) -> Iterable[tuple[str, torch.Tensor]]:
-            for name, weight in weights:
-                # Handle the special case for adaLN modulation parameters that require swapping shift and scale
-                if name.endswith(".scale"):
-                    name = name.replace(".scale", ".weight")
-                if name == "norm_out.linear.weight":
-                    shift, scale = weight.chunk(2, dim=0)
-                    weight = torch.cat([scale, shift], dim=0)
-                    yield name, weight
-                else:
-                    yield name, weight
-
-        weights = gguf_quant_weights_iterator(self.gguf_file)
-        weights = self.gguf_to_hf_mapper.apply(weights)
-        yield from custom_weights_adapter(weights)
diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py b/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py
deleted file mode 100644
index cc8ccea5c..000000000
--- a/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py
+++ /dev/null
@@ -1,27 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from collections.abc import Iterable
-
-import torch
-
-from .base import GGUFAdapter, gguf_quant_weights_iterator
-
-
-class QwenImageGGUFAdapter(GGUFAdapter):
-    """GGUF adapter for the Qwen-Image transformer family."""
-
-    @staticmethod
-    def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
-        model_class = od_config.model_class_name or ""
-        if model_class.startswith("QwenImage"):
-            return True
-        cfg = od_config.tf_model_config
-        if cfg is not None:
-            model_type = str(cfg.get("model_type", "")).lower()
-            if model_type.startswith("qwen_image"):
-                return True
-        return False
-
-    def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
-        yield from gguf_quant_weights_iterator(self.gguf_file)
diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py b/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py
deleted file mode 100644
index 7d8963355..000000000
--- a/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-from __future__ import annotations
-
-from collections.abc import Iterable
-
-import torch
-from vllm.model_executor.models.utils import WeightsMapper
-
-from .base import GGUFAdapter, gguf_quant_weights_iterator
-
-Z_IMAGE_KEYS_RENAME_DICT = {
-    "final_layer.": "all_final_layer.2-1.",
-    "x_embedder.": "all_x_embedder.2-1.",
-    ".attention.qkv": ".attention.to_qkv",
-    ".attention.k_norm": ".attention.norm_k",
-    ".attention.q_norm": ".attention.norm_q",
-    ".attention.out": ".attention.to_out.0",
-    "model.diffusion_model.": "",
-}
-
-
-class ZImageGGUFAdapter(GGUFAdapter):
-    """GGUF adapter for Z-Image models with QKV/FFN shard support."""
-
-    @staticmethod
-    def is_compatible(od_config, model: torch.nn.Module, source) -> bool:
-        model_class = od_config.model_class_name or ""
-        if model_class.startswith("ZImage"):
-            return True
-        cfg = od_config.tf_model_config
-        if cfg is not None:
-            model_type = str(cfg.get("model_type", "")).lower()
-            if model_type in {"z_image", "zimage", "z-image"}:
-                return True
-        return False
-
-    gguf_to_hf_mapper = WeightsMapper(
-        orig_to_new_substr=Z_IMAGE_KEYS_RENAME_DICT,
-    )
-
-    def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]:
-        weights = gguf_quant_weights_iterator(self.gguf_file)
-        yield from self.gguf_to_hf_mapper.apply(weights)
diff --git a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py b/vllm_omni/diffusion/models/audiox/pipeline_audiox.py
index 02639f1af..395e06dc8 100644
--- a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py
+++ b/vllm_omni/diffusion/models/audiox/pipeline_audiox.py
@@ -60,6 +60,15 @@ def get_audiox_post_process_func(od_config: OmniDiffusionConfig):
     """Convert the pipeline's float audio tensor to a CPU numpy array for serving."""
 
     def post_process_func(audio: torch.Tensor) -> Any:
+        if isinstance(audio, dict) and isinstance(audio.get("payload"), dict):
+            payload = dict(audio["payload"])
+            audio_payload = payload.get("audio")
+            if isinstance(audio_payload, torch.Tensor):
+                payload["audio"] = audio_payload.detach().cpu().float().numpy()
+            return {
+                "payload": payload,
+                "metadata": audio.get("metadata") or {},
+            }
         if isinstance(audio, torch.Tensor):
             return audio.detach().cpu().float().numpy()
         return audio
@@ -915,9 +924,10 @@ class AudioXPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMix
             audio = audio[..., :target_samples]
 
         stage_durations = self.stage_durations if getattr(self, "enable_diffusion_pipeline_profiler", False) else None
-        custom_output = {"audiox_task": task_norm}
         return DiffusionOutput(
-            output=audio,
-            custom_output=custom_output,
+            output={
+                "payload": {"audio": audio},
+                "metadata": {"audio": {"audiox_task": task_norm}},
+            },
             stage_durations=stage_durations,
         )
diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
index 453622765..9786e2b47 100644
--- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
+++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py
@@ -739,8 +739,10 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                         text_output = text_output.split("<|im_start|>")[-1]
 
             return DiffusionOutput(
-                output=text_output,
-                custom_output={"text_output": text_output},
+                output={
+                    "payload": {"text": text_output},
+                    "metadata": {"text": {"text_output": text_output}},
+                },
                 stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
             )
 
@@ -847,22 +849,28 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         if trajectory_log_probs:
             trajectory_log_probs_stacked = torch.stack(trajectory_log_probs)
 
-        custom = {}
+        payload = {"image": img}
+        metadata = {}
         if think_text is not None:
-            custom["think_text"] = think_text
-        # Mirror the PIL image into ``custom_output`` so callers reading via
-        # the orchestrator IPC boundary (which strips the bare ``output``
-        # field) can still recover the result.  ``video_frames`` already
-        # uses this pattern.
-        custom["image"] = img
+            metadata["text"] = {"think_text": think_text}
+        trajectory_payload = {}
+        if trajectory_latents_stacked is not None:
+            trajectory_payload["latents"] = trajectory_latents_stacked
+        if trajectory_timesteps_stacked is not None:
+            trajectory_payload["timesteps"] = trajectory_timesteps_stacked
+        if trajectory_log_probs_stacked is not None:
+            trajectory_payload["log_probs"] = trajectory_log_probs_stacked
+        if trajectory_decoded is not None:
+            trajectory_payload["decoded"] = trajectory_decoded
+        if trajectory_payload:
+            payload["trajectory"] = trajectory_payload
+            metadata["trajectory"] = {"type": "denoising"}
 
         return DiffusionOutput(
-            output=img,
-            trajectory_latents=trajectory_latents_stacked,
-            trajectory_timesteps=trajectory_timesteps_stacked,
-            trajectory_log_probs=trajectory_log_probs_stacked,
-            trajectory_decoded=trajectory_decoded,
-            custom_output=custom,
+            output={
+                "payload": payload,
+                "metadata": metadata,
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
diff --git a/vllm_omni/diffusion/models/cosmos3/__init__.py b/vllm_omni/diffusion/models/cosmos3/__init__.py
index 43be0c37a..af4eeaacf 100644
--- a/vllm_omni/diffusion/models/cosmos3/__init__.py
+++ b/vllm_omni/diffusion/models/cosmos3/__init__.py
@@ -3,16 +3,16 @@
 
 from .pipeline_cosmos3 import (
     Cosmos3OmniDiffusersPipeline,
-    get_cosmos3_action_post_process_func,
     get_cosmos3_post_process_func,
     get_cosmos3_pre_process_func,
 )
 from .transformer_cosmos3 import Cosmos3VFMTransformer
+from .transformer_cosmos3_edge import Cosmos3EdgeVFMTransformer
 
 __all__ = [
     "Cosmos3OmniDiffusersPipeline",
-    "get_cosmos3_action_post_process_func",
     "get_cosmos3_post_process_func",
     "get_cosmos3_pre_process_func",
     "Cosmos3VFMTransformer",
+    "Cosmos3EdgeVFMTransformer",
 ]
diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py
index 27f6d36c2..c5fe41a91 100644
--- a/vllm_omni/diffusion/models/cosmos3/action.py
+++ b/vllm_omni/diffusion/models/cosmos3/action.py
@@ -4,8 +4,8 @@
 
 These helpers cover action modes that use action tokens as an auxiliary output
 stream: ``policy``, ``forward_dynamics``, and ``inverse_dynamics``. The pipeline
-returns predicted actions through ``custom_output`` rather than as a normal
-``modalities`` entry.
+returns predicted actions through the diffusion output payload/metadata
+envelope rather than as a video frame stream.
 """
 
 from __future__ import annotations
diff --git a/vllm_omni/diffusion/models/cosmos3/guardrails.py b/vllm_omni/diffusion/models/cosmos3/guardrails.py
index e249d2826..3c6166e85 100644
--- a/vllm_omni/diffusion/models/cosmos3/guardrails.py
+++ b/vllm_omni/diffusion/models/cosmos3/guardrails.py
@@ -21,6 +21,7 @@ from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.progress_bar import _is_rank_zero
 from vllm_omni.errors import GuardrailViolationError
+from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -28,6 +29,28 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+# ---------------------------------------------------------------------------
+# NPU/XPU compatibility: cosmos_guardrail.cosmos_utils.load_model calls
+# ``torch.load(path, weights_only=True)`` without ``map_location``.
+# On NPU/XPU, ``torch.cuda.is_available()`` is False, so deserializing a
+# CUDA checkpoint with ``weights_only=True`` raises an UnpicklingError.
+# Work around by patching ``torch.load`` to default ``map_location="cpu"``
+# when CUDA is unavailable and no explicit map_location was given.
+# ---------------------------------------------------------------------------
+_original_torch_load = torch.load
+
+
+def _patched_torch_load(*args, **kwargs):
+    if (
+        "map_location" not in kwargs
+        and not torch.cuda.is_available()
+        and (current_omni_platform.is_npu() or current_omni_platform.is_xpu())
+    ):
+        kwargs["map_location"] = "cpu"
+    return _original_torch_load(*args, **kwargs)
+
+
+torch.load = _patched_torch_load  # type: ignore[assignment]
 
 try:
     from cosmos_guardrail import CosmosSafetyChecker
@@ -75,7 +98,7 @@ def _build_text_guardrail(checker: Any) -> TextGuardrailFn:
 
 def _build_video_guardrail(checker: Any, offload_to_cpu: bool) -> VideoGuardrailFn:
     video_models = _nn_models(checker.video_guardrail)
-    compute_device = "cuda"
+    compute_device = current_omni_platform.device_type
 
     def video_guardrail(frames: np.ndarray) -> np.ndarray:
         if offload_to_cpu:
@@ -112,9 +135,9 @@ def _init_default_guardrails(offload_to_cpu: bool = False) -> None:
     checker = CosmosSafetyChecker()
 
     # Place text models on their resting device permanently. Video models
-    # idle on CPU when offload is on and move to GPU per-call (handled in
+    # idle on CPU when offload is on and move to compute device per-call (handled in
     # the video guardrail closure).
-    idle_device = "cpu" if offload_to_cpu else "cuda"
+    idle_device = "cpu" if offload_to_cpu else current_omni_platform.device_type
     for m in _nn_models(checker.text_guardrail):
         m.to(idle_device)
     for m in _nn_models(checker.video_guardrail):
diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
index 4178115d5..41a3e80fc 100644
--- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py
@@ -20,7 +20,8 @@ Video generation is further specialized by inputs and extra args:
   transfer video generation.
 * ``action_mode``: action-capable video generation. RoboLab/OpenPI observation
   payloads in ``extra_args["robot_obs"]`` or ``extra_args["observation"]``
-  bypass normal video output and return action-only custom output.
+  bypass normal video output and return an action-only payload/metadata
+  envelope.
 
 Generated sound is video-only, cannot be combined with action or transfer, and
 is produced from sound latents rather than from ``multi_modal_data["audio"]``.
@@ -28,6 +29,7 @@ is produced from sound latents rather than from ``multi_modal_data["audio"]``.
 
 from __future__ import annotations
 
+import json
 import math
 import os
 import time
@@ -38,7 +40,6 @@ from typing import Any, ClassVar
 import numpy as np
 import PIL.Image
 import torch
-from diffusers import UniPCMultistepScheduler
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
 from torch import nn
@@ -57,8 +58,15 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL
 from vllm_omni.diffusion.models.interface import (
     ReferenceVideoDecodeSpec,
     SupportImageInput,
+    SupportsComponentDiscovery,
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
+from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
+    FlowUniPCMultistepScheduler,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -77,6 +85,7 @@ from .action import (
     normalize_action_mode,
     pad_action_to_dim,
     resolve_domain_id,
+    vision_condition_indexes,
 )
 from .transfer import (
     Cosmos3TransferConfig,
@@ -91,7 +100,8 @@ from .transfer import (
     transfer_max_frames_from_extra_args,
     uint8_cthw_to_normalized_5d,
 )
-from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen
+from .transformer_cosmos3 import Cosmos3VFMTransformer, _tf_config_get, resolve_sound_gen
+from .transformer_cosmos3_edge import COSMOS3_EDGE_BACKBONE_TYPE, Cosmos3EdgeVFMTransformer
 from .utils import (
     COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION,
     COSMOS3_VAE_TEMPORAL_COMPRESSION,
@@ -131,6 +141,7 @@ from .utils import (
 
 logger = init_logger(__name__)
 
+
 COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = (
     max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1
 )
@@ -150,6 +161,7 @@ COSMOS3_T2V_DEFAULT_WIDTH = 1280
 COSMOS3_T2V_DEFAULT_NUM_FRAMES = 189
 COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS = 35
 COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE = 6.0
+COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT = 10.0
 
 COSMOS3_T2I_DEFAULT_HEIGHT = 1024
 COSMOS3_T2I_DEFAULT_WIDTH = 1024
@@ -158,12 +170,73 @@ COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE = 7.0
 COSMOS3_T2I_DEFAULT_FLOW_SHIFT = 3.0
 COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0)
 
+COSMOS3_EDGE_T2V_DEFAULT_HEIGHT = 480
+COSMOS3_EDGE_T2V_DEFAULT_WIDTH = 832
+COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE = 5.0
+COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT = 3.0
+COSMOS3_EDGE_T2I_DEFAULT_HEIGHT = 640
+COSMOS3_EDGE_T2I_DEFAULT_WIDTH = 640
+
+COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS = "FlowMatchEulerDiscreteScheduler"
+
 # Truncation cap on the prompt token count (shared by T2I and T2V).  Prompts
 # are tokenized to their natural length (no padding); this only bounds the
 # UND pathway / GEN cross-attention cost for pathologically long prompts.
 COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
 
 
+def _ceil_video_num_frames(num_frames: int, temporal_compression_factor: int) -> int:
+    """Round a video length up to the causal VAE's ``factor * k + 1`` grid."""
+    if temporal_compression_factor <= 0:
+        raise ValueError(f"Cosmos3 temporal_compression_factor must be positive, got {temporal_compression_factor}.")
+    if num_frames <= 1:
+        return num_frames
+    latent_frames = math.ceil((num_frames - 1) / temporal_compression_factor)
+    return latent_frames * temporal_compression_factor + 1
+
+
+def _format_json_object_prompt(
+    prompt: str,
+    *,
+    num_frames: int,
+    frame_rate: float,
+    height: int,
+    width: int,
+    aspect_ratio: str | None,
+) -> str | None:
+    """Match Imaginaire4 metadata injection for JSON-object prompts."""
+    try:
+        prompt_obj = json.loads(prompt)
+    except (json.JSONDecodeError, TypeError, ValueError):
+        return None
+    if not isinstance(prompt_obj, dict):
+        return None
+
+    metadata: dict[str, Any] = {}
+    if num_frames > 1:
+        duration_seconds = int(num_frames / frame_rate) if frame_rate > 0 else 0
+        metadata.update({"duration": f"{duration_seconds}s", "fps": float(frame_rate)})
+    else:
+        prompt_obj.pop("duration", None)
+        prompt_obj.pop("fps", None)
+    metadata["resolution"] = {"H": int(height), "W": int(width)}
+    if aspect_ratio is not None:
+        metadata["aspect_ratio"] = aspect_ratio
+
+    prompt_obj.update(metadata)
+    return json.dumps(prompt_obj)
+
+
+def resolve_cosmos3_transformer_cls(model_config: Any) -> type[Cosmos3VFMTransformer]:
+    """Select the Cosmos3 transformer implementation from transformer/config.json."""
+    backbone_type = _tf_config_get(model_config, "backbone_type", None)
+    if backbone_type is None:
+        return Cosmos3VFMTransformer
+    if backbone_type == COSMOS3_EDGE_BACKBONE_TYPE:
+        return Cosmos3EdgeVFMTransformer
+    raise ValueError(f"Unsupported Cosmos3 transformer backbone_type={backbone_type!r}.")
+
+
 # ---------------------------------------------------------------------------
 # Post-process function (registered in registry.py)
 # ---------------------------------------------------------------------------
@@ -186,6 +259,9 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
     from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled
 
     video_processor = VideoProcessor(vae_scale_factor=16)
+    is_edge_model = (
+        _tf_config_get(getattr(od_config, "tf_model_config", None), "backbone_type", None) == COSMOS3_EDGE_BACKBONE_TYPE
+    )
     # Eager-load guardrail models at pipeline build time when the server-level
     # gate is on. Per-request overrides only decide whether the loaded models
     # are *invoked* — they cannot turn checks on without a server-side preload.
@@ -406,12 +482,17 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
         extra = _extra_args(request)
         transfer_requested = action_mode is None and has_transfer_hints(extra)
 
-        # Auto-calculate H/W from aspect ratio (720p max area)
+        # Resolve missing H/W.
         if transfer_requested:
             _set_transfer_size_from_image(request, image)
         elif request.sampling_params.height is None or request.sampling_params.width is None:
             if action_mode is not None:
                 _set_action_size_from_image(request, image)
+            elif is_edge_model:
+                if request.sampling_params.height is None:
+                    request.sampling_params.height = COSMOS3_EDGE_T2V_DEFAULT_HEIGHT
+                if request.sampling_params.width is None:
+                    request.sampling_params.width = COSMOS3_EDGE_T2V_DEFAULT_WIDTH
             else:
                 max_area = 720 * 1280
                 aspect_ratio = image.height / image.width
@@ -528,6 +609,37 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
         if output_type == "latent":
             return output
 
+        def _postprocess_action(action: Any, metadata: dict[str, Any]) -> Any:
+            internal_metadata = metadata.get("internal")
+            inputs = (
+                internal_metadata.get("robolab_action_postprocess") if isinstance(internal_metadata, dict) else None
+            )
+            if isinstance(inputs, RoboLabActionPostprocessInputs):
+                return postprocess_robolab_action(action, inputs)
+            return action
+
+        pending_action = None
+        pending_action_metadata: dict[str, Any] = {}
+        envelope_public_metadata: dict[str, Any] = {}
+        if isinstance(output, dict) and isinstance(output.get("payload"), dict):
+            envelope_payload = dict(output.get("payload") or {})
+            metadata = output.get("metadata") or {}
+            envelope_metadata = metadata if isinstance(metadata, dict) else {}
+            envelope_public_metadata = {key: value for key, value in envelope_metadata.items() if key != "internal"}
+            action = envelope_payload.pop("actions", None)
+            if action is not None:
+                pending_action = _postprocess_action(action, envelope_metadata)
+                pending_action_metadata = envelope_public_metadata
+                if not envelope_payload:
+                    return {
+                        "payload": {
+                            "video": [],
+                            "actions": pending_action,
+                        },
+                        "metadata": pending_action_metadata,
+                    }
+            output = envelope_payload
+
         audio = None
         audio_sample_rate = None
         if isinstance(output, dict):
@@ -566,11 +678,31 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
                 # check_video_safety expects a 5D tensor; re-add T axis.
                 checked = check_video_safety(image.unsqueeze(2))
                 image = checked.squeeze(2)
-            return video_processor.postprocess(image, output_type="pil")
-        if is_guardrails_enabled(od_config, sampling_params):
+            processed_image = video_processor.postprocess(image, output_type="pil")
+            if envelope_public_metadata:
+                return {
+                    "payload": {"image": processed_image},
+                    "metadata": envelope_public_metadata,
+                }
+            return processed_image
+        guardrails_enabled = is_guardrails_enabled(od_config, sampling_params)
+        if guardrails_enabled:
             video = check_video_safety(video)
         processed_video = video_processor.postprocess_video(video, output_type=output_type)
         if audio is None:
+            if pending_action is not None:
+                return {
+                    "payload": {
+                        "video": processed_video,
+                        "actions": pending_action,
+                    },
+                    "metadata": pending_action_metadata,
+                }
+            if envelope_public_metadata:
+                return {
+                    "payload": {"video": processed_video},
+                    "metadata": envelope_public_metadata,
+                }
             return processed_video
         if isinstance(audio, torch.Tensor):
             audio = audio.detach().cpu()
@@ -581,33 +713,38 @@ def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
         }
         if audio_sample_rate is not None:
             result["audio_sample_rate"] = int(audio_sample_rate)
+        if pending_action is not None:
+            return {
+                "payload": {
+                    "video": result["video"],
+                    "audio": result["audio"],
+                    "actions": pending_action,
+                },
+                "metadata": {
+                    **pending_action_metadata,
+                    "video": {"fps": result["fps"]},
+                    "audio": {"sample_rate": result.get("audio_sample_rate")},
+                },
+            }
+        if envelope_public_metadata:
+            return {
+                "payload": {
+                    "video": result["video"],
+                    "audio": result["audio"],
+                },
+                "metadata": {
+                    **envelope_public_metadata,
+                    "video": {"fps": result["fps"], **envelope_public_metadata.get("video", {})}
+                    if isinstance(envelope_public_metadata.get("video"), dict)
+                    else {"fps": result["fps"]},
+                    "audio": {"sample_rate": result.get("audio_sample_rate")},
+                },
+            }
         return result
 
     return post_process_func
 
 
-def get_cosmos3_action_post_process_func(od_config: OmniDiffusionConfig):
-    """Build the custom-output postprocessor for Cosmos3 action predictions.
-
-    Action modes return predicted action tensors in ``custom_output`` alongside
-    normal video output. RoboLab/OpenPI policy serving marks action-only output
-    and carries observation metadata used here to map model-space actions back
-    to the requested robot action representation.
-    """
-    del od_config
-
-    def action_post_process_func(action: Any, custom_output: dict[str, Any] | None = None, sampling_params=None):
-        del sampling_params
-        inputs = custom_output.get("robolab_action_postprocess") if isinstance(custom_output, dict) else None
-        if isinstance(inputs, RoboLabActionPostprocessInputs):
-            processed_action = postprocess_robolab_action(action, inputs)
-            custom_output.pop("robolab_action_postprocess", None)
-            return processed_action
-        return action
-
-    return action_post_process_func
-
-
 def get_cosmos3_ir_op_priority_func(od_config: OmniDiffusionConfig):
     del od_config
 
@@ -627,7 +764,12 @@ def get_cosmos3_ir_op_priority_func(od_config: OmniDiffusionConfig):
 # Pipeline
 # ---------------------------------------------------------------------------
 class Cosmos3OmniDiffusersPipeline(
-    nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin
+    nn.Module,
+    CFGParallelMixin,
+    SupportImageInput,
+    SupportsComponentDiscovery,
+    ProgressBarMixin,
+    DiffusionPipelineProfilerMixin,
 ):
     """Cosmos3 text/image/video/sound/action pipeline.
 
@@ -662,14 +804,19 @@ class Cosmos3OmniDiffusersPipeline(
       T2I, transfer, and action+sound are rejected.
     * **Action generation** when ``action_mode`` is provided. ``policy`` and
       ``forward_dynamics`` require an image or video input; ``inverse_dynamics``
-      requires video input. Action predictions are returned in ``custom_output``.
+      requires video input. Action predictions are returned in the diffusion
+      output payload/metadata envelope.
       RoboLab/OpenPI observations in ``extra_args['robot_obs']`` or
-      ``extra_args['observation']`` return action-only custom output.
+      ``extra_args['observation']`` return an action-only envelope.
     * **T2V** otherwise (default video generation).
     """
 
     support_image_input: ClassVar[bool] = True
     color_format: ClassVar[str] = "RGB"
+    _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"]
+    _encoder_modules: ClassVar[list[str]] = []
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []
 
     @classmethod
     def reference_video_decode_spec(
@@ -713,12 +860,6 @@ class Cosmos3OmniDiffusersPipeline(
         prefix: str = "",
     ) -> None:
         super().__init__()
-        if od_config.enable_cpu_offload:
-            raise ValueError(
-                "Cosmos3 has no separate text encoder, so CPU offloading "
-                "(transformer↔encoder swapping) is not supported. "
-                "Use --enable-layerwise-offload instead."
-            )
         self.od_config = od_config
         self.device = get_local_device()
         self.dtype = od_config.dtype
@@ -759,27 +900,52 @@ class Cosmos3OmniDiffusersPipeline(
             sound_latent_fps = self._sound_tokenizer.latent_fps
 
         # --- Transformer (weights loaded later via weights_sources) ---
-        self.transformer = Cosmos3VFMTransformer(
+        transformer_cls = resolve_cosmos3_transformer_cls(od_config.tf_model_config)
+        self.transformer = transformer_cls(
             od_config=od_config,
             temporal_compression_factor=self.vae_scale_factor_temporal,
             sound_gen=sound_gen,
             sound_dim=sound_dim,
             sound_latent_fps=sound_latent_fps,
         )
+        self.is_edge_model = transformer_cls is Cosmos3EdgeVFMTransformer
 
         # --- Scheduler ---
-        # Load from checkpoint to preserve solver_order, timestep_spacing,
-        # beta_schedule, sigma bounds, flow_shift, etc. Only override
-        # flow_shift when explicitly requested by the user.
-        self.scheduler = UniPCMultistepScheduler.from_pretrained(
+        # Distilled model differs from regular one only by scheduler,
+        # distilled one uses FlowMatchEulerDiscreteScheduler, while
+        # regular should use FlowUniPCMultistepScheduler
+
+        scheduler_config = FlowUniPCMultistepScheduler.load_config(
             model_path,
             subfolder="scheduler",
             local_files_only=local_files_only,
         )
-        if od_config.flow_shift is not None:
-            self.scheduler = UniPCMultistepScheduler.from_config(
-                self.scheduler.config,
-                flow_shift=od_config.flow_shift,
+
+        scheduler_class_name = scheduler_config.get("_class_name")
+
+        self.is_distilled_model = False
+        if scheduler_class_name == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS:
+            fixed_step_config = scheduler_config.get("fixed_step_sampler_config")
+            if not isinstance(fixed_step_config, dict) or fixed_step_config.get("sample_type") != "sde":
+                raise ValueError("Cosmos3 distilled scheduler requires fixed_step_sampler_config.sample_type=sde.")
+            t_list = fixed_step_config.get("t_list")
+            if not isinstance(t_list, list) or not t_list:
+                raise ValueError("Cosmos3 distilled scheduler requires a non-empty fixed_step_sampler_config.t_list.")
+            self.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+                scheduler_config,
+                stochastic_sampling=True,
+            )
+            self._scheduler_init_t_list = list(t_list)
+            self.is_distilled_model = True
+        else:
+            # Preserve compatible solver settings from the checkpoint, but keep
+            # the base shift neutral. The concrete request shift is applied when
+            # FlowUniPC builds its timesteps.
+            self.scheduler = FlowUniPCMultistepScheduler.from_config(
+                scheduler_config,
+                shift=1.0,
+                use_dynamic_shifting=False,
+                prediction_type="flow_prediction",
             )
         self._cpu_scheduler_state()
 
@@ -798,14 +964,16 @@ class Cosmos3OmniDiffusersPipeline(
             ),
         ]
 
-        # Snapshot the loaded scheduler config so we can rebuild the
-        # scheduler at request time when a per-request flow_shift override
-        # is supplied (T2I uses shift=3.0; T2V/I2V use the engine default).
-        self._base_scheduler_config = self.scheduler.config
-        self._engine_init_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0)
+        # An engine-level override becomes the video default; otherwise use
+        # the matching regular/Edge default. Per-request values still take
+        # precedence.
+        default_video_flow_shift = (
+            COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT if self.is_edge_model else COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT
+        )
+        self._engine_init_flow_shift = float(
+            od_config.flow_shift if od_config.flow_shift is not None else default_video_flow_shift
+        )
         self._current_flow_shift = self._engine_init_flow_shift
-        self._base_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas(self.scheduler.config)
-        self._current_scheduler_use_karras_sigmas = self._base_scheduler_use_karras_sigmas
 
         self._guidance_scale = None
         self._num_timesteps = None
@@ -822,6 +990,32 @@ class Cosmos3OmniDiffusersPipeline(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
         )
 
+    def enable_omni_model_cpu_offload(
+        self,
+        *,
+        device: torch.device,
+        pin_memory: bool = True,
+        use_hsdp: bool = False,
+    ) -> None:
+        """Enable Cosmos3 component-level model offload.
+
+        Cosmos3 has a nested reasoner/generator transformer instead of separate
+        text-encoder and DiT pipeline components, so the transformer owns the
+        mutual-exclusion swaps.  The VAE stays resident on GPU like the generic
+        model-level offloader.
+        """
+        self.vae.to(device, non_blocking=True)
+        if isinstance(self._sound_tokenizer, nn.Module):
+            self._sound_tokenizer.to(device)
+        self.transformer.enable_model_cpu_offload(
+            device=device,
+            pin_memory=pin_memory,
+            use_hsdp=use_hsdp,
+        )
+
+    def disable_omni_model_cpu_offload(self) -> None:
+        self.transformer.disable_model_cpu_offload()
+
     # -- Weight loading --------------------------------------------------------
 
     @staticmethod
@@ -899,6 +1093,7 @@ class Cosmos3OmniDiffusersPipeline(
             "self_attn.to_out.": f"{und_lp}.self_attn.to_out.",
             "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.",
             "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.",
+            "self_attn.k_norm_und_for_gen.": f"{und_lp}.self_attn.k_norm_und_for_gen.",
             # GEN attention
             "self_attn.add_q_proj.": f"{gen_lp}.cross_attention.to_q.",
             "self_attn.add_k_proj.": f"{gen_lp}.cross_attention.to_k.",
@@ -928,6 +1123,9 @@ class Cosmos3OmniDiffusersPipeline(
 
         return None
 
+    # Checkpoint adapters use this hook before model-specific weight loading.
+    remap_checkpoint_key = _remap_ckpt_key
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Stream-remap checkpoint weights and load via AutoWeightsLoader.
 
@@ -942,6 +1140,10 @@ class Cosmos3OmniDiffusersPipeline(
             total = kept = 0
             for name, tensor in weights:
                 total += 1
+                if name in allowed or name in tp_aware:
+                    kept += 1
+                    yield name, tensor
+                    continue
                 remapped = self._remap_ckpt_key(name)
                 if remapped is not None and (remapped in allowed or remapped in tp_aware):
                     kept += 1
@@ -957,6 +1159,7 @@ class Cosmos3OmniDiffusersPipeline(
         loaded = loader.load_weights(_remapped_weights())
         self.transformer.post_load_weights()
         self.transformer.eval()
+        self.transformer.validate_loaded_weights(loaded)
         if getattr(self.transformer, "sound_gen", False):
             sound_markers = ("audio_proj_in.", "audio_proj_out.", "audio_modality_embed")
             missing = [marker.rstrip(".") for marker in sound_markers if not any(marker in name for name in loaded)]
@@ -1262,6 +1465,8 @@ class Cosmos3OmniDiffusersPipeline(
         inputs: RoboLabPolicyInputs,
         pipeline_start: float,
     ) -> DiffusionOutput:
+        if self.is_distilled_model:
+            raise self._distilled_unsupported_error("RoboLab/action policy requests are unsupported.")
         if not getattr(self.transformer, "action_gen", False):
             raise ValueError(
                 "Cosmos3 RoboLab policy serving was requested, but the transformer "
@@ -1372,21 +1577,32 @@ class Cosmos3OmniDiffusersPipeline(
             guidance_interval=None,
             raw_action_dim=raw_action_dim,
             scheduler=scheduler,
+            generator=generator,
         )
 
         if _is_rank_zero():
             logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)
 
         action = action_latents[:, :, :raw_action_dim].detach().cpu()
-        custom_action_output: dict[str, Any] = {
-            "action": action,
-            "raw_action_dim": raw_action_dim,
-            "action_mode": action_mode,
-            "domain_id": domain_id,
-            "action_only_output": True,
-            "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
+        action_output: dict[str, Any] = {
+            "payload": {
+                "actions": action,
+            },
+            "metadata": {
+                "actions": {
+                    "raw_action_dim": raw_action_dim,
+                    "action_mode": action_mode,
+                    "domain_id": domain_id,
+                },
+                "common": {
+                    "action_only_output": True,
+                },
+                "internal": {
+                    "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
+                },
+            },
         }
-        return DiffusionOutput(output={}, custom_output=custom_action_output)
+        return DiffusionOutput(output=action_output)
 
     @staticmethod
     def _truthy(value) -> bool:
@@ -1453,38 +1669,34 @@ class Cosmos3OmniDiffusersPipeline(
             raise ValueError(f"Incorrect modality value in {modalities}, expected one of {accepted_modalities}.")
         return "image" in modalities
 
+    def _set_timesteps(self, num_inference_steps: int, device: str | torch.device, shift: float) -> None:
+        if self.is_distilled_model:
+            self.scheduler.set_timesteps(sigmas=self._scheduler_init_t_list, device=device)
+        else:
+            self.scheduler.set_timesteps(
+                num_inference_steps,
+                device=device,
+                shift=shift,
+            )
+
+    def _set_flow_shift(self, target_shift: float) -> None:
+        """Select the shift applied by FlowUniPC when timesteps are built."""
+        self._current_flow_shift = float(target_shift)
+
     @staticmethod
-    def _scheduler_use_karras_sigmas(config: Any) -> bool | None:
-        value = getattr(config, "use_karras_sigmas", None)
-        return None if value is None else bool(value)
+    def _resolve_seed(sp: OmniDiffusionSamplingParams, generator: torch.Generator | None, default: int = 42) -> int:
+        if sp.seed is not None:
+            return int(sp.seed)
+        if isinstance(generator, torch.Generator):
+            return int(generator.initial_seed())
+        return int(default)
 
-    def _set_flow_shift(self, target_shift: float, *, use_karras_sigmas: bool | None = None) -> None:
-        """Set UniPC scheduler mode for a concrete request.
-
-        The scheduler is rebuilt from the saved base config if
-        the target differs from the current shift or Karras-sigma mode.
-        Tracking explicit scheduler state is required because the previous
-        mode may have rebuilt the scheduler - we cannot rely on
-        ``self.scheduler.config`` reflecting the last requested target if a
-        rebuild was skipped via the equality check.
-        """
-        target = float(target_shift)
-        target_use_karras_sigmas = (
-            self._base_scheduler_use_karras_sigmas if use_karras_sigmas is None else bool(use_karras_sigmas)
-        )
-        if (
-            target == float(self._current_flow_shift)
-            and target_use_karras_sigmas == self._current_scheduler_use_karras_sigmas
-        ):
-            return
-
-        scheduler_kwargs: dict[str, Any] = {"flow_shift": target}
-        if use_karras_sigmas is not None:
-            scheduler_kwargs["use_karras_sigmas"] = bool(use_karras_sigmas)
-        self.scheduler = UniPCMultistepScheduler.from_config(self._base_scheduler_config, **scheduler_kwargs)
-        self._cpu_scheduler_state()
-        self._current_flow_shift = target
-        self._current_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas(self.scheduler.config)
+    def _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float:
+        if self.is_distilled_model:
+            return 1.0
+        if sp.guidance_scale_provided:
+            return float(sp.guidance_scale)
+        return float(default)
 
     def _cpu_scheduler_state(self) -> None:
         # We need to move scheduler tensors to CPU, as unipc from diffusers assumes they are on CPU.
@@ -1505,6 +1717,53 @@ class Cosmos3OmniDiffusersPipeline(
     def num_timesteps(self):
         return self._num_timesteps
 
+    @staticmethod
+    def _distilled_unsupported_error(detail: str) -> ValueError:
+        return ValueError(
+            f"Cosmos3 distilled checkpoints support only text-to-image and image-to-video generation; {detail}"
+        )
+
+    def _validate_distilled_generation_mode(
+        self,
+        *,
+        is_t2i: bool,
+        image_tensor: torch.Tensor | None,
+        action_enabled: bool,
+        transfer_config: Cosmos3TransferConfig | None,
+        is_v2v: bool,
+        sound_enabled: bool,
+    ) -> None:
+        if not self.is_distilled_model:
+            return
+        if action_enabled:
+            raise self._distilled_unsupported_error("action requests are unsupported.")
+        if transfer_config is not None:
+            raise self._distilled_unsupported_error("transfer requests are unsupported.")
+        if is_v2v:
+            raise self._distilled_unsupported_error("video-to-video requests are unsupported.")
+        if sound_enabled:
+            raise self._distilled_unsupported_error("sound generation is unsupported.")
+        if not is_t2i and image_tensor is None:
+            raise self._distilled_unsupported_error("text-to-video requests are unsupported.")
+        if is_t2i and image_tensor is not None:
+            raise self._distilled_unsupported_error("image-conditioned image generation is unsupported.")
+
+    def _validate_edge_generation_mode(
+        self,
+        *,
+        transfer_config: Cosmos3TransferConfig | None,
+        is_v2v: bool,
+        sound_enabled: bool,
+    ) -> None:
+        if not self.is_edge_model:
+            return
+        if sound_enabled:
+            raise ValueError("Cosmos3 Edge checkpoints do not support sound generation.")
+        if transfer_config is not None:
+            raise ValueError("Cosmos3 Edge checkpoints do not support transfer inference.")
+        if is_v2v:
+            raise ValueError("Cosmos3 Edge checkpoints do not support video-to-video generation.")
+
     # -- Prompt formatting -----------------------------------------------------
 
     @staticmethod
@@ -1521,6 +1780,10 @@ class Cosmos3OmniDiffusersPipeline(
         """
         Append duration and resolution metadata to a prompt.
         """
+        prompt = prompt.strip()
+        if duration_template is None and resolution_template is None:
+            return prompt
+
         parts: list[str] = []
         head = prompt.rstrip(".").strip()
         if head:
@@ -1796,15 +2059,26 @@ class Cosmos3OmniDiffusersPipeline(
             res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE
         else:
             res_tmpl = None
-        prompt = self._apply_metadata_templates(
+        json_prompt = _format_json_object_prompt(
             prompt,
-            num_frames,
-            frame_rate,
-            height,
-            width,
-            duration_template=dur_tmpl,
-            resolution_template=res_tmpl,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            height=height,
+            width=width,
+            aspect_ratio=self._get_sp_param(sp, "aspect_ratio", None),
         )
+        if json_prompt is not None:
+            prompt = json_prompt
+        else:
+            prompt = self._apply_metadata_templates(
+                prompt,
+                num_frames,
+                frame_rate,
+                height,
+                width,
+                duration_template=dur_tmpl,
+                resolution_template=res_tmpl,
+            )
         if _is_rank_zero():
             logger.info("Final prompt: '%s'", prompt)
 
@@ -2009,8 +2283,21 @@ class Cosmos3OmniDiffusersPipeline(
         generator: torch.Generator,
         image_size: Any | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Prepare video latents for action modes with mode-specific conditioning."""
+        """Prepare video latents for action modes with mode-specific conditioning.
+
+        Policy and forward-dynamics modes condition only latent frame zero.
+        The Wan VAE is temporally causal, so encoding only the first pixel
+        frame preserves that latent while avoiding unused future-frame work.
+        Inverse dynamics conditions every latent and keeps the full encode.
+        """
         del height, width
+        if video_tensor.ndim == 4:
+            video_tensor = video_tensor.unsqueeze(0)
+        if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
+            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
+        if video_tensor.shape[2] < 1:
+            raise ValueError("Cosmos3 action video tensor must contain at least one frame.")
+
         C = self.transformer.latent_channel_size
         T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
         latent_hw = self._latent_hw_from_image_size(image_size)
@@ -2026,12 +2313,24 @@ class Cosmos3OmniDiffusersPipeline(
             device=self.device,
             dtype=self.dtype,
         )
-        cond_latent = self._encode_video_tensor(video_tensor, image_size=image_size)
-        if cond_latent.shape[2:] != noise.shape[2:]:
+        condition_indexes = vision_condition_indexes(mode, num_frames, self.vae_scale_factor_temporal)
+        condition_video = video_tensor[:, :, :1] if condition_indexes == [0] else video_tensor
+        cond_prefix_latent = self._encode_video_tensor(condition_video, image_size=image_size)
+        expected_prefix = (1, C, max(condition_indexes) + 1, H_lat, W_lat)
+        if (
+            cond_prefix_latent.shape[0] != expected_prefix[0]
+            or cond_prefix_latent.shape[1] != expected_prefix[1]
+            or cond_prefix_latent.shape[2] < expected_prefix[2]
+            or cond_prefix_latent.shape[3:] != expected_prefix[3:]
+        ):
             raise ValueError(
                 "Cosmos3 action video latent shape mismatch: "
-                f"encoded={tuple(cond_latent.shape)}, expected={tuple(noise.shape)}."
+                f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}."
             )
+
+        condition_latents = torch.zeros_like(noise)
+        for index in condition_indexes:
+            condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1]
         condition_mask = build_vision_condition_mask(
             mode,
             num_frames,
@@ -2039,9 +2338,9 @@ class Cosmos3OmniDiffusersPipeline(
             device=self.device,
             dtype=self.dtype,
         )
-        latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise
+        latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
         velocity_mask = 1.0 - condition_mask
-        return latents, velocity_mask, cond_latent
+        return latents, velocity_mask, condition_latents
 
     def _prepare_action_latents(
         self,
@@ -2140,6 +2439,7 @@ class Cosmos3OmniDiffusersPipeline(
         guidance_interval: tuple[float, float] | None = None,
         raw_action_dim: int | None = None,
         scheduler: Any | None = None,
+        generator: torch.Generator | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Denoising loop with 3-mode CFG support (parallel, sequential, none).
 
@@ -2251,11 +2551,23 @@ class Cosmos3OmniDiffusersPipeline(
                 if raw_action_dim is not None and 0 < raw_action_dim < action_pred.shape[-1]:
                     action_pred[..., raw_action_dim:] = 0
             if action_latents is None and sound_latents is None:
-                latents = step_scheduler.step(video_pred, t, latents, return_dict=False)[0]
+                latents = step_scheduler.step(
+                    video_pred,
+                    t,
+                    latents,
+                    generator=generator,
+                    return_dict=False,
+                )[0]
             else:
                 packed_noise, shapes, numels = _pack_joint(video_pred, action_pred, sound_pred)
                 packed_latents, _, _ = _pack_joint(latents, action_latents, sound_latents)
-                packed_next = step_scheduler.step(packed_noise, t, packed_latents, return_dict=False)[0]
+                packed_next = step_scheduler.step(
+                    packed_noise,
+                    t,
+                    packed_latents,
+                    generator=generator,
+                    return_dict=False,
+                )[0]
                 unpacked = _unpack_joint(packed_next, shapes, numels)
                 latents = unpacked[0]
                 idx = 1
@@ -2481,6 +2793,7 @@ class Cosmos3OmniDiffusersPipeline(
         velocity_mask: torch.Tensor,
         condition_latents: torch.Tensor,
         guidance_interval: tuple[float, float] | None = None,
+        generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool:
             if interval is None:
@@ -2592,7 +2905,13 @@ class Cosmos3OmniDiffusersPipeline(
                 if isinstance(noise_pred, tuple):
                     raise ValueError("Cosmos3 transfer diffusion expects video-only tensor predictions.")
                 noise_pred = noise_pred * velocity_mask
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                latents = self.scheduler.step(
+                    noise_pred,
+                    t,
+                    latents,
+                    generator=generator,
+                    return_dict=False,
+                )[0]
                 latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
         finally:
             self._cosmos3_branch_caches = None
@@ -2609,6 +2928,8 @@ class Cosmos3OmniDiffusersPipeline(
         transfer_video_tensor: torch.Tensor | None,
         transfer_input_fps: float | None,
     ) -> DiffusionOutput:
+        if self.is_distilled_model:
+            raise self._distilled_unsupported_error("transfer requests are unsupported.")
         input_frames = None
         if transfer_video_tensor is not None:
             input_frames = normalized_video_to_uint8_cthw(transfer_video_tensor)
@@ -2674,7 +2995,7 @@ class Cosmos3OmniDiffusersPipeline(
         guidance_scale = (
             float(transfer_config.guidance_scale)
             if transfer_config.guidance_scale is not None
-            else float(sp.guidance_scale or COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE)
+            else self._resolve_guidance_scale(sp, COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE)
         )
         flow_shift_target = float(
             transfer_config.flow_shift
@@ -2689,11 +3010,11 @@ class Cosmos3OmniDiffusersPipeline(
 
         self._guidance_scale = guidance_scale
         self._num_timesteps = num_inference_steps
-        self._set_flow_shift(flow_shift_target, use_karras_sigmas=False)
+        self._set_flow_shift(flow_shift_target)
 
         generator = sp.generator
+        seed = self._resolve_seed(sp, generator)
         if generator is None:
-            seed = sp.seed if sp.seed is not None else 42
             generator = torch.Generator(device=self.device).manual_seed(seed)
 
         cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts(
@@ -2783,7 +3104,11 @@ class Cosmos3OmniDiffusersPipeline(
                 transfer_share_vision_temporal_positions=transfer_config.share_vision_temporal_positions,
             )
 
-            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+            self._set_timesteps(
+                num_inference_steps,
+                device=self.device,
+                shift=self._current_flow_shift,
+            )
             latents = self.diffuse_transfer(
                 latents=latents,
                 timesteps=self.scheduler.timesteps,
@@ -2798,6 +3123,7 @@ class Cosmos3OmniDiffusersPipeline(
                 shared_kwargs=shared_kwargs,
                 velocity_mask=velocity_mask,
                 condition_latents=condition_latents,
+                generator=generator,
             )
             output_video = self._decode_latents(latents).clamp(-1, 1)
             previous_output = output_video
@@ -2825,11 +3151,15 @@ class Cosmos3OmniDiffusersPipeline(
             full_output = torch.cat([normalized_input.to(full_output), full_output], dim=-1)
 
         return DiffusionOutput(
-            output={"video": full_output},
-            custom_output={
-                "transfer_controls": full_controls,
-                "transfer_hints": list(per_hint_frames),
-                "fps": frame_rate,
+            output={
+                "payload": {"video": full_output},
+                "metadata": {
+                    "transfer": {
+                        "controls": full_controls,
+                        "hints": list(per_hint_frames),
+                    },
+                    "video": {"fps": frame_rate},
+                },
             },
         )
 
@@ -2873,6 +3203,20 @@ class Cosmos3OmniDiffusersPipeline(
         action_enabled = action_mode is not None
         transfer_config = resolve_transfer_config(sp, prompt_data)
         action_video_tensor = video_tensor if action_enabled else None
+        is_v2v = video_tensor is not None and not is_t2i and not action_enabled
+        self._validate_distilled_generation_mode(
+            is_t2i=is_t2i,
+            image_tensor=image_tensor,
+            action_enabled=action_enabled,
+            transfer_config=transfer_config,
+            is_v2v=is_v2v,
+            sound_enabled=sound_enabled,
+        )
+        self._validate_edge_generation_mode(
+            transfer_config=transfer_config,
+            is_v2v=is_v2v,
+            sound_enabled=sound_enabled,
+        )
         if transfer_config is not None:
             if is_t2i:
                 raise ValueError("Cosmos3 transfer inference is supported only for video outputs.")
@@ -2920,31 +3264,36 @@ class Cosmos3OmniDiffusersPipeline(
             raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
         if video_tensor is not None and is_t2i:
             raise ValueError("Cosmos3 video-to-video generation is supported only for video outputs.")
-        is_v2v = video_tensor is not None and not is_t2i and not action_enabled
 
-        # T2I and T2V share the same model + forward path; only defaults
-        # differ:
-        #   T2I: 1024x1024, 50 steps, shift=3.0, guidance_interval=[400, 1000]
-        #   T2V: 720x1280,  35 steps, shift=engine-init, no interval
+        # T2I and T2V share the same model and forward path; their defaults are:
+        #   T2I: regular 1024x1024, Edge 640x640; 50 steps, shift=3.0,
+        #        guidance_interval=[400, 1000]
+        #   T2V: 189 frames and 35 steps; regular guidance=6, shift=10;
+        #        Edge guidance=5, shift=3;
+        #        no guidance interval
         if is_t2i:
-            height = sp.height or COSMOS3_T2I_DEFAULT_HEIGHT
-            width = sp.width or COSMOS3_T2I_DEFAULT_WIDTH
+            height = sp.height or (
+                COSMOS3_EDGE_T2I_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2I_DEFAULT_HEIGHT
+            )
+            width = sp.width or (COSMOS3_EDGE_T2I_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2I_DEFAULT_WIDTH)
             num_frames = 1
             num_inference_steps = sp.num_inference_steps or COSMOS3_T2I_DEFAULT_NUM_INFERENCE_STEPS
-            guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE
+            guidance_scale = self._resolve_guidance_scale(sp, COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE)
             default_flow_shift = COSMOS3_T2I_DEFAULT_FLOW_SHIFT
             default_guidance_interval: tuple[float, float] | None = COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL
             batch_size = max(1, int(sp.num_outputs_per_prompt or 1))
         else:
-            height = sp.height or COSMOS3_T2V_DEFAULT_HEIGHT
-            width = sp.width or COSMOS3_T2V_DEFAULT_WIDTH
+            height = sp.height or (
+                COSMOS3_EDGE_T2V_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2V_DEFAULT_HEIGHT
+            )
+            width = sp.width or (COSMOS3_EDGE_T2V_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2V_DEFAULT_WIDTH)
+            default_guidance_scale = (
+                COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE if self.is_edge_model else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE
+            )
             num_frames = sp.num_frames or COSMOS3_T2V_DEFAULT_NUM_FRAMES
             num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS
-            guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE
-            # Fall back to the engine-init shift, NOT None: passing None
-            # to ``_set_flow_shift`` would leak a prior T2I rebuild
-            # (shift=3.0) into a subsequent video request.
-            default_flow_shift = COSMOS3_V2V_DEFAULT_FLOW_SHIFT if is_v2v else self._engine_init_flow_shift
+            guidance_scale = self._resolve_guidance_scale(sp, default_guidance_scale)
+            default_flow_shift = self._engine_init_flow_shift
             default_guidance_interval = None
             batch_size = 1  # Existing video pipeline assumes B=1.
 
@@ -2967,9 +3316,23 @@ class Cosmos3OmniDiffusersPipeline(
                     f"or action_chunk_size + 1; got num_frames={num_frames}, action_chunk_size={action_chunk_size}."
                 )
             num_inference_steps = sp.num_inference_steps or 30
-            guidance_scale = sp.guidance_scale if sp.guidance_scale is not None else 1.0
+            guidance_scale = self._resolve_guidance_scale(sp, 1.0)
             default_flow_shift = 5.0
 
+        if not is_t2i and not action_enabled:
+            requested_num_frames = int(num_frames)
+            num_frames = _ceil_video_num_frames(
+                requested_num_frames,
+                self.vae_scale_factor_temporal,
+            )
+            if num_frames != requested_num_frames and _is_rank_zero():
+                logger.info(
+                    "Rounded Cosmos3 num_frames from %d to %d for temporal compression factor %d.",
+                    requested_num_frames,
+                    num_frames,
+                    self.vae_scale_factor_temporal,
+                )
+
         domain_id = None
         if action_enabled:
             domain_id = resolve_domain_id(
@@ -3010,13 +3373,12 @@ class Cosmos3OmniDiffusersPipeline(
         self._guidance_scale = guidance_scale
         self._num_timesteps = num_inference_steps
 
-        # Always resolve to a concrete target shift for this request, then
-        # update the shared Diffusers scheduler.
-        self._set_flow_shift(flow_shift_target, use_karras_sigmas=False if is_v2v else None)
+        # Always resolve to a concrete target shift for this request.
+        self._set_flow_shift(flow_shift_target)
 
         generator = sp.generator
+        seed = self._resolve_seed(sp, generator)
         if generator is None:
-            seed = sp.seed if sp.seed is not None else 42
             generator = torch.Generator(device=self.device).manual_seed(seed)
 
         # --- Format prompts & tokenize (B=1; reused across loop iterations
@@ -3157,7 +3519,11 @@ class Cosmos3OmniDiffusersPipeline(
             )
 
         def _run_diffusion(start_latents):
-            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
+            self._set_timesteps(
+                num_inference_steps,
+                device=self.device,
+                shift=self._current_flow_shift,
+            )
             scheduler = self.scheduler
             return self.diffuse(
                 latents=start_latents,
@@ -3178,6 +3544,7 @@ class Cosmos3OmniDiffusersPipeline(
                 guidance_interval=guidance_interval,
                 raw_action_dim=raw_action_dim,
                 scheduler=scheduler,
+                generator=generator,
             )
 
         if is_t2i and batch_size > 1:
@@ -3231,12 +3598,18 @@ class Cosmos3OmniDiffusersPipeline(
                 raise ValueError("Cosmos3 action generation finished without action latents.")
             action = action_latents[:, :, :raw_action_dim].detach().cpu()
             return DiffusionOutput(
-                output={"video": video},
-                custom_output={
-                    "action": action,
-                    "raw_action_dim": raw_action_dim,
-                    "action_mode": action_mode,
-                    "domain_id": domain_id,
+                output={
+                    "payload": {
+                        "video": video,
+                        "actions": action,
+                    },
+                    "metadata": {
+                        "actions": {
+                            "raw_action_dim": raw_action_dim,
+                            "action_mode": action_mode,
+                            "domain_id": domain_id,
+                        },
+                    },
                 },
             )
 
diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
index 261573335..de9767d6f 100644
--- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
+++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py
@@ -12,7 +12,9 @@ Ported from the TRT-LLM integration (tekit branch user/shreyasm/cosmos3).
 from __future__ import annotations
 
 import math
-from typing import Any
+from collections.abc import Iterator
+from contextlib import AbstractContextManager, contextmanager, nullcontext
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.distributed as dist
@@ -38,6 +40,9 @@ from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
 from vllm_omni.platforms import current_omni_platform
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+
 logger = init_logger(__name__)
 
 
@@ -391,7 +396,6 @@ class TimestepEmbedder(nn.Module):
         hidden_size: int,
         frequency_embedding_size: int = 256,
         max_period: int = 10000,
-        target_dtype: torch.dtype = torch.bfloat16,
     ) -> None:
         super().__init__()
         # Following diffusers naming pattern here for checkpoint compatibility.
@@ -402,7 +406,7 @@ class TimestepEmbedder(nn.Module):
         self.hidden_size = hidden_size
 
         half = frequency_embedding_size // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half)
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
         self.register_buffer("freqs", freqs, persistent=False)
 
     def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -582,6 +586,7 @@ class Cosmos3CrossAttention(nn.Module):
         head_dim: int,
         rms_norm_eps: float,
         quant_config: QuantizationConfig | None = None,
+        qk_norm: bool = True,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -631,8 +636,10 @@ class Cosmos3CrossAttention(nn.Module):
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
+            self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         self.attn = FrameworkAttention(
             num_heads=self.num_heads,
@@ -716,8 +723,9 @@ class Cosmos3CrossAttention(nn.Module):
         v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim)
 
         # Per-head QK norm
-        q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
-        k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
+        if self.qk_norm:
+            q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
+            k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
 
         # Qwen3-style RoPE
         q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
@@ -801,6 +809,8 @@ class Cosmos3GenDecoderLayer(nn.Module):
         head_dim: int,
         rms_norm_eps: float,
         quant_config: QuantizationConfig | None = None,
+        mlp_cls: type[nn.Module] = Cosmos3GatedMLP,
+        qk_norm: bool = True,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -812,11 +822,12 @@ class Cosmos3GenDecoderLayer(nn.Module):
             head_dim=head_dim,
             rms_norm_eps=rms_norm_eps,
             quant_config=quant_config,
+            qk_norm=qk_norm,
             prefix=f"{prefix}.cross_attention",
         )
         self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.mlp = Cosmos3GatedMLP(
+        self.mlp = mlp_cls(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
             quant_config=quant_config,
@@ -869,6 +880,8 @@ class Cosmos3LanguageModel(nn.Module):
     computed once and reused across all sampling steps.
     """
 
+    _layerwise_offload_blocks_attrs = ["layers"]
+
     def __init__(
         self,
         *,
@@ -959,7 +972,13 @@ class Cosmos3VFMTransformer(nn.Module):
     runs at each denoising step over the target video/image latent stream and
     optional transfer-control, action, and sound latent streams.
 
-    Layerwise offloading uses ``gen_layers`` as the block container.
+    Layerwise offloading uses ``gen_layers`` as the GEN block container.  The
+    nested UND language model declares its own ``layers`` container so the two
+    pathways are offloaded as independent rings.
+
+    Model-level CPU offload is Cosmos3-local: the UND ``reasoner`` and GEN
+    ``generator`` components are swapped at the phase boundaries marked by
+    ``with self._offload_context(...)``.
 
     Sequence parallelism uses ``_sp_plan`` to shard/gather the GEN pathway at
     module boundaries. ``Cosmos3CrossAttention`` checks
@@ -983,6 +1002,9 @@ class Cosmos3VFMTransformer(nn.Module):
 
     packed_modules_mapping = {}
 
+    _language_model_cls = Cosmos3LanguageModel
+    _gen_mlp_cls = Cosmos3GatedMLP
+
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:
         return ("gen_layers" in name or "language_model.layers" in name) and name.split(".")[-1].isdigit()
@@ -1019,6 +1041,25 @@ class Cosmos3VFMTransformer(nn.Module):
             if actual != expected:
                 raise ValueError(f"Unsupported Cosmos3 transformer config: {key}={actual!r}; expected {expected!r}.")
 
+    @classmethod
+    def _resolve_rms_norm_eps(cls, model_config: Any) -> float:
+        return float(_tf_config_get(model_config, "rms_norm_eps", 1e-6))
+
+    @classmethod
+    def _resolve_rope_theta(cls, model_config: Any) -> float:
+        return float(_tf_config_get(model_config, "rope_theta", 5_000_000))
+
+    @classmethod
+    def _resolve_mrope_section(cls, model_config: Any) -> list[int]:
+        rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {}
+        return list(rope_scaling.get("mrope_section", [24, 20, 20]))
+
+    def _language_model_kwargs(self) -> dict[str, Any]:
+        return {}
+
+    def validate_loaded_weights(self, loaded: set[str]) -> None:
+        del loaded
+
     def __init__(
         self,
         od_config: OmniDiffusionConfig,
@@ -1031,7 +1072,6 @@ class Cosmos3VFMTransformer(nn.Module):
         super().__init__()
         model_config = od_config.tf_model_config
         self._validate_supported_config(model_config)
-        rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {}
 
         self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096))
         self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36))
@@ -1040,9 +1080,10 @@ class Cosmos3VFMTransformer(nn.Module):
         self.head_dim = int(_tf_config_get(model_config, "head_dim", 128))
         self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288))
         self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936))
-        self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6))
-        self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000))
-        self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20]))
+        self.rms_norm_eps = self._resolve_rms_norm_eps(model_config)
+        self.rope_theta = self._resolve_rope_theta(model_config)
+        self.mrope_section = self._resolve_mrope_section(model_config)
+        self.qk_norm_for_diffusion = bool(_tf_config_get(model_config, "qk_norm_for_diffusion", True))
         self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2))
         self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48))
         self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001))
@@ -1082,10 +1123,12 @@ class Cosmos3VFMTransformer(nn.Module):
         )
         self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size
 
+        self.use_k_norm_und_for_gen = _tf_config_get(model_config, "use_k_norm_und_for_gen", None)
+
         dtype = od_config.dtype
         quant_config = getattr(od_config, "quantization_config", None) if od_config else None
 
-        self.language_model = Cosmos3LanguageModel(
+        self.language_model = self._language_model_cls(
             hidden_size=self.hidden_size,
             intermediate_size=self.intermediate_size,
             num_hidden_layers=self.num_hidden_layers,
@@ -1098,12 +1141,13 @@ class Cosmos3VFMTransformer(nn.Module):
             mrope_section=self.mrope_section,
             quant_config=quant_config,
             prefix="language_model",
+            **self._language_model_kwargs(),
         )
 
         # Video projection layers are small; not worth quantizing.
         self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size)
         self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim)
-        self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=dtype)
+        self.time_embedder = TimestepEmbedder(self.hidden_size)
         if self.action_gen:
             self.action_proj_in = DomainAwareLinear(
                 self.action_dim,
@@ -1134,6 +1178,8 @@ class Cosmos3VFMTransformer(nn.Module):
                     head_dim=self.head_dim,
                     rms_norm_eps=self.rms_norm_eps,
                     quant_config=quant_config,
+                    mlp_cls=self._gen_mlp_cls,
+                    qk_norm=self.qk_norm_for_diffusion,
                     prefix=f"gen_layers.{i}",
                 )
                 for i in range(self.num_hidden_layers)
@@ -1148,10 +1194,125 @@ class Cosmos3VFMTransformer(nn.Module):
         self.cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None
         self.cached_freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None
 
+        self._model_cpu_offload_enabled = False
+        self._model_cpu_offload_device: torch.device | None = None
+        self._model_cpu_offload_mover: SequentialOffloadHook | None = None
+        self._active_model_cpu_offload_component: str | None = None
+
     @property
     def device(self) -> torch.device:
+        offload_device = getattr(self, "_model_cpu_offload_device", None)
+        if getattr(self, "_model_cpu_offload_enabled", False) and offload_device is not None:
+            return offload_device
         return next(self.parameters()).device
 
+    def _model_cpu_offload_components(self) -> dict[str, list[nn.Module]]:
+        """Cosmos3's mutually-exclusive reasoner/generator component sets."""
+        return {
+            "reasoner": [self.language_model.layers],
+            "generator": [self.gen_layers],
+        }
+
+    def _model_cpu_offload_component_tensor_ids(self) -> set[int]:
+        component_tensors: set[int] = set()
+        for modules in self._model_cpu_offload_components().values():
+            for module in modules:
+                component_tensors.update(id(param) for param in module.parameters())
+                component_tensors.update(id(buffer) for buffer in module.buffers())
+        return component_tensors
+
+    def _move_model_cpu_offload_residents(self) -> None:
+        """Keep non-reasoner/non-generator weights resident on the target device."""
+        component_tensors = self._model_cpu_offload_component_tensor_ids()
+        device = self._model_cpu_offload_device
+        if device is None:
+            return
+        with torch.no_grad():
+            for param in self.parameters():
+                if id(param) not in component_tensors and param.data.device != device:
+                    param.data = param.data.to(device, non_blocking=False)
+            for buffer in self.buffers():
+                if id(buffer) not in component_tensors and buffer.device != device:
+                    buffer.data = buffer.data.to(device, non_blocking=False)
+
+    def _offload_model_cpu_component(self, name: str) -> None:
+        mover = self._model_cpu_offload_mover
+        if mover is None:
+            raise RuntimeError("Cosmos3 model CPU offload is not enabled")
+        for module in self._model_cpu_offload_components()[name]:
+            mover._to_cpu(module)
+
+    def _load_model_cpu_component(self, name: str) -> None:
+        mover = self._model_cpu_offload_mover
+        if mover is None:
+            raise RuntimeError("Cosmos3 model CPU offload is not enabled")
+        for module in self._model_cpu_offload_components()[name]:
+            mover._to_gpu(module)
+
+    def enable_model_cpu_offload(
+        self,
+        *,
+        device: torch.device,
+        pin_memory: bool = True,
+        use_hsdp: bool = False,
+    ) -> None:
+        """Enable Cosmos3 reasoner/generator CPU swapping inside ``forward``."""
+        if getattr(self, "_model_cpu_offload_enabled", False):
+            return
+
+        from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+
+        self._model_cpu_offload_device = torch.device(device)
+        self._model_cpu_offload_mover = SequentialOffloadHook(
+            offload_targets=[],
+            device=self._model_cpu_offload_device,
+            pin_memory=pin_memory,
+            use_hsdp=use_hsdp,
+        )
+        self._move_model_cpu_offload_residents()
+        for name in self._model_cpu_offload_components():
+            self._offload_model_cpu_component(name)
+        self._model_cpu_offload_enabled = True
+        self._active_model_cpu_offload_component = None
+        logger.info("Cosmos3 component-level CPU offload enabled on %s", self._model_cpu_offload_device)
+
+    def disable_model_cpu_offload(self) -> None:
+        if not getattr(self, "_model_cpu_offload_enabled", False):
+            return
+        for name in self._model_cpu_offload_components():
+            self._load_model_cpu_component(name)
+        if self._model_cpu_offload_device is not None and self._model_cpu_offload_device.type != "cpu":
+            current_omni_platform.synchronize()
+        self._model_cpu_offload_enabled = False
+        self._model_cpu_offload_device = None
+        self._model_cpu_offload_mover = None
+        self._active_model_cpu_offload_component = None
+
+    def _activate_model_cpu_offload_component(self, name: str) -> None:
+        if not getattr(self, "_model_cpu_offload_enabled", False):
+            return
+        components = self._model_cpu_offload_components()
+        if name not in components:
+            raise ValueError(f"Unknown Cosmos3 offload component: {name!r} (known: {list(components)})")
+        if self._active_model_cpu_offload_component == name:
+            return
+        self._active_model_cpu_offload_component = None
+        for other in components:
+            if other != name:
+                self._offload_model_cpu_component(other)
+        self._load_model_cpu_component(name)
+        self._active_model_cpu_offload_component = name
+
+    @contextmanager
+    def _model_cpu_offload_context(self, name: str) -> Iterator[None]:
+        self._activate_model_cpu_offload_component(name)
+        yield
+
+    def _offload_context(self, name: str) -> AbstractContextManager[None]:
+        if not getattr(self, "_model_cpu_offload_enabled", False):
+            return nullcontext()
+        return self._model_cpu_offload_context(name)
+
     # -- Patchify / Unpatchify -----------------------------------------------
 
     def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]:
@@ -1483,26 +1644,14 @@ class Cosmos3VFMTransformer(nn.Module):
         # Query Ulysses state at runtime
         ulysses_size, _, _ = _get_ulysses_state()
 
-        # Patchify latents and project to hidden space
-        hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w))
-        s_video = hidden_video.shape[1]
-        s_control = 0
-        hidden_controls: list[torch.Tensor] = []
-        for idx, control in enumerate(control_latent_list):
-            if control.shape != hidden_states.shape:
-                raise ValueError(
-                    "Cosmos3 transfer control latent shape must match target latent shape: "
-                    f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}."
-                )
-            hidden_control = self.proj_in(
-                self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w)
-            )
-            hidden_controls.append(hidden_control)
-            s_control += hidden_control.shape[1]
+        # Pack action/sound tokens (no learned weights) up front so the UND
+        # cache sizing knows their token lengths.  The modality projections are
+        # deferred into the generator offload context below, so model-level
+        # offload does not stage GEN weights before swapping to the reasoner.
+        action_tokens = None
+        sound_tokens = None
         s_action = 0
-        hidden_action = None
         s_sound = 0
-        hidden_sound = None
         if action_latents is not None:
             if action_latents.shape[0] != hidden_states.shape[0]:
                 raise ValueError(
@@ -1511,60 +1660,16 @@ class Cosmos3VFMTransformer(nn.Module):
                 )
             if action_domain_ids is None:
                 action_domain_ids = torch.zeros(action_latents.shape[0], dtype=torch.long, device=action_latents.device)
-            hidden_action = self.action_proj_in(self.pack_action(action_latents), action_domain_ids)
-            hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype)
-            s_action = hidden_action.shape[1]
+            action_tokens = self.pack_action(action_latents)
+            s_action = action_tokens.shape[1]
         if sound_latents is not None:
             if sound_latents.shape[0] != hidden_states.shape[0]:
                 raise ValueError(
                     "Cosmos3 sound and video batch sizes must match: "
                     f"video={hidden_states.shape[0]}, sound={sound_latents.shape[0]}."
                 )
-            hidden_sound = self.audio_proj_in(self.pack_sound(sound_latents))
-            hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype)
-            s_sound = hidden_sound.shape[1]
-
-        # Timestep embedding (fp32 for precision).
-        # For I2V: only add to noisy tokens, not conditioned ones.
-        # Conditioned frames are clean context and should not receive
-        # the diffusion timestep signal.
-        with torch.autocast(current_omni_platform.device_type, enabled=False):
-            time_embed = self.time_embedder((timestep * self.timestep_scale).float())
-        time_embed = time_embed.to(hidden_states.dtype)
-
-        if noisy_frame_mask is not None:
-            # Build per-token mask from per-frame mask.
-            # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1]
-            token_noisy_mask = (
-                noisy_frame_mask[:, 0, :, 0, 0]  # [B, t]
-                .unsqueeze(-1)  # [B, t, 1]
-                .expand(-1, -1, hp * wp)  # [B, t, hp*wp]
-                .reshape(hidden_video.shape[0], -1, 1)  # [B, t*hp*wp, 1]
-            )
-            hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask
-        else:
-            hidden_video = hidden_video + time_embed.unsqueeze(1)
-
-        if hidden_action is not None:
-            if action_noisy_mask is None:
-                hidden_action = hidden_action + time_embed.unsqueeze(1)
-            else:
-                if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1):
-                    raise ValueError(
-                        "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], "
-                        f"got {tuple(action_noisy_mask.shape)}."
-                    )
-                action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device)
-                hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask
-
-        if hidden_sound is not None:
-            hidden_sound = hidden_sound + time_embed.unsqueeze(1)
-        hidden_parts = [*hidden_controls, hidden_video]
-        if hidden_action is not None:
-            hidden_parts.append(hidden_action)
-        if hidden_sound is not None:
-            hidden_parts.append(hidden_sound)
-        hidden_gen = torch.cat(hidden_parts, dim=1)
+            sound_tokens = self.pack_sound(sound_latents)
+            s_sound = sound_tokens.shape[1]
 
         # Run UND pathway once and cache K/V (replicated across all ranks)
         if self.cached_kv is None:
@@ -1583,90 +1688,162 @@ class Cosmos3VFMTransformer(nn.Module):
                 num_vision_items=len(control_latent_list) + 1,
                 share_vision_temporal_positions=transfer_share_vision_temporal_positions,
             )
-            cached_kv_full = self.language_model(text_ids, freqs_und)
+            with self._offload_context("reasoner"):
+                cached_kv_full = self.language_model(text_ids, freqs_und)
             self.cached_freqs_gen = freqs_gen
 
             # Trim to real text length (remove padding).  K/V stay replicated;
             # the framework Attention layer head-slices them via joint_key/value.
             self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full]
 
-        # Run GEN layers.  UND K/V (replicated) is passed to each layer;
-        # the Cosmos3CrossAttention forwards them as joint_key/value so the
-        # framework Attention handles the Ulysses head-slicing internally.
-        if self.cached_kv is None or self.cached_freqs_gen is None:
-            raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
-        self._validate_gen_sequence_parallel(
-            s_gen=hidden_gen.shape[1],
-            s_video=s_video,
-            s_control=s_control,
-            s_action=s_action,
-            s_sound=s_sound,
-            has_action=has_action,
-            has_sound=has_sound,
-            has_control=has_control,
-            ulysses_size=ulysses_size,
-        )
-        freqs_cos, freqs_sin = self.cached_freqs_gen
-        hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)
-        freqs_gen = (freqs_cos, freqs_sin)
-
-        if len(self.gen_layers) == len(self.cached_kv):
-            for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True):
-                hidden_gen = layer(
-                    hidden_gen,
-                    k_und=k_und,
-                    v_und=v_und,
-                    freqs_cos=freqs_cos,
-                    freqs_sin=freqs_sin,
+        with self._offload_context("generator"):
+            # Patchify latents and project to hidden space after UND cache
+            # construction, so model-level offload does not stage GEN weights
+            # before immediately swapping to the reasoner.
+            hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w))
+            s_video = hidden_video.shape[1]
+            s_control = 0
+            hidden_controls: list[torch.Tensor] = []
+            for idx, control in enumerate(control_latent_list):
+                if control.shape != hidden_states.shape:
+                    raise ValueError(
+                        "Cosmos3 transfer control latent shape must match target latent shape: "
+                        f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}."
+                    )
+                hidden_control = self.proj_in(
+                    self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w)
                 )
-                # Cache-dit's block wrapper may return a tuple; unwrap it.
-                if isinstance(hidden_gen, tuple):
-                    hidden_gen = hidden_gen[0]
-        else:
-            # Cache-dit patches gen_layers to a grouped wrapper.
-            for layer in self.gen_layers:
-                hidden_gen = layer(
-                    hidden_gen,
-                    cached_kv=self.cached_kv,
-                    freqs_gen=freqs_gen,
+                hidden_controls.append(hidden_control)
+                s_control += hidden_control.shape[1]
+            hidden_action = None
+            hidden_sound = None
+            if action_tokens is not None:
+                assert action_domain_ids is not None
+                hidden_action = self.action_proj_in(action_tokens, action_domain_ids)
+                hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype)
+            if sound_tokens is not None:
+                hidden_sound = self.audio_proj_in(sound_tokens)
+                hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype)
+
+            # Timestep embedding (fp32 for precision).
+            # For I2V: only add to noisy tokens, not conditioned ones.
+            # Conditioned frames are clean context and should not receive
+            # the diffusion timestep signal.
+            with torch.autocast(current_omni_platform.device_type, enabled=False):
+                time_embed = self.time_embedder((timestep * self.timestep_scale).float())
+            time_embed = time_embed.to(hidden_states.dtype)
+
+            if noisy_frame_mask is not None:
+                # Build per-token mask from per-frame mask.
+                # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1]
+                token_noisy_mask = (
+                    noisy_frame_mask[:, 0, :, 0, 0]  # [B, t]
+                    .unsqueeze(-1)  # [B, t, 1]
+                    .expand(-1, -1, hp * wp)  # [B, t, hp*wp]
+                    .reshape(hidden_video.shape[0], -1, 1)  # [B, t*hp*wp, 1]
                 )
-                if isinstance(hidden_gen, tuple):
-                    hidden_gen = hidden_gen[0]
+                hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask
+            else:
+                hidden_video = hidden_video + time_embed.unsqueeze(1)
 
-        hidden_gen = self.gen_sp_gather(hidden_gen)
+            if hidden_action is not None:
+                if action_noisy_mask is None:
+                    hidden_action = hidden_action + time_embed.unsqueeze(1)
+                else:
+                    if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1):
+                        raise ValueError(
+                            "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], "
+                            f"got {tuple(action_noisy_mask.shape)}."
+                        )
+                    action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device)
+                    hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask
 
-        # Final norm and project back to latent space
-        hidden_gen = self.norm_moe_gen(hidden_gen)
-        if not has_action and not has_sound and not has_control:
-            return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
+            if hidden_sound is not None:
+                hidden_sound = hidden_sound + time_embed.unsqueeze(1)
+            hidden_parts = [*hidden_controls, hidden_video]
+            if hidden_action is not None:
+                hidden_parts.append(hidden_action)
+            if hidden_sound is not None:
+                hidden_parts.append(hidden_sound)
+            hidden_gen = torch.cat(hidden_parts, dim=1)
 
-        split_sizes = []
-        if has_control:
-            split_sizes.append(s_control)
-        split_sizes.append(s_video)
-        if has_action:
-            split_sizes.append(s_action)
-        if has_sound:
-            split_sizes.append(s_sound)
-        split_hidden = hidden_gen.split(split_sizes, dim=1)
-        split_idx = 0
-        if has_control:
+            # Run GEN layers.  UND K/V (replicated) is passed to each layer;
+            # the Cosmos3CrossAttention forwards them as joint_key/value so the
+            # framework Attention handles the Ulysses head-slicing internally.
+            if self.cached_kv is None or self.cached_freqs_gen is None:
+                raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
+            self._validate_gen_sequence_parallel(
+                s_gen=hidden_gen.shape[1],
+                s_video=s_video,
+                s_control=s_control,
+                s_action=s_action,
+                s_sound=s_sound,
+                has_action=has_action,
+                has_sound=has_sound,
+                has_control=has_control,
+                ulysses_size=ulysses_size,
+            )
+            freqs_cos, freqs_sin = self.cached_freqs_gen
+            hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)
+            freqs_gen = (freqs_cos, freqs_sin)
+
+            if len(self.gen_layers) == len(self.cached_kv):
+                for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True):
+                    hidden_gen = layer(
+                        hidden_gen,
+                        k_und=k_und,
+                        v_und=v_und,
+                        freqs_cos=freqs_cos,
+                        freqs_sin=freqs_sin,
+                    )
+                    # Cache-dit's block wrapper may return a tuple; unwrap it.
+                    if isinstance(hidden_gen, tuple):
+                        hidden_gen = hidden_gen[0]
+            else:
+                # Cache-dit patches gen_layers to a grouped wrapper.
+                for layer in self.gen_layers:
+                    hidden_gen = layer(
+                        hidden_gen,
+                        cached_kv=self.cached_kv,
+                        freqs_gen=freqs_gen,
+                    )
+                    if isinstance(hidden_gen, tuple):
+                        hidden_gen = hidden_gen[0]
+
+            hidden_gen = self.gen_sp_gather(hidden_gen)
+
+            # Final norm and project back to latent space
+            hidden_gen = self.norm_moe_gen(hidden_gen)
+            if not has_action and not has_sound and not has_control:
+                return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
+
+            split_sizes = []
+            if has_control:
+                split_sizes.append(s_control)
+            split_sizes.append(s_video)
+            if has_action:
+                split_sizes.append(s_action)
+            if has_sound:
+                split_sizes.append(s_sound)
+            split_hidden = hidden_gen.split(split_sizes, dim=1)
+            split_idx = 0
+            if has_control:
+                split_idx += 1
+            hidden_video = split_hidden[split_idx]
             split_idx += 1
-        hidden_video = split_hidden[split_idx]
-        split_idx += 1
-        video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w)
-        if has_control:
-            return video_pred
-        outputs: list[torch.Tensor] = [video_pred]
-        if has_action:
-            hidden_action = split_hidden[split_idx]
-            split_idx += 1
-            assert action_domain_ids is not None
-            outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids)))
-        if has_sound:
-            hidden_sound = split_hidden[split_idx]
-            outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))
-        return tuple(outputs)
+            video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w)
+            if has_control:
+                return video_pred
+            outputs: list[torch.Tensor] = [video_pred]
+            if has_action:
+                hidden_action = split_hidden[split_idx]
+                split_idx += 1
+                assert action_domain_ids is not None
+                outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids)))
+            if has_sound:
+                hidden_sound = split_hidden[split_idx]
+                outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))
+            return tuple(outputs)
 
     def post_load_weights(self) -> None:
         """Post-load processing: ensure correct dtypes."""
diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py
new file mode 100644
index 000000000..6a1b4313b
--- /dev/null
+++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py
@@ -0,0 +1,396 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cosmos3 Edge transformer variant with a Nemotron dense UND backbone."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
+
+from .transformer_cosmos3 import (
+    Cosmos3VFMTransformer,
+    Qwen3VLTextRotaryEmbedding,
+    RMSNorm,
+    _apply_rotary_pos_emb,
+    _nested_get,
+    _tf_config_get,
+)
+
+COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense"
+
+
+def _tf_config_nested_get(config: Any, key: str, default: Any = None) -> Any:
+    if config is None:
+        return default
+    if isinstance(config, dict):
+        value = _nested_get(config, key)
+    elif hasattr(config, "to_dict"):
+        value = _nested_get(config.to_dict(), key)
+    elif hasattr(config, "params"):
+        value = _nested_get(config.params, key)
+    else:
+        value = getattr(config, key, None)
+    return default if value is None else value
+
+
+class Cosmos3Relu2MLP(nn.Module):
+    """Nemotron dense MLP: down_proj(relu(up_proj(x)) ** 2)."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.up_proj = ColumnParallelLinear(
+            hidden_size,
+            intermediate_size,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        hidden = F.relu(self.up_proj(x))
+        return self.down_proj(hidden * hidden)
+
+
+class Cosmos3EdgeCausalAttention(nn.Module):
+    """Edge UND causal attention.
+
+    Reasoner self-attention uses raw UND Q/K with RoPE because
+    ``qk_norm_for_text=false``.  The GEN-facing UND key cache is a separate
+    normalized-then-RoPE'd view of the raw K tensor when GEN QK norm is enabled;
+    no-GEN-QK-norm checkpoints reuse the raw RoPE'd UND K.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        rms_norm_eps: float,
+        use_k_norm_und_for_gen: bool,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_attention_heads
+        self.num_kv_heads = num_key_value_heads
+        self.head_dim = head_dim
+
+        tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads_local = self.num_heads // tp_size
+        self.num_kv_heads_local = self.num_kv_heads // tp_size
+
+        self.to_q = ColumnParallelLinear(
+            hidden_size,
+            self.num_heads * self.head_dim,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_q",
+        )
+        self.to_k = ColumnParallelLinear(
+            hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_k",
+        )
+        self.to_v = ColumnParallelLinear(
+            hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_v",
+        )
+        self.to_out = RowParallelLinear(
+            self.num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_out",
+        )
+
+        self.k_norm_und_for_gen = RMSNorm(self.head_dim, eps=rms_norm_eps) if use_k_norm_und_for_gen else None
+        self.attn = FrameworkAttention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            causal=True,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            num_kv_heads=self.num_kv_heads,
+            skip_sequence_parallel=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        freqs_cos: torch.Tensor,
+        freqs_sin: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        B, S, _ = hidden_states.shape
+
+        q = self.to_q(hidden_states).view(B, S, self.num_heads_local, self.head_dim)
+        k = self.to_k(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
+        v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
+
+        q_reasoner, k_reasoner = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
+
+        if self.k_norm_und_for_gen is not None:
+            k_gen = F.rms_norm(
+                k,
+                (self.head_dim,),
+                self.k_norm_und_for_gen.weight,
+                eps=self.k_norm_und_for_gen.variance_epsilon,
+            )
+            _, k_gen = _apply_rotary_pos_emb(q, k_gen, freqs_cos, freqs_sin)
+        else:
+            k_gen = k_reasoner
+
+        out = self.attn(q_reasoner, k_reasoner, v).reshape(B, S, -1)
+        return self.to_out(out), k_gen, v
+
+
+class Cosmos3EdgeUndDecoderLayer(nn.Module):
+    """Edge UND decoder layer: causal self-attention + ReLU2 MLP."""
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        intermediate_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        rms_norm_eps: float,
+        use_k_norm_und_for_gen: bool,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.self_attn = Cosmos3EdgeCausalAttention(
+            hidden_size=hidden_size,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            rms_norm_eps=rms_norm_eps,
+            use_k_norm_und_for_gen=use_k_norm_und_for_gen,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.mlp = Cosmos3Relu2MLP(
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        freqs: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        cos, sin = freqs
+        attn_out, k, v = self.self_attn(hidden_states, cos, sin)
+        hidden_states = residual + attn_out
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + self.mlp(hidden_states)
+
+        return hidden_states, k, v
+
+
+class Cosmos3EdgeLanguageModel(nn.Module):
+    """Nemotron dense UND tower that returns GEN-facing K plus raw V per layer."""
+
+    _layerwise_offload_blocks_attrs = ["layers"]
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        intermediate_size: int,
+        num_hidden_layers: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        vocab_size: int,
+        rms_norm_eps: float,
+        rope_theta: float,
+        mrope_section: list[int],
+        use_k_norm_und_for_gen: bool,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
+        self.rotary_emb = Qwen3VLTextRotaryEmbedding(
+            head_dim=head_dim,
+            rope_theta=rope_theta,
+            mrope_section=mrope_section,
+        )
+        self.layers = nn.ModuleList(
+            [
+                Cosmos3EdgeUndDecoderLayer(
+                    hidden_size=hidden_size,
+                    intermediate_size=intermediate_size,
+                    num_attention_heads=num_attention_heads,
+                    num_key_value_heads=num_key_value_heads,
+                    head_dim=head_dim,
+                    rms_norm_eps=rms_norm_eps,
+                    use_k_norm_und_for_gen=use_k_norm_und_for_gen,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{i}",
+                )
+                for i in range(num_hidden_layers)
+            ]
+        )
+        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
+
+    def forward(
+        self,
+        text_ids: torch.Tensor,
+        freqs: tuple[torch.Tensor, torch.Tensor],
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        hidden = self.embed_tokens(text_ids)
+
+        cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for layer in self.layers:
+            hidden, k, v = layer(hidden, freqs)
+            cached_kv.append((k, v))
+
+        return cached_kv
+
+
+class Cosmos3EdgeVFMTransformer(Cosmos3VFMTransformer):
+    """Cosmos3 Edge variant with Nemotron dense UND and shared GEN diffusion."""
+
+    _language_model_cls = Cosmos3EdgeLanguageModel
+    _gen_mlp_cls = Cosmos3Relu2MLP
+
+    @staticmethod
+    def _validate_supported_config(model_config: Any) -> None:
+        expected_values = {
+            "qk_norm_for_text": False,
+            "position_embedding_type": "unified_3d_mrope",
+            "unified_3d_mrope_reset_spatial_ids": True,
+            "joint_attn_implementation": "two_way",
+        }
+        for key, expected in expected_values.items():
+            actual = _tf_config_get(model_config, key, expected)
+            if actual != expected:
+                raise ValueError(
+                    f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}."
+                )
+
+        backbone_type = _tf_config_get(model_config, "backbone_type", None)
+        if backbone_type is None:
+            raise ValueError(
+                f"Cosmos3 Edge transformer config must declare backbone_type={COSMOS3_EDGE_BACKBONE_TYPE!r}."
+            )
+        if backbone_type != COSMOS3_EDGE_BACKBONE_TYPE:
+            raise ValueError(
+                "Unsupported Cosmos3 Edge transformer config: "
+                f"backbone_type={backbone_type!r}; expected {COSMOS3_EDGE_BACKBONE_TYPE!r}."
+            )
+
+        required_values = {
+            "latent_channel": 48,
+            "latent_patch_size": 2,
+            "temporal_compression_factor": 4,
+        }
+        for key, expected in required_values.items():
+            actual = _tf_config_get(model_config, key, None)
+            if actual is None:
+                raise ValueError(f"Cosmos3 Edge transformer config must declare {key}={expected}.")
+            if int(actual) != expected:
+                raise ValueError(
+                    f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}."
+                )
+
+    @classmethod
+    def _resolve_rms_norm_eps(cls, model_config: Any) -> float:
+        return float(
+            _tf_config_get(
+                model_config,
+                "rms_norm_eps",
+                _tf_config_get(model_config, "layer_norm_epsilon", 1e-5),
+            )
+        )
+
+    @classmethod
+    def _resolve_rope_theta(cls, model_config: Any) -> float:
+        return float(_tf_config_get(model_config, "rope_theta", 100_000_000))
+
+    @classmethod
+    def _resolve_mrope_section(cls, model_config: Any) -> list[int]:
+        mrope_section = _tf_config_get(model_config, "mrope_section", None)
+        if mrope_section is None:
+            mrope_section = _tf_config_nested_get(model_config, "mrope_section", [24, 20, 20])
+        return list(mrope_section)
+
+    def _language_model_kwargs(self) -> dict[str, Any]:
+        return {"use_k_norm_und_for_gen": self.use_k_norm_und_for_gen}
+
+    def validate_loaded_weights(self, loaded: set[str]) -> None:
+        missing: list[str] = []
+        for layer_idx in range(self.num_hidden_layers):
+            required_markers = (
+                f"language_model.layers.{layer_idx}.mlp.up_proj.",
+                f"language_model.layers.{layer_idx}.mlp.down_proj.",
+                f"gen_layers.{layer_idx}.mlp.up_proj.",
+                f"gen_layers.{layer_idx}.mlp.down_proj.",
+            )
+            if self.use_k_norm_und_for_gen:
+                required_markers = (
+                    *required_markers,
+                    f"language_model.layers.{layer_idx}.self_attn.k_norm_und_for_gen.",
+                )
+            missing.extend(
+                marker.rstrip(".") for marker in required_markers if not any(marker in name for name in loaded)
+            )
+        if missing:
+            preview = ", ".join(missing[:10])
+            suffix = " ..." if len(missing) > 10 else ""
+            raise ValueError(
+                "Cosmos3 Edge transformer checkpoint is missing required weights for "
+                f"{preview}{suffix}. Use a converted Edge checkpoint with ReLU2 UND/GEN MLP weights."
+            )
diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py
index 4d48053a6..94fcb6699 100644
--- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py
+++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py
@@ -24,6 +24,11 @@ from torch import nn
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import BasePipelineUtils, get_pipeline_utils
+from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import (
+    apply_diffusers_quantization_config,
+    convert_diffusers_quantization_config,
+    ensure_supported_diffusers_quantization,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniPromptType, OmniTextPrompt
@@ -31,6 +36,20 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
+_DIFFUSERS_CONFIG_LOAD_KWARGS = {
+    "cache_dir",
+    "dduf_entries",
+    "force_download",
+    "local_dir",
+    "local_dir_use_symlinks",
+    "local_files_only",
+    "proxies",
+    "revision",
+    "subfolder",
+    "token",
+    "user_agent",
+}
+
 
 class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
     """Black-box adapter that delegates full pipeline execution to a diffusers pipeline.
@@ -80,12 +99,19 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             "torch_dtype": dtype,
             **self.od_config.diffusers_load_kwargs,
         }
-        logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
+        convert_diffusers_quantization_config(load_kwargs)
 
         pipeline_class = self.od_config.diffusers_pipeline_cls
         pipeline_class_name = pipeline_class.__name__ if pipeline_class is not None else None
         self._pipeline_utils = get_pipeline_utils(pipeline_class_name)
         self._pipeline_utils.update_load_kwargs(self.od_config, load_kwargs)
+        component_names = (
+            self._load_diffusers_component_names(model_id, load_kwargs)
+            if self.od_config.quantization_config is not None and "quantization_config" not in load_kwargs
+            else {}
+        )
+        apply_diffusers_quantization_config(self.od_config, load_kwargs, component_names)
+        logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
 
         self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs)
         self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config)
@@ -190,10 +216,38 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                 "Eager execution is not supported with the diffusers backend. "
                 "Use a native pipeline for continuous batching mode."
             )
-        if self.od_config.quantization_config is not None:
-            raise NotImplementedError(
-                "Quantization is not supported with the diffusers backend. Use a native pipeline for quantization."
+        if (
+            self.od_config.quantization_config is not None
+            and "quantization_config" not in self.od_config.diffusers_load_kwargs
+        ):
+            ensure_supported_diffusers_quantization(self.od_config.quantization_config)
+            if self.od_config.diffusers_load_kwargs.get("dduf_file"):
+                raise NotImplementedError(
+                    "Diffusers backend quantization conversion does not support "
+                    "diffusers_load_kwargs.dduf_file yet. The preflight component "
+                    "discovery would need to mirror Diffusers' DDUF config loading. "
+                    "Use diffusers_load_kwargs.quantization_config for a native "
+                    "Diffusers quantization config, or omit dduf_file for vLLM-Omni "
+                    "quantization conversion."
+                )
+
+    def _load_diffusers_component_names(
+        self,
+        model_id: str,
+        load_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        config_load_kwargs = {k: load_kwargs[k] for k in _DIFFUSERS_CONFIG_LOAD_KWARGS if k in load_kwargs}
+        pipeline_config = DiffusionPipeline.load_config(model_id, **config_load_kwargs)
+        return {
+            name: value
+            for name, value in pipeline_config.items()
+            if (
+                isinstance(value, list)
+                and len(value) > 0
+                and value[0] is not None
+                and (name not in load_kwargs or load_kwargs[name] is not None)
             )
+        }
 
     # ------------------------------------------------------------------
     # Wrap settings, inputs, and outputs
diff --git a/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py b/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py
new file mode 100644
index 000000000..a7228d21e
--- /dev/null
+++ b/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py
@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantization helpers for the Diffusers backend."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from diffusers import PipelineQuantizationConfig, TorchAoConfig
+
+logger = logging.getLogger(__name__)
+
+_DIFFUSERS_TRANSFORMER_COMPONENTS = ("transformer", "transformer_2")
+_QuantizationConfigSpec = tuple[Callable[[Any], None], str]
+
+
+def ensure_supported_diffusers_quantization(quant_config: Any) -> None:
+    """Validate that a vLLM-Omni quantization config has a Diffusers mapping."""
+
+    method = _get_quant_method_name(quant_config)
+    validator, _ = _get_quantization_config_spec(method)
+    validator(quant_config)
+
+
+def convert_diffusers_quantization_config(load_kwargs: dict[str, Any]) -> None:
+    """Build Diffusers-native quantization config objects from dict kwargs."""
+
+    quant_config = load_kwargs.get("quantization_config")
+    if isinstance(quant_config, dict):
+        load_kwargs["quantization_config"] = _build_diffusers_native_quantization_config(quant_config)
+
+
+def apply_diffusers_quantization_config(
+    od_config: Any,
+    load_kwargs: dict[str, Any],
+    component_names: dict[str, Any],
+) -> bool:
+    """Inject a courtesy-converted quantization_config into load kwargs.
+
+    ``diffusers_load_kwargs`` is the canonical Diffusers backend configuration
+    path, so an explicit ``quantization_config`` already present in ``load_kwargs``
+    is never replaced. Returns whether a config was injected by this helper.
+    """
+
+    quant_config = getattr(od_config, "quantization_config", None)
+    if quant_config is None:
+        return False
+
+    if "quantization_config" in load_kwargs:
+        logger.warning(
+            "Both vLLM-Omni quantization_config and diffusers_load_kwargs.quantization_config "
+            "were provided for the diffusers backend. Using the Diffusers-native "
+            "quantization_config from diffusers_load_kwargs."
+        )
+        return False
+
+    transformer_components = _get_diffusers_transformer_components(component_names)
+    load_kwargs["quantization_config"] = _build_diffusers_quantization_config(
+        quant_config,
+        transformer_components,
+    )
+    return True
+
+
+def _normalize_method_name(method: Any) -> str:
+    method = getattr(method, "value", method)
+    return str(method).lower().replace("-", "_")
+
+
+def _get_quant_method_name(quant_config: Any) -> str:
+    get_name = getattr(quant_config, "get_name", None)
+    if get_name is not None:
+        return _normalize_method_name(get_name() if callable(get_name) else get_name)
+
+    method = getattr(quant_config, "quant_method", None)
+    if method is None:
+        method = getattr(quant_config, "method", None)
+    if method is None:
+        raise NotImplementedError(
+            "Diffusers backend quantization conversion requires a quantization "
+            "config with get_name(), quant_method, or method."
+        )
+    return _normalize_method_name(method)
+
+
+def _ensure_no_ignored_layers(quant_config: Any) -> None:
+    ignored_layers = getattr(quant_config, "ignored_layers", None)
+    if not ignored_layers:
+        ignored_layers = getattr(quant_config, "modules_to_not_convert", None)
+    if ignored_layers:
+        raise NotImplementedError(
+            "Diffusers backend quantization conversion does not map vLLM "
+            "ignored_layers/modules_to_not_convert names to Diffusers module "
+            "names. Use diffusers_load_kwargs for a native Diffusers config."
+        )
+
+
+def _get_torchao_quant_type_cls(class_name: str) -> type[Any]:
+    try:
+        import torchao.quantization as torchao_quantization
+    except ImportError as exc:
+        raise ImportError(
+            "Diffusers backend quantization conversion for fp8/int8 requires "
+            "torchao. Install torchao or pass a Diffusers-native "
+            "quantization_config through diffusers_load_kwargs."
+        ) from exc
+
+    try:
+        return getattr(torchao_quantization, class_name)
+    except AttributeError as exc:
+        raise ImportError(f"torchao.quantization.{class_name} is required for this quantization mapping.") from exc
+
+
+def _build_diffusers_native_quantization_config(config: dict[str, Any]) -> Any:
+    config = config.copy()
+    quant_mapping = config.get("quant_mapping")
+    if isinstance(quant_mapping, dict):
+        config["quant_mapping"] = {
+            component_name: _build_diffusers_component_quant_config(component_config)
+            for component_name, component_config in quant_mapping.items()
+        }
+    return PipelineQuantizationConfig(**config)
+
+
+def _build_diffusers_component_quant_config(config: Any) -> Any:
+    if not isinstance(config, dict):
+        return config
+
+    quant_method = config.get("quant_method")
+    if quant_method is None:
+        raise ValueError("Diffusers quant_mapping entries must provide quant_method.")
+
+    quant_method = _normalize_method_name(quant_method)
+    config_cls = _get_diffusers_quantization_config_cls(quant_method)
+    from_dict = getattr(config_cls, "from_dict", None)
+    if from_dict is not None:
+        return from_dict(config)
+
+    init_kwargs = config.copy()
+    init_kwargs.pop("quant_method", None)
+    return config_cls(**init_kwargs)
+
+
+def _get_diffusers_quantization_config_cls(quant_method: str) -> type[Any]:
+    from diffusers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING as DIFFUSERS_QUANT_CONFIG_MAPPING
+
+    config_cls = DIFFUSERS_QUANT_CONFIG_MAPPING.get(quant_method)
+    if config_cls is not None:
+        return config_cls
+
+    try:
+        from transformers.quantizers.auto import (
+            AUTO_QUANTIZATION_CONFIG_MAPPING as TRANSFORMERS_QUANT_CONFIG_MAPPING,
+        )
+    except ImportError:
+        TRANSFORMERS_QUANT_CONFIG_MAPPING = {}
+
+    config_cls = TRANSFORMERS_QUANT_CONFIG_MAPPING.get(quant_method)
+    if config_cls is None:
+        raise ValueError(f"Diffusers quant_mapping entry uses unknown quant_method={quant_method!r}.")
+    return config_cls
+
+
+def _build_torchao_pipeline_quant_config(
+    torchao_quant_type_name: str,
+    transformer_components: list[str],
+) -> Any:
+    """Build a Diffusers config after the caller has validated the vLLM config."""
+
+    quant_type_cls = _get_torchao_quant_type_cls(torchao_quant_type_name)
+    return PipelineQuantizationConfig(
+        quant_mapping={
+            component_name: TorchAoConfig(
+                quant_type=quant_type_cls(),
+            )
+            for component_name in transformer_components
+        }
+    )
+
+
+def _get_diffusers_transformer_components(component_names: dict[str, Any]) -> list[str]:
+    transformer_components = [
+        component_name for component_name in component_names if component_name in _DIFFUSERS_TRANSFORMER_COMPONENTS
+    ]
+    if transformer_components:
+        return transformer_components
+
+    available_components = ", ".join(sorted(component_names.keys())) or "<none>"
+    raise NotImplementedError(
+        "Diffusers backend quantization conversion currently only supports "
+        "pipelines with transformer or transformer_2 components. "
+        f"Found pipeline components: {available_components}. Use "
+        "diffusers_load_kwargs for a native Diffusers quantization config."
+    )
+
+
+def _validate_fp8_quant_config(quant_config: Any) -> None:
+    if getattr(quant_config, "is_checkpoint_fp8_serialized", False):
+        raise NotImplementedError(
+            "Diffusers backend fp8 conversion only supports online/dynamic "
+            "TorchAO quantization; serialized vLLM fp8 checkpoints are not mapped."
+        )
+
+    activation_scheme = getattr(quant_config, "activation_scheme", "dynamic")
+    if activation_scheme != "dynamic":
+        raise NotImplementedError(
+            f"Diffusers backend fp8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}."
+        )
+
+    weight_block_size = getattr(quant_config, "weight_block_size", None)
+    if weight_block_size is not None:
+        raise NotImplementedError(
+            "Diffusers backend fp8 conversion does not map vLLM weight_block_size "
+            "to TorchAO. Use diffusers_load_kwargs for a native Diffusers config."
+        )
+
+    _ensure_no_ignored_layers(quant_config)
+
+
+def _validate_int8_quant_config(quant_config: Any) -> None:
+    if getattr(quant_config, "is_checkpoint_int8_serialized", False):
+        raise NotImplementedError(
+            "Diffusers backend int8 conversion only supports online/dynamic "
+            "TorchAO quantization; serialized vLLM int8 checkpoints are not mapped."
+        )
+
+    activation_scheme = getattr(quant_config, "activation_scheme", "dynamic")
+    if activation_scheme != "dynamic":
+        raise NotImplementedError(
+            f"Diffusers backend int8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}."
+        )
+
+    _ensure_no_ignored_layers(quant_config)
+
+
+_QUANTIZATION_CONFIG_SPECS: dict[str, _QuantizationConfigSpec] = {
+    "fp8": (_validate_fp8_quant_config, "Float8DynamicActivationFloat8WeightConfig"),
+    "int8": (_validate_int8_quant_config, "Int8DynamicActivationInt8WeightConfig"),
+}
+
+
+def _get_quantization_config_spec(method: str) -> _QuantizationConfigSpec:
+    spec = _QUANTIZATION_CONFIG_SPECS.get(method)
+    if spec is None:
+        raise NotImplementedError(
+            f"Diffusers backend quantization conversion does not support {method!r}. "
+            "Use diffusers_load_kwargs for a native Diffusers quantization config, "
+            "or use a native vLLM-Omni pipeline for this quantization method."
+        )
+    return spec
+
+
+def _build_diffusers_quantization_config(
+    quant_config: Any,
+    transformer_components: list[str],
+) -> Any:
+    """Build a Diffusers PipelineQuantizationConfig from a supported config."""
+
+    method = _get_quant_method_name(quant_config)
+    _, torchao_quant_type_name = _get_quantization_config_spec(method)
+    return _build_torchao_pipeline_quant_config(
+        torchao_quant_type_name,
+        transformer_components,
+    )
diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
index ad3175e1e..ca9253891 100644
--- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
+++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py
@@ -37,6 +37,20 @@ from vllm_omni.diffusion.models.dreamzero.action_encoder import (
     MultiEmbodimentActionEncoder,
 )
 
+# AR-Diffusion paged self-attention (in-tree experimental engine). Import at
+# module level so the isinstance check + custom-op call trace cleanly inside
+# the fullgraph-compiled DiT block (an import inside the traced region would
+# graph-break). The model still works without the engine: the payload type is
+# only ever constructed by the AR-Diffusion runner.
+try:
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+        ARDiffusionPagedLayerInputs,
+        paged_write_attn,
+    )
+except ImportError:  # pragma: no cover - experimental package always ships in-tree
+    ARDiffusionPagedLayerInputs = None
+    paged_write_attn = None
+
 # ── RoPE utilities ──────────────────────────────────────────────────
 
 
@@ -510,7 +524,7 @@ class CausalWanSelfAttention(nn.Module):
         freqs_action: torch.Tensor,
         freqs_state: torch.Tensor,
         action_register_length: int | None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -567,24 +581,44 @@ class CausalWanSelfAttention(nn.Module):
             action_v = v[:, -action_register_length:]
             v = v[:, :-action_register_length]
 
-        updated_k = kv_cache[0]
-        updated_v = kv_cache[1]
-        new_k = torch.cat([updated_k, roped_key], dim=1)
-        new_v = torch.cat([updated_v, v], dim=1)
-        new_k = new_k[:, -self.max_attention_size :]
-        new_v = new_v[:, -self.max_attention_size :]
-
-        if action_register_length is not None:
-            q_cat = torch.cat([roped_query, roped_action_query], dim=1)
-            k_cat = torch.cat([new_k, roped_action_key], dim=1)
-            v_cat = torch.cat([new_v, action_v], dim=1)
+        if ARDiffusionPagedLayerInputs is not None and isinstance(kv_cache, ARDiffusionPagedLayerInputs):
+            # Fused write+attend custom op: one opaque node in the compiled
+            # graph (slot writes + FlashAttention block-table kernel inside).
+            # Metadata tensors were prepared once per forward in _forward_blocks.
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_act, v_act = roped_action_key[0], action_v[0]
+            else:
+                q_cat = roped_query
+                k_act = v_act = None
+            x = paged_write_attn(
+                kv_cache,
+                q_cat[0],
+                roped_key[0],
+                v[0],
+                k_act,
+                v_act,
+                self.head_dim**-0.5,
+            ).unsqueeze(0)
         else:
-            q_cat = roped_query
-            k_cat = new_k
-            v_cat = new_v
+            updated_k = kv_cache[0]
+            updated_v = kv_cache[1]
+            new_k = torch.cat([updated_k, roped_key], dim=1)
+            new_v = torch.cat([updated_v, v], dim=1)
+            new_k = new_k[:, -self.max_attention_size :]
+            new_v = new_v[:, -self.max_attention_size :]
 
-        x = self.attn(q_cat, k_cat, v_cat)
-        updated_kv_cache = torch.stack([new_k, new_v], dim=0)
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_cat = torch.cat([new_k, roped_action_key], dim=1)
+                v_cat = torch.cat([new_v, action_v], dim=1)
+            else:
+                q_cat = roped_query
+                k_cat = new_k
+                v_cat = new_v
+
+            x = self.attn(q_cat, k_cat, v_cat)
+            updated_kv_cache = torch.stack([new_k, new_v], dim=0)
 
         x = x.flatten(2)
         x = self.o(x)
@@ -646,7 +680,7 @@ class CausalWanAttentionBlock(nn.Module):
         freqs_state: torch.Tensor,
         context: torch.Tensor,
         action_register_length: int | None = None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         crossattn_cache: dict | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
@@ -925,10 +959,10 @@ class CausalWanModel(nn.Module):
         action: torch.Tensor | None,
         timestep_action: torch.Tensor | None,
         state: torch.Tensor | None,
-        kv_cache: list[torch.Tensor],
-        crossattn_cache: list[dict] | None,
+        kv_cache: list[Any],
         current_start_frame: int,
-    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+        crossattn_cache: list[dict] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor | None]]:
         x = x.flatten(start_dim=2).transpose(1, 2)
         B = x.shape[0]
         F_t = timestep.shape[1]
@@ -966,7 +1000,28 @@ class CausalWanModel(nn.Module):
             clip_embedding = self.img_emb(clip_feature)
             context = torch.cat([clip_embedding, context], dim=1)
 
-        updated_kv_caches: list[torch.Tensor] = []
+        # AR-Diffusion paged path: host-side prep ONCE per branch forward, outside
+        # the compiled blocks — allocate current video/action slots, build the
+        # padded block-table metadata all 40 layers share, and hand the compiled
+        # region plain tensors (ARDiffusionPagedLayerInputs) instead of a Python
+        # context object. Lazy-allocation contract preserved: only the branch this
+        # CFG-parallel rank executes reaches its _forward_blocks.
+        if kv_cache and getattr(kv_cache[0], "is_ar_diffusion_paged_context", False):
+            if B != 1:
+                raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1")
+            fctx = kv_cache[0].forward_ctx
+            if seq_len != fctx.seq_len:
+                raise RuntimeError(
+                    f"AR-Diffusion paged context seq_len={fctx.seq_len} but current video KV has {seq_len} tokens"
+                )
+            fctx.prepare(
+                device=x.device,
+                action_len=int(action_register_length or 0),
+                query_len=int(x.shape[1]),
+            )
+            kv_cache = [c.to_layer_inputs() for c in kv_cache]
+
+        updated_kv_caches: list[torch.Tensor | None] = []
         for block_index, block in enumerate(self.blocks):
             x, updated_kv_cache = block(
                 x=x,
@@ -1000,9 +1055,9 @@ class CausalWanModel(nn.Module):
         timestep: torch.Tensor,
         context: torch.Tensor,
         seq_len: int,
-        kv_cache: list[torch.Tensor],
-        crossattn_cache: list[torch.Tensor],
+        kv_cache: list[Any],
         current_start_frame: int,
+        crossattn_cache: list[dict] | None = None,
         y: torch.Tensor | None = None,
         clip_feature: torch.Tensor | None = None,
         action: torch.Tensor | None = None,
diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
index 65e1aac70..88a1a4d47 100644
--- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py
@@ -3,7 +3,7 @@
 
 """DreamZero pipeline for vllm-omni.
 
-Entry point for DiffusionEngine.step() → pipeline.forward(req)
+Entry point for DiffusionEngine.step() -> pipeline.forward(req)
 """
 
 from __future__ import annotations
@@ -94,17 +94,102 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     Multi-output: predict_noise() returns (video_pred, action_pred).
     CFG: video gets standard CFG, action takes positive branch only.
-    State: DreamZeroState manages KV cache + frame buffer across forward() calls.
+
+    KV is managed by the AR-Diffusion engine: ``self._ar_diffusion_kv_state`` is set by the
+    runner before ``forward()`` and the pipeline routes all KV access (get / update /
+    commit / reset) through the pool-backed state. Purely duck-typed (no engine import).
     """
 
+    _ar_diffusion_kv_state = None  # set by the runner before each forward
+
+    def _kv_get(self, state, is_negative, seq_len=None, update_kv_cache=False):
+        return self._ar_diffusion_kv_state.get_kv_caches(
+            is_negative,
+            seq_len=seq_len,
+            commit_current=update_kv_cache,
+        )
+
+    def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
+        # The engine owns all KV allocation: self-attn is allocated lazily from
+        # paged contexts, and cross-attn is populated eagerly in _kv_populate_cross.
+        # Nothing is created model-side.
+        return
+
+    def _kv_commit(self, is_negative: bool):
+        self._ar_diffusion_kv_state.commit_paged_context(is_negative)
+
+    def _kv_get_cross(self, state, is_negative):
+        """Cross-attn cache from the engine pool (text k/v + I2V image k_img/v_img)."""
+        return self._ar_diffusion_kv_state.get_cross_kv_caches(is_negative)
+
+    def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None:
+        """Eagerly project cross-attn K/V for all layers into the AR-Diffusion pool.
+
+        Caches the session-invariant cross-attn projections once, per half: the text
+        ``k``/``v`` from ``text_embedding(context)`` survive window-boundary resets
+        (prompt unchanged within a session — only session resets clear them), while
+        the I2V image-token ``k_img``/``v_img`` from ``img_emb(clip_feature)`` (the
+        257 image tokens the forward splits off, cached model-side by #4154) are
+        re-projected on every window restart from the fresh CLIP features. Must run
+        after the image is encoded so ``clip_feature`` is available.
+        """
+        s = self._ar_diffusion_kv_state
+        need_text = not s._cross_text_populated.get(is_negative, False)
+        need_img = (
+            clip_feature is not None
+            and getattr(self.transformer, "model_type", "t2v") == "i2v"
+            and not s._cross_img_populated.get(is_negative, False)
+        )
+        if not need_text and not need_img:
+            return
+        projected = self.transformer.text_embedding(context) if need_text else None
+        img_ctx = self.transformer.img_emb(clip_feature) if need_img else None
+        for i, block in enumerate(self.transformer.blocks):
+            ca = block.cross_attn
+            n, d = ca.tp_num_heads, ca.head_dim
+            k = v = None
+            if projected is not None:
+                k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
+                v = ca.v(projected).unflatten(2, (n, d))
+            k_img = v_img = None
+            if img_ctx is not None:
+                k_img = ca.norm_k_img(ca.k_img(img_ctx)).unflatten(2, (n, d))
+                v_img = ca.v_img(img_ctx).unflatten(2, (n, d))
+            s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img)
+        if need_text:
+            s._cross_text_populated[is_negative] = True
+        if need_img:
+            s._cross_img_populated[is_negative] = True
+        logger.info(
+            "AR-Diffusion CROSS POPULATE [%s]: %d layers, text=%s img=%s",
+            "neg" if is_negative else "pos",
+            len(self.transformer.blocks),
+            "kept" if projected is None else tuple(context.shape),
+            None if img_ctx is None else tuple(img_ctx.shape),
+        )
+
+    def _kv_reset(self, state, *, clear_video_latents: bool = True):
+        """Reset the engine's pooled session window plus the model's non-KV state.
+
+        DreamZero resets at the attention-window boundary; the engine pool drops the
+        same window so the next forward starts fresh. ``clear_video_latents=False``
+        keeps the accumulated video latents for export.
+
+        ``clear_video_latents=False`` also marks a window ("inference") reset: the
+        prompt is unchanged, so the pool keeps the text cross-attn K/V and only the
+        image half repopulates on the restart forward.
+        """
+        state.reset(clear_video_latents=clear_video_latents)
+        self._ar_diffusion_kv_state.reset(keep_cross_text=not clear_video_latents)
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.
 
         DreamZero root checkpoint layout (GEAR-Dreams/DreamZero-DROID):
-          config.json                     — root config (action_head_cfg, architectures, etc.)
-          model-*.safetensors             — all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
-          experiment_cfg/metadata.json    — per-embodiment action normalization stats
-          vae/                            — symlink to Wan2.1 VAE (diffusers-compatible)
+          config.json                     -- root config (action_head_cfg, architectures, etc.)
+          model-*.safetensors             -- all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
+          experiment_cfg/metadata.json    -- per-embodiment action normalization stats
+          vae/                            -- symlink to Wan2.1 VAE (diffusers-compatible)
 
         Components are instantiated from config (not from_pretrained), then filled
         by load_weights() which reads root safetensors and remaps key prefixes.
@@ -119,6 +204,18 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         """
         super().__init__()
 
+        # DreamZero is engine-only: every KV access in forward() routes through
+        # the AR-Diffusion engine's pool-backed state. Fail fast here — a stale
+        # or programmatic config that leaves engine_backend="default" would
+        # otherwise only crash mid-forward on the first KV access.
+        engine_backend = str(getattr(od_config, "engine_backend", "") or "")
+        if "ar_diffusion" not in engine_backend.lower().replace("-", "_"):
+            raise ValueError(
+                "DreamZeroPipeline requires the AR-Diffusion engine; set "
+                "engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine "
+                f"in the deploy config (got engine_backend={engine_backend!r})."
+            )
+
         model_path = od_config.model
         model_config = od_config.model_config
         local_files_only = os.path.exists(model_path)
@@ -229,8 +326,11 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         self.max_action_dim: int = ah_config["max_action_dim"]
 
         self.negative_prompt: str = model_config.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT)
+        # The negative prompt is a model constant: encode it once, lazily, and
+        # reuse across every forward/session (UMT5 encode is deterministic).
+        self._negative_prompt_embeds_cache: torch.Tensor | None = None
 
-        # Embodiment name → numeric ID mapping (model knowledge)
+        # Embodiment name -> numeric ID mapping (model knowledge)
         self.embodiment_name_to_id: dict[str, int] = model_config.get(
             "embodiment_name_to_id",
             DEFAULT_EMBODIMENT_NAME_TO_ID,
@@ -345,11 +445,17 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             state=kwargs.get("state_features"),
             embodiment_id=kwargs.get("embodiment_id"),
         )
-        if kwargs.get("update_kv_cache", False) and updated_kv_caches:
-            state = kwargs.get("dreamzero_state", self.state)
+        if kwargs.get("update_kv_cache", False):
             is_neg = kwargs.get("is_negative", False)
-            for i, kv in enumerate(updated_kv_caches):
-                state.update_kv_cache(i, kv, is_negative=is_neg)
+            logger.debug(
+                "AR-Diffusion pipeline predict_noise -> commit paged context: "
+                "is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
+                is_neg,
+                kwargs.get("seq_len"),
+                kwargs.get("current_start_frame"),
+                len(updated_kv_caches),
+            )
+            self._kv_commit(is_neg)
 
         return video_pred, action_pred
 
@@ -363,6 +469,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False``
         on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and
         compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop).
+        Incremental VAE encode (``_vae_encode_encoder_chunk``) stays eager because
+        Wan ``feat_cache`` mutation is incompatible with CUDAGraph capture.
         DiT blocks use per-block ``fullgraph=True``.
         """
         if not torch.cuda.is_available():
@@ -377,6 +485,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
         # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
+        # The AR-Diffusion paged self-attention is a registered custom op, so the
+        # block stays fullgraph even on that path.
         dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
 
         logger.info(
@@ -397,9 +507,9 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
 
         try:
-            self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
+            self.vae._encode = torch.compile(self.vae._encode, **compile_ro)
         except Exception as exc:
-            logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
+            logger.warning("DreamZero: vae._encode compile failed (%s); skipping.", exc)
 
         compiled_blocks = 0
         for block in self.transformer.blocks:
@@ -438,10 +548,24 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
             try:
                 image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
-                self._encode_image(image, self.num_frames, 180, 320)
+                self._encode_image(image, self.num_frames, 180, 320, state=self.state)
             except Exception as exc:
                 logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
 
+            try:
+                self.state.reset_vae_encoder_stream()
+                dummy_video = torch.zeros(1, 3, 1, 180, 320, dtype=torch.bfloat16, device=device)
+                self._vae_stream_seed(self.state, dummy_video[:, :, :1])
+                for _ in range(4):
+                    self._vae_stream_append_frame(self.state, dummy_video[:, :, :1])
+                self._vae_stream_get_observation_latents(
+                    self.state,
+                    self.num_frame_per_block,
+                    dtype=torch.bfloat16,
+                )
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (vae encode stream) skipped: %s", exc)
+
             try:
                 latent_h, latent_w = 180 // 8, 320 // 8
                 dummy_latent = torch.zeros(
@@ -500,7 +624,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     # -----------------------------------------------------------------------
 
     def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor:
-        """uint8 [B,T,H,W,C] → bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
+        """uint8 [B,T,H,W,C] -> bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
         videos = videos.permute(0, 4, 1, 2, 3)
         if videos.dtype == torch.uint8:
             videos = videos.float() / 255.0
@@ -547,6 +671,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         num_frames: int,
         height: int,
         width: int,
+        *,
+        state: DreamZeroState | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Encode first frame via CLIP + VAE.
         Returns: (clip_feas, ys, image_latent)
@@ -582,11 +708,147 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             new_image = y[:, :, 0:1]
             y = torch.concat([msk, y], dim=1)
 
+            if state is not None:
+                if not state.vae_stream_initialized:
+                    # Seed the AR streaming encoder after the compiled full encode above;
+                    # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph.
+                    self._vae_stream_seed(state, image_input[:, :, :1])
+                else:
+                    # Window ("inference") restart with a live stream: keep the real
+                    # frame history in the Wan feat_cache and append the restart
+                    # observation instead of reseeding from scratch.
+                    self._cudagraph_mark_step_begin()
+                    self._vae_stream_append_frame(state, image_input[:, :, :1])
+
         return clip_context, y, new_image
 
+    def _vae_patchify(self, videos: torch.Tensor) -> torch.Tensor:
+        if self.vae.config.patch_size is not None:
+            from diffusers.models.autoencoders.autoencoder_kl_wan import patchify
+
+            return patchify(videos, patch_size=self.vae.config.patch_size)
+        return videos
+
+    @staticmethod
+    def _vae_clone_feat_map(feat_map: list[torch.Tensor | None]) -> list[torch.Tensor | None]:
+        return [entry.clone() if isinstance(entry, torch.Tensor) else entry for entry in feat_map]
+
+    def _vae_init_enc_feat_map(self) -> list[torch.Tensor | None]:
+        self.vae.clear_cache()
+        return self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_encode_encoder_chunk(
+        self,
+        chunk: torch.Tensor,
+        feat_map: list[torch.Tensor | None],
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """Run one Wan encoder chunk while mutating ``feat_map`` in place.
+
+        Must stay eager (not ``torch.compile``): Wan causal ``feat_cache`` updates
+        conflict with ``reduce-overhead`` CUDAGraph capture.
+        """
+        self.vae._enc_feat_map = feat_map
+        self.vae._enc_conv_idx = [0]
+        out = self.vae.encoder(
+            chunk,
+            feat_cache=self.vae._enc_feat_map,
+            feat_idx=self.vae._enc_conv_idx,
+        )
+        return out, self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_quantize_encoder_out(self, encoder_out: torch.Tensor) -> torch.Tensor:
+        enc = self.vae.quant_conv(encoder_out)
+        mu, _ = enc.chunk(2, dim=1)
+        mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
+        inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype)
+        return (mu - mean) * inv_std
+
+    def _vae_stream_seed(self, state: DreamZeroState, first_frame: torch.Tensor) -> None:
+        """Seed incremental VAE encode with the first observation frame."""
+        state.reset_vae_encoder_stream()
+        feat_map = self._vae_init_enc_feat_map()
+        chunk = self._vae_patchify(first_frame.to(dtype=self.vae.dtype))
+        self._cudagraph_mark_step_begin()
+        encoder_out, feat_map = self._vae_encode_encoder_chunk(chunk[:, :, :1], feat_map)
+        state.vae_enc_feat_map = feat_map
+        state.vae_encoder_out = encoder_out
+        state.vae_pending_body_frames = None
+        state.vae_stream_initialized = True
+
+    def _vae_stream_append_frame(self, state: DreamZeroState, new_frame: torch.Tensor) -> None:
+        """Append one pixel frame and encode a 4-frame body chunk when ready."""
+        if not state.vae_stream_initialized or state.vae_enc_feat_map is None or state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream is not initialized.")
+
+        frame = new_frame.to(dtype=self.vae.dtype)
+        if state.vae_pending_body_frames is None:
+            state.vae_pending_body_frames = frame
+        else:
+            state.vae_pending_body_frames = torch.cat([state.vae_pending_body_frames, frame], dim=2)
+
+        while state.vae_pending_body_frames is not None and state.vae_pending_body_frames.shape[2] >= 4:
+            body = state.vae_pending_body_frames[:, :, :4]
+            if state.vae_pending_body_frames.shape[2] > 4:
+                state.vae_pending_body_frames = state.vae_pending_body_frames[:, :, 4:]
+            else:
+                state.vae_pending_body_frames = None
+            chunk = self._vae_patchify(body)
+            encoder_chunk, feat_map = self._vae_encode_encoder_chunk(chunk, state.vae_enc_feat_map)
+            state.vae_encoder_out = torch.cat([state.vae_encoder_out, encoder_chunk], dim=2)
+            state.vae_enc_feat_map = feat_map
+
+    def _vae_stream_get_observation_latents(
+        self,
+        state: DreamZeroState,
+        num_latent_frames: int,
+        *,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream has no accumulated encoder output.")
+        latents = self._vae_quantize_encoder_out(state.vae_encoder_out).to(dtype=dtype)
+        if latents.shape[2] >= num_latent_frames:
+            return latents[:, :, -num_latent_frames:]
+        pad_count = num_latent_frames - latents.shape[2]
+        pad = latents[:, :, -1:].expand(-1, -1, pad_count, -1, -1)
+        return torch.cat([pad, latents], dim=2)
+
+    def _preprocess_vae_observation_window(self, videos: torch.Tensor) -> torch.Tensor:
+        _, _, num_frames_raw, _, _ = videos.shape
+        if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
+            return videos
+        if num_frames_raw // 4 != self.num_frame_per_block:
+            repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
+            videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
+            first_frame = videos[:, :, 0:1]
+            return torch.cat([first_frame, videos], dim=2)
+        first_frame = videos[:, :, 0:1]
+        return torch.cat([first_frame, videos], dim=2)
+
+    def _encode_observation_latents(
+        self,
+        state: DreamZeroState,
+        videos: torch.Tensor,
+        *,
+        latent_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Encode current robot observation into normalized VAE latents."""
+        if state.vae_stream_initialized:
+            self._cudagraph_mark_step_begin()
+            self._vae_stream_append_frame(state, videos[:, :, -1:])
+            return self._vae_stream_get_observation_latents(
+                state,
+                self.num_frame_per_block,
+                dtype=latent_dtype,
+            )
+
+        videos = self._preprocess_vae_observation_window(videos)
+        return self._encode_vae_latents(videos).to(dtype=latent_dtype)
+
     def _encode_vae_latents(self, videos: torch.Tensor) -> torch.Tensor:
         """Encode videos into normalized VAE latents."""
         input_dtype = videos.dtype
+        self._cudagraph_mark_step_begin()
         hidden = self.vae._encode(videos.to(dtype=self.vae.dtype))
         mu, _ = hidden.chunk(2, dim=1)
         mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
@@ -636,7 +898,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     ) -> None:
         """Prefill KV cache with first frame and/or current observation.
 
-        Uses predict_noise_maybe_with_cfg() for CFG parallel — same path as
+        Uses predict_noise_maybe_with_cfg() for CFG parallel -- same path as
         the denoise loop. The mixin handles rank dispatch automatically.
         KV cache update happens as a side effect inside predict_noise().
         """
@@ -647,7 +909,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         head_dim = self.transformer.dim // self.transformer.num_heads
 
         if state.current_start_frame == 0:
-            state.create_kv_caches(
+            self._kv_create(
+                state,
                 batch_size,
                 dtype,
                 device,
@@ -672,16 +935,16 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                kv_cache=self._kv_get(state, False, seq_len=frame_seqlen, update_kv_cache=True),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    kv_cache=self._kv_get(state, True, seq_len=frame_seqlen, update_kv_cache=True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -722,16 +985,16 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=True),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -764,8 +1027,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         For each timestep:
           1. Build positive_kwargs / negative_kwargs
-          2. predict_noise_maybe_with_cfg()    → (video_pred, action_pred)
-          3. scheduler_step_maybe_with_cfg()   → VideoActionScheduler
+          2. predict_noise_maybe_with_cfg()    -> (video_pred, action_pred)
+          3. scheduler_step_maybe_with_cfg()   -> VideoActionScheduler
           4. _synchronize_cfg_parallel_step_output()
         """
         seq_len = kwargs["seq_len"]
@@ -824,8 +1087,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                     hidden_states=noisy_input.transpose(1, 2),
                     timestep_video=timestep,
                     encoder_hidden_states=prompt_embeds,
-                    kv_cache=state.get_kv_caches(False),
-                    crossattn_cache=state.get_crossattn_caches(False),
+                    kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=False),
+                    crossattn_cache=self._kv_get_cross(state, False),
                     y=y,
                     clip_feature=state.clip_feas,
                     action=noisy_input_action,
@@ -839,8 +1102,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                         hidden_states=noisy_input.transpose(1, 2),
                         timestep_video=timestep,
                         encoder_hidden_states=negative_prompt_embeds,
-                        kv_cache=state.get_kv_caches(True),
-                        crossattn_cache=state.get_crossattn_caches(True),
+                        kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=False),
+                        crossattn_cache=self._kv_get_cross(state, True),
                         y=y,
                         clip_feature=state.clip_feas,
                         action=noisy_input_action,
@@ -937,7 +1200,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             device=device,
         )
 
-        # State: raw from transform → pad to (B, state_horizon=1, max_state_dim)
+        # State: raw from transform -> pad to (B, state_horizon=1, max_state_dim)
         raw_state = unified_obs["state"]
         state_for_postprocess = None
         if raw_state is not None:
@@ -977,38 +1240,66 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # Explicit reset from OpenPI serving is carried by `extra_args["reset"]`
         # on the next inference request after websocket reset/session switch.
         if extra_args.get("reset", False):
-            state.reset()
+            self._kv_reset(state)
         else:
+            # Auto-reset based on model state (before accumulation). Both a "session"
+            # reset and an "inference" (window-boundary) reset clear the model-local
+            # KV window, so route both through _kv_reset to drop the AR-Diffusion pool window
+            # too and stay parity-exact; "inference" keeps the accumulated video
+            # latents (reset_inference_state == reset(clear_video_latents=False)).
             reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
             if reset_reason == "session":
-                state.reset()
+                self._kv_reset(state)
             elif reset_reason == "inference":
-                state.reset_inference_state()
+                self._kv_reset(state, clear_video_latents=False)
         state.language = text_tokens
 
-        # Frame accumulation: stitched single frame → multi-frame video
+        # Frame accumulation: stitched single frame -> multi-frame video
         video_frames = state.accumulate_frames(stitched)  # (T, H, W, C)
         videos = torch.from_numpy(video_frames).unsqueeze(0).to(device)  # (B=1, T, H, W, C)
 
-        videos = self._preprocess_video(videos)  # → [B,C,T,H,W] bf16
+        videos = self._preprocess_video(videos)  # -> [B,C,T,H,W] bf16
         _, _, num_frames_raw, height, width = videos.shape
 
-        prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        # Optional phase timing (DZ_PHASE_TIMING=1): logs per-forward stage costs
+        # (text encode / obs VAE encode / KV prefill / denoise) at INFO. Each mark
+        # synchronizes CUDA, so leave it off for timed benchmark runs.
+        _pt = None
+        if os.environ.get("DZ_PHASE_TIMING"):
+            import time as _time
+
+            torch.accelerator.synchronize()
+            _pt = {"time": _time, "t0": _time.perf_counter(), "marks": []}
+
+        def _pt_mark(name: str) -> None:
+            if _pt is not None:
+                torch.accelerator.synchronize()
+                _pt["marks"].append((name, _pt["time"].perf_counter()))
+
+        # Prompt embeds are constant within a session (a prompt change triggers a
+        # "session" reset above, which clears this cache alongside state.language).
+        if state.prompt_embeds is None:
+            state.prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        prompt_embeds = state.prompt_embeds
         # Negative prompt for CFG uncond branch (model constant)
         negative_prompt_embeds = None
         if self.cfg_scale > 1.0:
-            neg_inputs = self.tokenizer(
-                self.negative_prompt,
-                max_length=512,
-                padding="max_length",
-                truncation=True,
-                return_tensors="pt",
-                add_special_tokens=True,
-            )
-            negative_prompt_embeds = self._encode_text(
-                neg_inputs["input_ids"].to(device),
-                neg_inputs["attention_mask"].to(device),
-            )
+            if self._negative_prompt_embeds_cache is None:
+                neg_inputs = self.tokenizer(
+                    self.negative_prompt,
+                    max_length=512,
+                    padding="max_length",
+                    truncation=True,
+                    return_tensors="pt",
+                    add_special_tokens=True,
+                )
+                self._negative_prompt_embeds_cache = self._encode_text(
+                    neg_inputs["input_ids"].to(device),
+                    neg_inputs["attention_mask"].to(device),
+                )
+            negative_prompt_embeds = self._negative_prompt_embeds_cache
+
+        _pt_mark("text_encode")
 
         # Extract first/last frame for CLIP + VAE encoding
         if num_frames_raw == 4 or num_frames_raw == 9:
@@ -1022,27 +1313,24 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 self.num_frames,
                 height,
                 width,
+                state=state,
             )
             state.clip_feas = clip_feas.to(dtype=image.dtype)
             state.ys = ys.to(dtype=image.dtype)
 
-        if state.current_start_frame != 0:
-            # Subsequent calls: encode current observation via VAE
-            if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
-                pass
-            elif num_frames_raw // 4 != self.num_frame_per_block:
-                repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
-                videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
-            else:
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
+            # Eager cross-attn population (AR-Diffusion only): cache text + image-token K/V
+            # into the pool now that the image is encoded (clip_feas available).
+            # Runs on the first forward of a session and after each window-boundary
+            # reset (current_start_frame returns to 0); the populate guards (text/img halves) gate re-entry.
+            self._kv_populate_cross(prompt_embeds, state.clip_feas, is_negative=False)
+            if negative_prompt_embeds is not None:
+                self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True)
 
+        if state.current_start_frame != 0:
             latent_dtype = videos.dtype
             with torch.no_grad():
-                image = self._encode_vae_latents(videos)
-            image = image.to(dtype=latent_dtype)
+                image = self._encode_observation_latents(state, videos, latent_dtype=latent_dtype)
+        _pt_mark("obs_vae_encode")
 
         batch_size = image.shape[0]
         generator = torch.Generator(device=device).manual_seed(self.seed)
@@ -1083,6 +1371,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             do_true_cfg,
             state,
         )
+        _pt_mark("prefill_kv")
 
         sample_scheduler = copy.deepcopy(self.scheduler)
         sample_scheduler_action = copy.deepcopy(self.scheduler)
@@ -1124,6 +1413,19 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             state_features=state_features,
             embodiment_id=embodiment_id,
         )
+        if _pt is not None:
+            _pt_mark("diffuse")
+            prev_t = _pt["t0"]
+            parts = []
+            for name, t in _pt["marks"]:
+                parts.append(f"{name}={1000 * (t - prev_t):.1f}ms")
+                prev_t = t
+            logger.info(
+                "DZ_PHASE_TIMING csf=%s total=%.1fms %s",
+                state.current_start_frame,
+                1000 * (prev_t - _pt["t0"]),
+                " ".join(parts),
+            )
 
         if state.current_start_frame == 1:
             video_out = torch.cat([image, video_out], dim=1)
@@ -1134,7 +1436,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # q99 denorm: [-1,1] → real values
         action_out = self._denormalize_action(action_out.float(), embodiment_name)
 
-        # Relative → absolute: only for relative_action_keys (joint_position only)
+        # Relative -> absolute: only for relative_action_keys (joint_position only)
         # gripper_position is NOT relative, so don't add state back to it
         if self.relative_action and state_for_postprocess is not None:
             n_relative = self.relative_action_dim  # 7 for DROID (joint only)
@@ -1146,7 +1448,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 action_out[..., :n_relative] + last_state.unsqueeze(1)  # broadcast over horizon
             )
 
-        # Squeeze batch dim for output: (B, horizon, dim) → (horizon, dim)
+        # Squeeze batch dim for output: (B, horizon, dim) -> (horizon, dim)
         actions_np = action_out.squeeze(0).float().cpu().numpy()  # (horizon, max_action_dim)
         actions_np = transform.transform_action_output(actions_np)
 
@@ -1246,7 +1548,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         q01 = stats["q01"].to(device=action.device, dtype=action.dtype)
         q99 = stats["q99"].to(device=action.device, dtype=action.dtype)
         # action shape: (B, horizon, action_dim) or (B, horizon, max_action_dim)
-        # q01/q99 shape: (actual_action_dim,) — only denorm actual dims
+        # q01/q99 shape: (actual_action_dim,) -- only denorm actual dims
         actual_dim = q01.shape[0]
         action_real = action.clone()
         action_real[..., :actual_dim] = (action[..., :actual_dim] + 1) / 2 * (q99 - q01) + q01
diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
index 5815d93d4..4c9209808 100644
--- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
+++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py
@@ -115,17 +115,24 @@ class DreamZeroState:
         # Concatenated along T before decode (matches upstream video_across_time).
         self.video_latents_across_time = saved_video_latents
 
-        # KV cache once robot-policy diffusion supports that integration.
-        self.kv_cache = None
-        self.kv_cache_neg = None
-        self.crossattn_cache = None
-        self.crossattn_cache_neg = None
         self.current_start_frame = 0
 
         self.clip_feas = None
         self.ys = None
         if clear_video_latents:
+            # Session reset: drop the prompt-embed cache and VAE encoder stream.
             self.language = None
+            self.prompt_embeds = None
+            self.reset_vae_encoder_stream()
+        # Window ("inference") resets keep both: the prompt is unchanged, and the
+        # Wan feat_cache history is independent of the DiT attention window.
+
+    def reset_vae_encoder_stream(self) -> None:
+        """Clear incremental Wan VAE encoder state used across AR steps."""
+        self.vae_stream_initialized = False
+        self.vae_enc_feat_map: list[torch.Tensor | None] | None = None
+        self.vae_encoder_out: torch.Tensor | None = None
+        self.vae_pending_body_frames: torch.Tensor | None = None
 
     def reset_inference_state(self) -> None:
         """Reset KV/frame state after local attention rolls without dropping video latents."""
@@ -166,57 +173,3 @@ class DreamZeroState:
     def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
         """Determine if state should be reset before this forward()."""
         return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None
-
-    # ------------------------------------------------------------------
-    # KV cache management
-    # ------------------------------------------------------------------
-
-    def create_kv_caches(
-        self,
-        batch_size: int,
-        dtype: torch.dtype,
-        device: torch.device,
-        num_layers: int,
-        num_heads: int,
-        head_dim: int,
-    ) -> None:
-        """Initialize empty KV caches and cross-attention caches."""
-        self.kv_cache = [
-            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
-        ]
-        self.kv_cache_neg = [
-            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
-        ]
-
-        self.crossattn_cache = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-        self.crossattn_cache_neg = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-
-    def update_kv_cache(
-        self,
-        layer_index: int,
-        updated_kv: torch.Tensor,
-        is_negative: bool = False,
-    ) -> None:
-        """Update a single layer's KV cache after prefill."""
-        cache = self.kv_cache_neg if is_negative else self.kv_cache
-        if cache is None:
-            raise RuntimeError("KV caches not initialized, call create_kv_caches first.")
-        cache[layer_index] = updated_kv.clone()
-
-    def get_kv_caches(self, is_negative: bool = False) -> list[torch.Tensor]:
-        """Get KV caches for the specified branch."""
-        cache = self.kv_cache_neg if is_negative else self.kv_cache
-        if cache is None:
-            raise RuntimeError("KV caches not initialized.")
-        return cache
-
-    def get_crossattn_caches(self, is_negative: bool = False) -> list[dict[str, bool | torch.Tensor | None]]:
-        """Get cross-attention caches for the specified branch."""
-        cache = self.crossattn_cache_neg if is_negative else self.crossattn_cache
-        if cache is None:
-            raise RuntimeError("Cross-attn caches not initialized.")
-        return cache
diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
index d039e53c8..ba980a33e 100644
--- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
+++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py
@@ -9,6 +9,35 @@ import torch
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
+    def gpu_mem_stats(self) -> dict:
+        """Peak GPU memory (GiB) for this worker's device, for profiling.
+
+        ``max_memory_reserved`` is the caching-allocator high-water mark (the best
+        in-process proxy for "peak VRAM"); ``max_memory_allocated`` is the live
+        tensor high-water mark. Both are monotonic from process start, so the
+        end-of-run values capture the whole rollout's peak.
+        """
+        dev = torch.accelerator.current_device_index()
+        return {
+            "device": int(dev),
+            "peak_reserved_gib": torch.accelerator.max_memory_reserved(dev) / (1024**3),
+            "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3),
+        }
+
+    def ar_diffusion_perf_stats(self, reset: bool = True) -> dict:
+        """Per-request server-side forward E2E times (s) recorded by the AR-Diffusion runner.
+
+        The worker-side compute time per ``execute_model`` — the analog of the
+        upstream server's per-request E2E — which excludes the engine<->worker IPC
+        that the client's ``omni.generate`` wall time includes. Returns ``[]`` when
+        the runner is the base (KV-disabled) runner that does not record timings.
+        """
+        runner = self.model_runner
+        times = list(getattr(runner, "_perf_e2e_times", []) or [])
+        if reset and hasattr(runner, "_perf_e2e_times"):
+            runner._perf_e2e_times = []
+        return {"server_e2e_s": times}
+
     @staticmethod
     def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
         decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py
index bec8d3a14..e7f97597b 100644
--- a/vllm_omni/diffusion/models/flux/pipeline_flux.py
+++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py
@@ -502,38 +502,14 @@ class FluxPipeline(
             return False
         return True
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        prompt_2: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        negative_prompt_2: str | list[str] | None = None,
-        true_cfg_scale: float = 1.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 28,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 3.5,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.FloatTensor | None = None,
-        pooled_prompt_embeds: torch.FloatTensor | None = None,
-        negative_prompt_embeds: torch.FloatTensor | None = None,
-        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         """Forward pass for flux."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
@@ -541,40 +517,39 @@ class FluxPipeline(
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
-        sigmas = common_sampling_params.sigmas or sigmas
-        guidance_scale = (
-            common_sampling_params.guidance_scale
-            if common_sampling_params.guidance_scale is not None
-            else guidance_scale
-        )
-        true_cfg_scale = common_sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = common_sampling_params.num_inference_steps or 28
+        sigmas = common_sampling_params.sigmas
+        guidance_scale = common_sampling_params.guidance_scale
+        true_cfg_scale = common_sampling_params.true_cfg_scale or 1.0
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        max_sequence_length = common_sampling_params.max_sequence_length or 512
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
-                "prompt_embeds": prompt_embeds,
-                "negative_prompt_embeds": negative_prompt_embeds,
-                "pooled_prompt_embeds": pooled_prompt_embeds,
-                "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds,
+                "prompt_embeds": None,
+                "negative_prompt_embeds": None,
+                "pooled_prompt_embeds": None,
+                "negative_pooled_prompt_embeds": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
         negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
         pooled_prompt_embeds = prompt_fields["pooled_prompt_embeds"]
         negative_pooled_prompt_embeds = prompt_fields["negative_pooled_prompt_embeds"]
+        prompt_2 = None
+        negative_prompt_2 = None
         if prompt_embeds is not None:
             prompt = None
-            prompt_2 = None
         if negative_prompt_embeds is not None:
             negative_prompt = None
-            negative_prompt_2 = None
+
+        output_type = common_sampling_params.output_type or "pil"
+        joint_attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
@@ -598,9 +573,7 @@ class FluxPipeline(
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
+        if prompt is not None:
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py
index 5a65db257..52e0a4430 100644
--- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py
+++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py
@@ -1034,7 +1034,7 @@ class Flux2Transformer2DModel(nn.Module):
             name = original_name
             if name not in params_dict and ".to_out.0." in name:
                 name = name.replace(".to_out.0.", ".to_out.")
-            # Some GGUF checkpoints include quantized tensors for modules that
+            # Some quantized checkpoints include tensors for modules that
             # are intentionally left unquantized in this model.
             param = params_dict.get(name)
             if param is None:
diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
index cf2c9f78a..e62fccba0 100644
--- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
+++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
@@ -1044,7 +1044,7 @@ class Flux2Transformer2DModel(nn.Module):
             name = original_name
             if name not in params_dict and ".to_out.0." in name:
                 name = name.replace(".to_out.0.", ".to_out.")
-            # Some GGUF checkpoints include quantized tensors for modules that
+            # Some quantized checkpoints include tensors for modules that
             # are intentionally left unquantized in this model.
             param = params_dict.get(name)
             if param is None:
diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
index 81a6f50f1..184b5d402 100644
--- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
+++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
@@ -750,108 +750,7 @@ class Flux2KleinPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        reference_image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        mask_image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        strength: float = 1.0,
-        guidance_scale: float | None = 4.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,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-        text_encoder_out_layers: tuple[int, ...] = (9, 18, 27),
-        padding_mask_crop: int | None = None,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list of these):
-                `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
-                numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
-                or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
-                list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
-                latents as `image`, but if passing latents directly it is not encoded again.
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            guidance_scale (`float`, *optional*, defaults to 4.0):
-                Guidance scale as defined in [Classifier-Free Diffusion
-                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
-                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
-                `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
-                the text `prompt`, usually at the expense of lower image quality. For step-wise distilled models,
-                `guidance_scale` is ignored.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`List[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Note that "" is used as the negative prompt in this pipeline.
-                If not provided, will be generated from "".
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
-            text_encoder_out_layers (`Tuple[int]`):
-                Layer indices to use in the `text_encoder` to derive the final prompt embeddings.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.flux2.Flux2PipelineOutput`] or `tuple`: [`~pipelines.flux2.Flux2PipelineOutput`] if
-            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
-            generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         if len(req.prompts) > 1:
             logger.warning(
                 """This model only supports a single prompt, not a batched request.""",
@@ -878,39 +777,32 @@ class Flux2KleinPipeline(
         else:
             image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else cast(PIL.Image.Image, raw_image)
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        height = req.sampling_params.height
+        width = req.sampling_params.width
         num_inference_steps = (
-            req.sampling_params.num_inference_steps
-            if req.sampling_params.num_inference_steps is not None
-            else num_inference_steps
+            50 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
         )
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        generator = req.sampling_params.generator or generator
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        generator = req.sampling_params.generator
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers)
+        max_sequence_length = req.sampling_params.max_sequence_length or 512
+        latents = req.sampling_params.latents
+        strength = req.sampling_params.strength if req.sampling_params.strength is not None else 1.0
 
-        req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
-        if any(p is not None for p in req_prompt_embeds):
-            # If at list one prompt is provided as an embedding,
-            # Then assume that the user wants to provide embeddings for all prompts, and enter this if block
-            # If the user in fact provides mixed input format, req_prompt_embeds will have some None's
-            # And `torch.stack` automatically raises an exception for us
-            prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore # intentionally expect TypeError
-
-        req_negative_prompt_embeds = [
-            p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
-        ]
-        if any(p is not None for p in req_negative_prompt_embeds):
-            negative_prompt_embeds = torch.stack(req_negative_prompt_embeds)  # type: ignore # intentionally expect TypeError
+        text_encoder_out_layers: tuple[int, ...] = req.sampling_params.extra_args.get(
+            "text_encoder_out_layers", (9, 18, 27)
+        )
+        padding_mask_crop: int | None = req.sampling_params.extra_args.get("padding_mask_crop")
+        attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
+        output_type = req.sampling_params.output_type or "pil"
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
@@ -930,12 +822,7 @@ class Flux2KleinPipeline(
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
         device = self._execution_device
 
diff --git a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
index ba128f450..ade5eac83 100644
--- a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
+++ b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py
@@ -19,6 +19,7 @@ from transformers.utils import cached_file
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig
+from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
 from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag
 from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor
 from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable
@@ -552,4 +553,9 @@ class Gr00tN1d7Processor(ProcessorMixin):
         return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs)
 
 
-AutoProcessor.register("Gr00tN1d7", Gr00tN1d7Processor)
+# transformers >= 5.5 requires the config *class* (not a model_type string) as the
+# first arg to AutoProcessor.register — it does `key.__module__` internally. Passing
+# the string "Gr00tN1d7" raises AttributeError: 'str' object has no attribute
+# '__module__'. Match the AutoModel.register(Gr00tN1d7Config, ...) call in
+# modeling/gr00t_n1d7.py.
+AutoProcessor.register(Gr00tN1d7Config, Gr00tN1d7Processor)
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py
index f9f114513..fc1317a48 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py
@@ -3,6 +3,8 @@
 
 from typing import Any
 
+import torch
+import vllm.distributed.parallel_state as _vllm_ps
 import vllm.forward_context as _vllm_fc
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
@@ -24,6 +26,30 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None:
         forward_context.in_profile_run = False
 
 
+def _set_forward_context_dp_metadata(num_tokens: int) -> None:
+    """Populate vLLM MoE DP token metadata for diffusion-only DP mappings."""
+    if not _vllm_fc.is_forward_context_available():
+        return
+    forward_context = _vllm_fc.get_forward_context()
+    if getattr(forward_context, "dp_metadata", None) is not None:
+        return
+
+    dp_group = getattr(_vllm_ps, "_DP", None)
+    if dp_group is None or getattr(dp_group, "world_size", 1) <= 1:
+        return
+
+    gathered_num_tokens: list[int | None] = [None] * dp_group.world_size
+    torch.distributed.all_gather_object(
+        gathered_num_tokens,
+        int(num_tokens),
+        group=dp_group.cpu_group,
+    )
+    if any(count is None for count in gathered_num_tokens):
+        raise RuntimeError(f"Failed to gather MoE DP token counts: {gathered_num_tokens}")
+
+    forward_context.dp_metadata = _vllm_fc.DPMetadata(torch.tensor(gathered_num_tokens, dtype=torch.int64))
+
+
 class HunyuanFusedMoEDefault:
     """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3.
 
@@ -62,6 +88,7 @@ class HunyuanFusedMoEDefault:
                 hidden_states = args[0]
             if hidden_states is not None:
                 _set_forward_context_num_tokens(hidden_states.shape[0])
+                _set_forward_context_dp_metadata(hidden_states.shape[0])
 
         moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True)
 
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
index 1212497e6..f38d1ca7c 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
@@ -1099,6 +1099,12 @@ class ImageKVCacheManager:
 
         neg_kv = (key.reshape(-1, num_kv_heads, head_dim), value.reshape(-1, num_kv_heads, head_dim))
         self._injected_ar_kv = [self._injected_ar_kv[0], neg_kv]
+        logger.debug(
+            "[AR KV Reuse] neg branch: reused shared_prefix_len=%d, neg_kv_total_len=%d (q_len_actual=%d)",
+            shared_prefix_len,
+            key.shape[1],
+            q_len_actual,
+        )
         return key, value
 
     def __call__(
@@ -2235,13 +2241,8 @@ class HunyuanImage3Model(nn.Module):
             # processed with quantization, LoRA, fine-tuning, etc.
             if self.config.tie_word_embeddings and "lm_head.weight" in name:
                 continue
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache scales for compressed-tensors quantization
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                continue
+            # KV-cache scales are renamed via maybe_remap_kv_scale_name below;
+            # quant_config.get_cache_scale was removed in vLLM v0.23.0 (see #4810).
 
             is_found = False
             for param_name, weight_name, shard_id in stacked_params_mapping:
@@ -2924,6 +2925,10 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
     ):
         ar_kv_data = model_kwargs.pop("ar_kv_data", None)
         if ar_kv_data is None:
+            logger.debug(
+                "[AR KV Reuse] cfg_rank=%s: no AR KV received, fallback to full recompute (reuse_len=0)",
+                cfg_rank,
+            )
             return input_ids, 0
 
         # 1. positive prefix len
diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
index 73e9ff666..c0476b879 100644
--- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
@@ -2209,7 +2209,6 @@ class HunyuanImage3Pipeline(
         if output_type == "latent":
             return DiffusionOutput(
                 output=latents,
-                custom_output={},
                 stage_durations=getattr(self, "stage_durations", None),
             )
 
@@ -2232,12 +2231,14 @@ class HunyuanImage3Pipeline(
         )
 
         cot_text_list = state.extra.get(_STEP_COT_TEXT_LIST) or []
-        custom_output = {}
+        metadata = {}
         if any(text is not None for text in cot_text_list):
-            custom_output["ar_generated_text"] = cot_text_list[0]
+            metadata["text"] = {"ar_generated_text": cot_text_list[0]}
         return DiffusionOutput(
-            output=image[0],
-            custom_output=custom_output,
+            output={
+                "payload": {"image": image[0]},
+                "metadata": metadata,
+            },
             stage_durations=getattr(self, "stage_durations", None),
         )
 
@@ -2302,12 +2303,14 @@ class HunyuanImage3Pipeline(
 
         outputs = self._generate(**model_inputs, **kwargs)
         image = outputs[0]
-        custom_output = {}
+        metadata = {}
         if any(t is not None for t in cot_text_list):
-            custom_output["ar_generated_text"] = cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list
+            metadata["text"] = {"ar_generated_text": cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list}
         stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None
         return DiffusionOutput(
-            output=image,
-            custom_output=custom_output,
+            output={
+                "payload": {"image": image},
+                "metadata": metadata,
+            },
             stage_durations=stage_durations,
         )
diff --git a/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py b/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py
index 6c28e760c..10498620c 100644
--- a/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py
+++ b/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py
@@ -247,11 +247,13 @@ class InternVLAA1Pipeline(nn.Module, DiffusionPipelineProfilerMixin):
             noise=extra_args.get("noise"),
             decode_image=bool(extra_args.get("decode_image", False)),
         )
-        custom_output: dict[str, Any] = {}
+        metadata: dict[str, Any] = {}
         if decoded is not None:
-            custom_output["decoded"] = decoded
+            metadata["actions"] = {"decoded": decoded}
         return DiffusionOutput(
-            output=output,
-            custom_output=custom_output,
+            output={
+                "payload": {"actions": output},
+                "metadata": metadata,
+            },
             post_process_func=get_internvla_a1_post_process_func(self.od_config),
         )
diff --git a/vllm_omni/diffusion/models/krea2/__init__.py b/vllm_omni/diffusion/models/krea2/__init__.py
new file mode 100644
index 000000000..e802b0984
--- /dev/null
+++ b/vllm_omni/diffusion/models/krea2/__init__.py
@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Krea 2 diffusion model components."""
+
+from vllm_omni.diffusion.models.krea2.krea2_transformer import Krea2Transformer2DModel
+from vllm_omni.diffusion.models.krea2.pipeline_krea2 import (
+    Krea2Pipeline,
+    get_krea2_post_process_func,
+)
+
+__all__ = [
+    "Krea2Pipeline",
+    "Krea2Transformer2DModel",
+    "get_krea2_post_process_func",
+]
diff --git a/vllm_omni/diffusion/models/krea2/krea2_transformer.py b/vllm_omni/diffusion/models/krea2/krea2_transformer.py
new file mode 100644
index 000000000..e8c0247b5
--- /dev/null
+++ b/vllm_omni/diffusion/models/krea2/krea2_transformer.py
@@ -0,0 +1,571 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.embeddings import apply_rotary_emb, get_1d_rotary_pos_embed
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+logger = init_logger(__name__)
+
+# Devices that do not support float64; RoPE frequencies fall back to float32 there.
+_FP64_UNSUPPORTED_DEVICES = frozenset({"mps", "npu", "neuron"})
+
+
+def _rope_freqs_dtype(device: torch.device) -> torch.dtype:
+    return torch.float32 if device.type in _FP64_UNSUPPORTED_DEVICES else torch.float64
+
+
+def _join_prefix(prefix: str, name: str) -> str:
+    return f"{prefix}.{name}" if prefix else name
+
+
+def _linear(
+    in_features: int,
+    out_features: int,
+    bias: bool,
+    quant_config: "QuantizationConfig | None",
+    prefix: str,
+) -> ReplicatedLinear:
+    # ReplicatedLinear (not nn.Linear) so the diffusion LoRA manager can wrap these projections.
+    return ReplicatedLinear(
+        in_features,
+        out_features,
+        bias=bias,
+        quant_config=quant_config,
+        prefix=prefix,
+        return_bias=False,
+    )
+
+
+class Krea2RMSNorm(nn.Module):
+    """RMSNorm with a zero-centered scale: the effective multiplier is ``1 + weight``, matching the Krea 2 checkpoint
+    format. Activations are upcast so the normalization runs in float32."""
+
+    def __init__(self, dim: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        dtype = hidden_states.dtype
+        hidden_states = F.rms_norm(hidden_states.float(), (self.dim,), weight=self.weight + 1.0, eps=self.eps)
+        return hidden_states.to(dtype)
+
+
+class Krea2Attention(nn.Module):
+    """Self-attention with grouped-query projections, q/k RMSNorm, rotary embeddings and a sigmoid output gate.
+
+    Q/K/V layout is ``[B, seq, heads, head_dim]``; ``attention_mask`` is a 2D boolean key-padding mask
+    ``(batch, key_seq_len)`` that the backend broadcasts.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int | None = None,
+        eps: float = 1e-5,
+        role: str = "self",
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size={hidden_size} must be divisible by num_heads={num_heads}")
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.head_dim = hidden_size // num_heads
+
+        q_dim = self.head_dim * self.num_heads
+        kv_dim = self.head_dim * self.num_kv_heads
+        self.to_q = _linear(hidden_size, q_dim, False, quant_config, _join_prefix(prefix, "to_q"))
+        self.to_k = _linear(hidden_size, kv_dim, False, quant_config, _join_prefix(prefix, "to_k"))
+        self.to_v = _linear(hidden_size, kv_dim, False, quant_config, _join_prefix(prefix, "to_v"))
+        self.to_gate = _linear(hidden_size, hidden_size, False, quant_config, _join_prefix(prefix, "to_gate"))
+        self.norm_q = Krea2RMSNorm(self.head_dim, eps=eps)
+        self.norm_k = Krea2RMSNorm(self.head_dim, eps=eps)
+        self.to_out = nn.ModuleList(
+            [_linear(hidden_size, hidden_size, False, quant_config, _join_prefix(prefix, "to_out.0")), nn.Dropout(0.0)]
+        )
+
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.num_kv_heads,
+            role=role,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        query = self.to_q(hidden_states).unflatten(-1, (self.num_heads, self.head_dim))
+        key = self.to_k(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim))
+        value = self.to_v(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim))
+        gate = self.to_gate(hidden_states)
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+        query = query.to(value.dtype)
+        key = key.to(value.dtype)
+
+        attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
+        hidden_states = self.attn(query, key, value, attn_metadata)
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states * torch.sigmoid(gate)
+        return self.to_out[0](hidden_states)
+
+
+class Krea2SwiGLU(nn.Module):
+    """SwiGLU feed-forward network."""
+
+    def __init__(
+        self, dim: int, hidden_dim: int, quant_config: "QuantizationConfig | None" = None, prefix: str = ""
+    ) -> None:
+        super().__init__()
+        self.gate = _linear(dim, hidden_dim, False, quant_config, _join_prefix(prefix, "gate"))
+        self.up = _linear(dim, hidden_dim, False, quant_config, _join_prefix(prefix, "up"))
+        self.down = _linear(hidden_dim, dim, False, quant_config, _join_prefix(prefix, "down"))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(hidden_states)) * self.up(hidden_states))
+
+
+class Krea2TextFusionBlock(nn.Module):
+    """Pre-norm transformer block (no rotary embeddings, no time modulation) used by the text fusion stage."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        intermediate_size: int,
+        eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.norm1 = Krea2RMSNorm(dim, eps=eps)
+        self.norm2 = Krea2RMSNorm(dim, eps=eps)
+        self.attn = Krea2Attention(
+            dim,
+            num_heads,
+            num_kv_heads,
+            eps=eps,
+            role="self",
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
+        )
+        self.ff = Krea2SwiGLU(dim, intermediate_size, quant_config=quant_config, prefix=_join_prefix(prefix, "ff"))
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.norm1(hidden_states), attention_mask=attention_mask)
+        hidden_states = hidden_states + self.ff(self.norm2(hidden_states))
+        return hidden_states
+
+
+class Krea2TextFusion(nn.Module):
+    """Fuses the stack of tapped text-encoder hidden states into a single sequence of text features.
+
+    Two ``layerwise_blocks`` attend across the ``num_text_layers`` axis independently for every token, a linear
+    ``projector`` collapses that axis, and two ``refiner_blocks`` attend across the token sequence.
+    """
+
+    def __init__(
+        self,
+        num_text_layers: int,
+        dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        intermediate_size: int,
+        num_layerwise_blocks: int,
+        num_refiner_blocks: int,
+        eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.layerwise_blocks = nn.ModuleList(
+            [
+                Krea2TextFusionBlock(
+                    dim,
+                    num_heads,
+                    num_kv_heads,
+                    intermediate_size,
+                    eps,
+                    quant_config=quant_config,
+                    prefix=_join_prefix(prefix, f"layerwise_blocks.{i}"),
+                )
+                for i in range(num_layerwise_blocks)
+            ]
+        )
+        self.projector = _linear(num_text_layers, 1, False, quant_config, _join_prefix(prefix, "projector"))
+        self.refiner_blocks = nn.ModuleList(
+            [
+                Krea2TextFusionBlock(
+                    dim,
+                    num_heads,
+                    num_kv_heads,
+                    intermediate_size,
+                    eps,
+                    quant_config=quant_config,
+                    prefix=_join_prefix(prefix, f"refiner_blocks.{i}"),
+                )
+                for i in range(num_refiner_blocks)
+            ]
+        )
+
+    def forward(self, encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        batch_size, seq_len, num_text_layers, dim = encoder_hidden_states.shape
+
+        hidden_states = encoder_hidden_states.reshape(batch_size * seq_len, num_text_layers, dim)
+        for block in self.layerwise_blocks:
+            hidden_states = block(hidden_states.contiguous())
+
+        hidden_states = hidden_states.reshape(batch_size, seq_len, num_text_layers, dim).permute(0, 1, 3, 2)
+        hidden_states = self.projector(hidden_states).squeeze(-1)
+
+        for block in self.refiner_blocks:
+            hidden_states = block(hidden_states, attention_mask=attention_mask)
+
+        return hidden_states
+
+
+class Krea2TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        norm_eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.zeros(6, hidden_size))
+        self.norm1 = Krea2RMSNorm(hidden_size, eps=norm_eps)
+        self.norm2 = Krea2RMSNorm(hidden_size, eps=norm_eps)
+        self.attn = Krea2Attention(
+            hidden_size,
+            num_heads,
+            num_kv_heads,
+            eps=norm_eps,
+            role="self",
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
+        )
+        self.ff = Krea2SwiGLU(
+            hidden_size, intermediate_size, quant_config=quant_config, prefix=_join_prefix(prefix, "ff")
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # temb: (B, 1, 6 * hidden_size), shared across all blocks; each block only learns an additive table.
+        modulation = temb.unflatten(-1, (6, -1)) + self.scale_shift_table
+        prescale, preshift, pregate, postscale, postshift, postgate = modulation.unbind(-2)
+
+        attn_out = self.attn(
+            (1.0 + prescale) * self.norm1(hidden_states) + preshift,
+            attention_mask=attention_mask,
+            image_rotary_emb=image_rotary_emb,
+        )
+        hidden_states = hidden_states + pregate * attn_out
+        ff_out = self.ff((1.0 + postscale) * self.norm2(hidden_states) + postshift)
+        hidden_states = hidden_states + postgate * ff_out
+        return hidden_states
+
+
+class Krea2TimestepEmbedding(nn.Module):
+    """Sinusoidal flow-time embedding (cos-first, input scaled by 1000) followed by a two-layer MLP.
+
+    Keeps the sequence dimension at size 1 so the per-block modulations broadcast over tokens.
+    """
+
+    def __init__(
+        self, embed_dim: int, hidden_size: int, quant_config: "QuantizationConfig | None" = None, prefix: str = ""
+    ) -> None:
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.linear_1 = _linear(embed_dim, hidden_size, True, quant_config, _join_prefix(prefix, "linear_1"))
+        self.linear_2 = _linear(hidden_size, hidden_size, True, quant_config, _join_prefix(prefix, "linear_2"))
+
+    def forward(self, timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        half = self.embed_dim // 2
+        freqs = torch.exp(-math.log(1e4) * torch.arange(half, dtype=torch.float32, device=timestep.device) / half)
+        args = (timestep.float() * 1e3)[:, None, None] * freqs
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(dtype)
+        return self.linear_2(F.gelu(self.linear_1(emb), approximate="tanh"))
+
+
+class Krea2TextProjection(nn.Module):
+    """Projects the fused text features into the transformer width."""
+
+    def __init__(
+        self,
+        text_dim: int,
+        hidden_size: int,
+        eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.norm = Krea2RMSNorm(text_dim, eps=eps)
+        self.linear_1 = _linear(text_dim, hidden_size, True, quant_config, _join_prefix(prefix, "linear_1"))
+        self.linear_2 = _linear(hidden_size, hidden_size, True, quant_config, _join_prefix(prefix, "linear_2"))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.linear_1(self.norm(hidden_states))
+        return self.linear_2(F.gelu(hidden_states, approximate="tanh"))
+
+
+class Krea2FinalLayer(nn.Module):
+    """Final adaptive RMSNorm and output projection."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        out_channels: int,
+        eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.zeros(2, hidden_size))
+        self.norm = Krea2RMSNorm(hidden_size, eps=eps)
+        self.linear = _linear(hidden_size, out_channels, True, quant_config, _join_prefix(prefix, "linear"))
+
+    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        modulation = temb + self.scale_shift_table
+        scale, shift = modulation.chunk(2, dim=1)
+        hidden_states = (1.0 + scale) * self.norm(hidden_states) + shift
+        return self.linear(hidden_states)
+
+
+class Krea2RotaryPosEmbed(nn.Module):
+    """Multi-axis (t, h, w) rotary position embedding, following the Flux/Krea 2 convention."""
+
+    def __init__(self, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        n_axes = ids.shape[-1]
+        cos_out = []
+        sin_out = []
+        pos = ids.float()
+        freqs_dtype = _rope_freqs_dtype(ids.device)
+        for i in range(n_axes):
+            cos, sin = get_1d_rotary_pos_embed(
+                self.axes_dim[i],
+                pos[:, i],
+                theta=self.theta,
+                repeat_interleave_real=True,
+                use_real=True,
+                freqs_dtype=freqs_dtype,
+            )
+            cos_out.append(cos)
+            sin_out.append(sin)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+class Krea2Transformer2DModel(nn.Module):
+    r"""The single-stream MMDiT flow-matching backbone used by the Krea 2 pipeline (vLLM-Omni port).
+
+    Text conditioning enters as a stack of hidden states tapped from several layers of a multimodal text encoder. A
+    small text-fusion transformer collapses the layer axis and refines the token sequence; the result is concatenated
+    with the patchified image latents into a single ``[text, image]`` sequence processed by the transformer blocks.
+    The timestep conditions every block through one shared modulation vector plus per-block learned tables.
+    """
+
+    _repeated_blocks = ["Krea2TransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "transformer_blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+
+    def __init__(
+        self,
+        in_channels: int = 64,
+        num_layers: int = 28,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 48,
+        num_key_value_heads: int = 12,
+        intermediate_size: int = 16384,
+        timestep_embed_dim: int = 256,
+        text_hidden_dim: int = 2560,
+        num_text_layers: int = 12,
+        text_num_attention_heads: int = 20,
+        text_num_key_value_heads: int = 20,
+        text_intermediate_size: int = 6912,
+        num_layerwise_text_blocks: int = 2,
+        num_refiner_text_blocks: int = 2,
+        axes_dims_rope: tuple[int, int, int] = (32, 48, 48),
+        rope_theta: float = 1000.0,
+        norm_eps: float = 1e-5,
+        od_config: "OmniDiffusionConfig | None" = None,
+        quant_config: "QuantizationConfig | None" = None,
+    ) -> None:
+        super().__init__()
+        # The pipeline reads ``self.dtype`` to cast inputs before the transformer forward.
+        self.dtype = od_config.dtype if od_config is not None else torch.get_default_dtype()
+        self.od_config = od_config
+
+        hidden_size = attention_head_dim * num_attention_heads
+        if sum(axes_dims_rope) != attention_head_dim:
+            raise ValueError(
+                f"sum(axes_dims_rope)={sum(axes_dims_rope)} must equal attention_head_dim={attention_head_dim}"
+            )
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+
+        self.img_in = _linear(in_channels, hidden_size, True, quant_config, "img_in")
+        self.time_embed = Krea2TimestepEmbedding(timestep_embed_dim, hidden_size, quant_config, "time_embed")
+        self.time_mod_proj = _linear(hidden_size, 6 * hidden_size, True, quant_config, "time_mod_proj")
+        self.text_fusion = Krea2TextFusion(
+            num_text_layers=num_text_layers,
+            dim=text_hidden_dim,
+            num_heads=text_num_attention_heads,
+            num_kv_heads=text_num_key_value_heads,
+            intermediate_size=text_intermediate_size,
+            num_layerwise_blocks=num_layerwise_text_blocks,
+            num_refiner_blocks=num_refiner_text_blocks,
+            eps=norm_eps,
+            quant_config=quant_config,
+            prefix="text_fusion",
+        )
+        self.txt_in = Krea2TextProjection(text_hidden_dim, hidden_size, norm_eps, quant_config, "txt_in")
+        self.rotary_emb = Krea2RotaryPosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope))
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                Krea2TransformerBlock(
+                    hidden_size=hidden_size,
+                    intermediate_size=intermediate_size,
+                    num_heads=num_attention_heads,
+                    num_kv_heads=num_key_value_heads,
+                    norm_eps=norm_eps,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.final_layer = Krea2FinalLayer(hidden_size, in_channels, norm_eps, quant_config, "final_layer")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        position_ids: torch.Tensor,
+        encoder_attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        r"""Predict the flow-matching velocity for the image tokens.
+
+        Args:
+            hidden_states: Packed (patchified) noisy image latents ``(batch, image_seq_len, in_channels)``.
+            encoder_hidden_states: Tapped text-encoder hidden states
+                ``(batch, text_seq_len, num_text_layers, text_hidden_dim)``.
+            timestep: Flow-matching time in ``[0, 1]`` of shape ``(batch,)``.
+            position_ids: ``(t, h, w)`` rotary coordinates ``(text_seq_len + image_seq_len, 3)``.
+            encoder_attention_mask: Boolean mask marking valid text tokens ``(batch, text_seq_len)`` or ``None``.
+        """
+        if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
+            raise ValueError(f"`position_ids` must have shape (sequence_length, 3), got {tuple(position_ids.shape)}.")
+
+        batch_size, image_seq_len, _ = hidden_states.shape
+
+        temb = self.time_embed(timestep, dtype=hidden_states.dtype)
+        temb_mod = self.time_mod_proj(F.gelu(temb, approximate="tanh"))
+
+        # 2D key-padding masks: padded text tokens are excluded as keys; their lanes are dropped at the output slice.
+        text_attention_mask = None
+        attention_mask = None
+        if encoder_attention_mask is not None:
+            text_attention_mask = encoder_attention_mask.bool()
+            image_mask = text_attention_mask.new_ones((batch_size, image_seq_len))
+            attention_mask = torch.cat([text_attention_mask, image_mask], dim=1)
+
+        encoder_hidden_states = self.text_fusion(encoder_hidden_states, attention_mask=text_attention_mask)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+        text_seq_len = encoder_hidden_states.shape[1]
+        hidden_states = self.img_in(hidden_states)
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        image_rotary_emb = self.rotary_emb(position_ids)
+
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states, temb_mod, image_rotary_emb, attention_mask)
+
+        hidden_states = hidden_states[:, text_seq_len:]
+        return self.final_layer(hidden_states, temb)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Module/parameter names mirror the diffusers ``Krea2Transformer2DModel`` exactly, so weights map 1:1.
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if name not in params_dict:
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
diff --git a/vllm_omni/diffusion/models/krea2/pipeline_krea2.py b/vllm_omni/diffusion/models/krea2/pipeline_krea2.py
new file mode 100644
index 000000000..4aa0617c7
--- /dev/null
+++ b/vllm_omni/diffusion/models/krea2/pipeline_krea2.py
@@ -0,0 +1,551 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import logging
+import os
+from collections.abc import Iterable
+from typing import ClassVar
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import AutoTokenizer, Qwen3VLModel
+from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.transformers_utils.config import get_hf_file_to_dict
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.models.krea2.krea2_transformer import Krea2Transformer2DModel
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+
+logger = logging.getLogger(__name__)
+
+# Text-encoder layer taps used when model_index.json does not specify ``text_encoder_select_layers``.
+DEFAULT_TEXT_ENCODER_SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+
+
+def get_krea2_post_process_func(od_config: OmniDiffusionConfig):
+    model_name = od_config.model
+    # Read only the VAE config (``get_hf_file_to_dict`` resolves a local dir via
+    # ``Path(model)/file`` and a hub repo via a single-file download) instead of
+    # pulling all weights just to compute ``vae_scale_factor``.
+    vae_config = get_hf_file_to_dict("vae/config.json", model_name) or {}
+    vae_scale_factor = 2 ** len(vae_config["temperal_downsample"]) if "temperal_downsample" in vae_config else 8
+
+    patch_size = 2
+    try:
+        model_index = get_hf_file_to_dict("model_index.json", model_name) or {}
+        patch_size = int(model_index.get("patch_size", patch_size))
+    except Exception:  # noqa: BLE001
+        pass
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * patch_size)
+
+    def post_process_func(images: torch.Tensor):
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, int]:
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class Krea2Pipeline(nn.Module, DiffusionPipelineProfilerMixin, ProgressBarMixin, SupportsComponentDiscovery):
+    """The Krea 2 text-to-image pipeline (vLLM-Omni port of ``diffusers.Krea2Pipeline``).
+
+    Components:
+    - ``scheduler``: ``FlowMatchEulerDiscreteScheduler`` (resolution-aware exponential time shift).
+    - ``vae``: the Qwen-Image VAE (f8, 16 latent channels).
+    - ``text_encoder``: a Qwen3-VL model; the pipeline consumes a stack of hidden states tapped from several
+      decoder layers rather than the last hidden state.
+    - ``tokenizer``: the tokenizer paired with the text encoder.
+    - ``transformer``: the Krea 2 single-stream MMDiT that predicts the flow-matching velocity.
+    """
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    _PROFILER_TARGETS: ClassVar[list[str]] = ["text_encoder.forward", "diffuse", "vae.decode"]
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+        # Only the transformer goes through the standard weights loader; text encoder and VAE load eagerly below.
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=od_config.revision,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+
+        self.device = get_local_device()
+        model = od_config.model
+        local_files_only = os.path.isdir(model)
+
+        # Pipeline-level config lives in model_index.json, not the transformer config.
+        model_index = get_hf_file_to_dict("model_index.json", model) or {}
+        select_layers = model_index.get("text_encoder_select_layers") or DEFAULT_TEXT_ENCODER_SELECT_LAYERS
+        self.text_encoder_select_layers = tuple(int(i) for i in select_layers)
+        self.is_distilled = bool(model_index.get("is_distilled", False))
+        self.patch_size = int(model_index.get("patch_size", 2))
+
+        subfolders = ["scheduler", "text_encoder", "vae", "tokenizer"]
+        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        prefetch_subfolders(model, subfolders, local_files_only=local_files_only)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen3VLModel.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=od_config.dtype,
+        )
+        # Drop the unused Qwen3-VL vision tower before moving to GPU so it never consumes GPU memory.
+        if hasattr(self.text_encoder, "visual"):
+            del self.text_encoder.visual
+        else:
+            logger.warning("Krea2: vision tower not found on text encoder; skipping drop")
+        self.text_encoder = self.text_encoder.to(self.device)
+
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLQwenImage.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=od_config.dtype,
+        ).to(self.device)
+
+        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Krea2Transformer2DModel)
+        self.transformer = Krea2Transformer2DModel(
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+
+        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * self.patch_size)
+        self.default_sample_size = 1024
+
+        # Chat template tokenized as a fixed-length block: prompt padded to a fixed length first, assistant suffix
+        # appended after the padding; the first ``prompt_template_encode_start_idx`` system tokens are later dropped.
+        self.prompt_template_encode_prefix = (
+            "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, "
+            "spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"
+        )
+        self.prompt_template_encode_suffix = "<|im_end|>\n<|im_start|>assistant\n"
+        self.prompt_template_encode_start_idx = 34
+        self.prompt_template_encode_num_suffix_tokens = 5
+
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
+    def get_text_hidden_states(
+        self,
+        prompt: str | list[str],
+        max_sequence_length: int = 512,
+        device: torch.device | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize ``prompt`` into the fixed-length Krea 2 layout and tap the selected encoder hidden states.
+
+        Returns ``(hidden_states, attention_mask)`` of shapes ``(batch, text_seq_len, num_text_layers,
+        text_hidden_dim)`` and ``(batch, text_seq_len)`` (bool).
+        """
+        device = device or self.device
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prefix_idx = self.prompt_template_encode_start_idx
+        text = [self.prompt_template_encode_prefix + e for e in prompt]
+        text_tokens = self.tokenizer(
+            text,
+            truncation=True,
+            padding="max_length",
+            max_length=max_sequence_length + prefix_idx - self.prompt_template_encode_num_suffix_tokens,
+            return_tensors="pt",
+        ).to(device)
+        suffix_tokens = self.tokenizer([self.prompt_template_encode_suffix] * len(text), return_tensors="pt").to(device)
+
+        input_ids = torch.cat([text_tokens.input_ids, suffix_tokens.input_ids], dim=1)
+        attention_mask = torch.cat([text_tokens.attention_mask, suffix_tokens.attention_mask], dim=1).bool()
+
+        # Padding sits in the middle (``[prefix | prompt | PAD | suffix]``), so positions must count only real tokens
+        # (padding consumes no position). Broadcast across the 3 mRoPE axes (T/H/W equal for text).
+        position_ids = (attention_mask.long().cumsum(dim=-1) - 1).clamp(min=0)
+        position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+        outputs = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_hidden_states=True,
+        )
+        hidden_states = torch.stack([outputs.hidden_states[i] for i in self.text_encoder_select_layers], dim=2)
+
+        hidden_states = hidden_states[:, prefix_idx:]
+        attention_mask = attention_mask[:, prefix_idx:]
+        return hidden_states, attention_mask
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        device: torch.device | None = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        max_sequence_length: int = 512,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = device or self.device
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self.get_text_hidden_states(prompt, max_sequence_length, device)
+
+        batch_size, seq_len, num_text_layers, dim = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, num_text_layers, dim)
+        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt)
+        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+        return prompt_embeds, prompt_embeds_mask
+
+    def _pack_latents(self, latents, batch_size, num_channels_latents, height, width):
+        p = self.patch_size
+        latents = latents.view(batch_size, num_channels_latents, height // p, p, width // p, p)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // p) * (width // p), num_channels_latents * p * p)
+        return latents
+
+    def _unpack_latents(self, latents, height, width):
+        batch_size, _, channels = latents.shape
+        p = self.patch_size
+
+        height = p * (int(height) // (self.vae_scale_factor * p))
+        width = p * (int(width) // (self.vae_scale_factor * p))
+
+        latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+        latents = latents.reshape(batch_size, channels // (p * p), 1, height, width)
+        return latents
+
+    @staticmethod
+    def prepare_position_ids(text_seq_len: int, grid_height: int, grid_width: int, device: torch.device):
+        """Build the ``(text_seq_len + grid_height * grid_width, 3)`` rotary coordinates: text tokens sit at the
+        origin, image tokens carry their ``(0, h, w)`` latent-grid coordinates."""
+        text_ids = torch.zeros(text_seq_len, 3, device=device)
+        image_ids = torch.zeros(grid_height, grid_width, 3, device=device)
+        image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None]
+        image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :]
+        image_ids = image_ids.reshape(grid_height * grid_width, 3)
+        return torch.cat([text_ids, image_ids], dim=0)
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        latent_height = height // self.vae_scale_factor
+        latent_width = width // self.vae_scale_factor
+        shape = (batch_size, num_channels_latents, latent_height, latent_width)
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self._pack_latents(latents, batch_size, num_channels_latents, latent_height, latent_width)
+        return latents
+
+    def _extract_prompts(self, prompts):
+        """Extract prompt and negative_prompt from the request's OmniPromptType list."""
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
+            negative_prompt = None
+        elif prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
+        else:
+            negative_prompt = None
+        return prompt, negative_prompt
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 0
+
+    def diffuse(
+        self,
+        latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        position_ids: torch.Tensor,
+        prompt_embeds: torch.Tensor,
+        prompt_embeds_mask: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds_mask: torch.Tensor | None,
+        guidance_scale: float,
+    ) -> torch.Tensor:
+        """Run the flow-matching denoising loop and return the final latents."""
+        self.scheduler.set_begin_index(0)
+        transformer_dtype = self.transformer.dtype
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for _, t in enumerate(timesteps):
+                timestep = (t / self.scheduler.config.num_train_timesteps).expand(latents.shape[0])
+                timestep = timestep.to(transformer_dtype)
+                latent_input = latents.to(transformer_dtype)
+
+                noise_pred = self.transformer(
+                    hidden_states=latent_input,
+                    encoder_hidden_states=prompt_embeds.to(transformer_dtype),
+                    timestep=timestep,
+                    position_ids=position_ids,
+                    encoder_attention_mask=prompt_embeds_mask,
+                )
+
+                if self.do_classifier_free_guidance:
+                    neg_noise_pred = self.transformer(
+                        hidden_states=latent_input,
+                        encoder_hidden_states=negative_prompt_embeds.to(transformer_dtype),
+                        timestep=timestep,
+                        position_ids=position_ids,
+                        encoder_attention_mask=negative_prompt_embeds_mask,
+                    )
+                    noise_pred = noise_pred + guidance_scale * (noise_pred - neg_noise_pred)
+
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        return latents
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 28,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        output_type: str | None = "pil",
+        max_sequence_length: int = 512,
+    ) -> DiffusionOutput:
+        extracted_prompt, extracted_negative = self._extract_prompts(req.prompts)
+        prompt = extracted_prompt or prompt
+        if extracted_negative is not None:
+            negative_prompt = extracted_negative
+
+        sp = req.sampling_params
+        height = sp.height or height
+        width = sp.width or width
+        num_inference_steps = sp.num_inference_steps or num_inference_steps
+        sigmas = sp.sigmas or sigmas
+        max_sequence_length = sp.max_sequence_length or max_sequence_length
+        generator = sp.generator or generator
+        if sp.guidance_scale_provided:
+            guidance_scale = sp.guidance_scale
+        else:
+            # Distilled checkpoint is trained guidance-free (and the request layer coerces an explicit
+            # guidance_scale=0 to "not provided"), so default it to no CFG.
+            guidance_scale = 0.0 if self.is_distilled else guidance_scale
+        num_images_per_prompt = sp.num_outputs_per_prompt if sp.num_outputs_per_prompt > 0 else num_images_per_prompt
+
+        multiple = self.vae_scale_factor * self.patch_size
+        if height % multiple != 0 or width % multiple != 0:
+            rounded_height = ((height + multiple - 1) // multiple) * multiple
+            rounded_width = ((width + multiple - 1) // multiple) * multiple
+            logger.warning(
+                "`height` and `width` must be multiples of %d; rounding up from %dx%d to %dx%d.",
+                multiple,
+                height,
+                width,
+                rounded_height,
+                rounded_width,
+            )
+            height, width = rounded_height, rounded_width
+
+        device = self.device
+        self._guidance_scale = guidance_scale
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = 1
+
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        negative_prompt_embeds = None
+        negative_prompt_embeds_mask = None
+        if self.do_classifier_free_guidance:
+            if negative_prompt is None:
+                negative_prompt = ""
+            if isinstance(negative_prompt, str):
+                negative_prompt = [negative_prompt] * batch_size
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                prompt=negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        num_channels_latents = self.transformer.in_channels // (self.patch_size**2)
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        grid_height = height // (self.vae_scale_factor * self.patch_size)
+        grid_width = width // (self.vae_scale_factor * self.patch_size)
+        position_ids = self.prepare_position_ids(prompt_embeds.shape[1], grid_height, grid_width, device)
+
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        if self.is_distilled:
+            mu = 1.15
+        else:
+            mu = calculate_shift(
+                image_seq_len,
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 6400),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.15),
+            )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu
+        )
+
+        latents = self.diffuse(
+            latents=latents,
+            timesteps=timesteps,
+            position_ids=position_ids,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            guidance_scale=guidance_scale,
+        )
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width)
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+
+        return DiffusionOutput(
+            output=image,
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        loaded_weights = loader.load_weights(weights)
+        # VAE and text encoder load eagerly in __init__; record their params so strict load checks pass.
+        loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
+        if self.text_encoder is not None:
+            loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
+        return loaded_weights
diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py
index 6f7e9b95e..dc885ac44 100644
--- a/vllm_omni/diffusion/models/lance/pipeline_lance.py
+++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py
@@ -836,8 +836,10 @@ class LancePipeline(BagelPipeline):
         frames = [Image.fromarray(f) for f in frames_np]
         logger.info("Lance t2v: decoded %d frames at %dx%d", len(frames), frames[0].width, frames[0].height)
         return DiffusionOutput(
-            output=frames[0],
-            custom_output={"video_frames": frames, "video_shape": video_shape},
+            output={
+                "payload": {"video": frames},
+                "metadata": {"video": {"shape": video_shape}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1069,13 +1071,11 @@ class LancePipeline(BagelPipeline):
             )
 
         img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)
-        # Put the PIL image into ``custom_output`` too — the orchestrator
-        # serializes ``custom_output`` across the IPC boundary but strips
-        # the bare ``output`` field, so consumers that only see
-        # ``outputs[0].output`` (e.g. gradio_demo) get ``None`` otherwise.
         return DiffusionOutput(
-            output=img,
-            custom_output={"image": img, "image_shape": image_shape},
+            output={
+                "payload": {"image": img},
+                "metadata": {"image": {"shape": image_shape}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1361,8 +1361,10 @@ class LancePipeline(BagelPipeline):
         frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], out_shape)
         frames = [Image.fromarray(f) for f in frames_np]
         return DiffusionOutput(
-            output=frames[0],
-            custom_output={"video_frames": frames, "video_shape": out_shape},
+            output={
+                "payload": {"video": frames},
+                "metadata": {"video": {"shape": out_shape}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1622,8 +1624,10 @@ class LancePipeline(BagelPipeline):
         frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape)
         frames = [Image.fromarray(f) for f in frames_np]
         return DiffusionOutput(
-            output=frames[0],
-            custom_output={"video_frames": frames, "video_shape": video_shape},
+            output={
+                "payload": {"video": frames},
+                "metadata": {"video": {"shape": video_shape}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1705,8 +1709,10 @@ class LancePipeline(BagelPipeline):
         text_output = text_output.lstrip("\n")
         logger.info("Lance x2t_video: generated %d tokens", token_ids.shape[0])
         return DiffusionOutput(
-            output=text_output,
-            custom_output={"text_output": text_output},
+            output={
+                "payload": {"text": text_output},
+                "metadata": {"text": {"text_output": text_output}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1801,8 +1807,10 @@ class LancePipeline(BagelPipeline):
         text_output = text_output.lstrip("\n")
         logger.info("Lance x2t_image: generated %d tokens", token_ids.shape[0])
         return DiffusionOutput(
-            output=text_output,
-            custom_output={"text_output": text_output},
+            output={
+                "payload": {"text": text_output},
+                "metadata": {"text": {"text_output": text_output}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py
index 2884adfca..b95f59228 100644
--- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py
+++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py
@@ -495,66 +495,35 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
         noise_pred = comb_pred * scale
         return noise_pred
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 4.5,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        enable_cfg_renorm: bool | None = True,
-        cfg_renorm_min: float | None = 0.0,
-        enable_prompt_rewrite: bool | None = True,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
-        height = req.sampling_params.height or height or self.default_sample_size * self.vae_scale_factor
-        width = req.sampling_params.width or width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        generator = req.sampling_params.generator or generator
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        generator = req.sampling_params.generator
+        guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt is not None
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        enable_prompt_rewrite = req.sampling_params.extra_args.get("enable_prompt_rewrite", enable_prompt_rewrite)
-        enable_cfg_renorm = req.sampling_params.extra_args.get("enable_cfg_renorm", enable_cfg_renorm)
-        cfg_renorm_min = req.sampling_params.extra_args.get("cfg_renorm_min", cfg_renorm_min)
+        latents = req.sampling_params.latents
 
-        req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
-        if any(p is not None for p in req_prompt_embeds):
-            # If at list one prompt is provided as an embedding,
-            # Then assume that the user wants to provide embeddings for all prompts, and enter this if block
-            # If the user in fact provides mixed input format, req_prompt_embeds will have some None's
-            # And `torch.stack` automatically raises an exception for us
-            prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore # intentionally expect TypeError
-
-        req_negative_prompt_embeds = [
-            p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
-        ]
-        if any(p is not None for p in req_negative_prompt_embeds):
-            negative_prompt_embeds = torch.stack(req_negative_prompt_embeds)  # type: ignore # intentionally expect TypeError
+        extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
+        enable_prompt_rewrite: bool = extra_args.get("enable_prompt_rewrite", True)
+        enable_cfg_renorm: bool = extra_args.get("enable_cfg_renorm", True)
+        cfg_renorm_min: float = extra_args.get("cfg_renorm_min", 0.0)
+        joint_attention_kwargs: dict[str, Any] | None = None
+        output_type = req.sampling_params.output_type or "pil"
 
         self.check_inputs(
             prompt,
@@ -571,12 +540,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = len(prompt)
 
         device = self.device
         if enable_prompt_rewrite and prompt is not None:
diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py
index 195aa89a4..4821425b0 100644
--- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py
+++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py
@@ -6,7 +6,7 @@ import math
 import os
 import re
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -542,24 +542,7 @@ class LongCatImageEditPipeline(
                 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
             )
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 3.5,
-        num_images_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         if len(req.prompts) > 1:
@@ -570,27 +553,22 @@ class LongCatImageEditPipeline(
         first_prompt = req.prompts[0]
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
         negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
-        prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("prompt_embeds")
-        negative_prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt_embeds")  # type: ignore # Why it is list[torch.Tensor] in OmniTokenInputs or OmniEmbedsPrompt? Doesn't make sense
 
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt is not None
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.sampling_params.generator or generator
+        generator = req.sampling_params.generator
         height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        output_type = req.sampling_params.output_type or "pil"
 
-        if prompt is not None:
-            batch_size = 1 if isinstance(prompt, str) else len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
 
         if not isinstance(first_prompt, str) and "preprocessed_image" in (
             additional_information := first_prompt.get("additional_information", {})
@@ -600,13 +578,7 @@ class LongCatImageEditPipeline(
             calculated_height = additional_information.get("calculated_height", height)
             calculated_width = additional_information.get("calculated_width", width)
         else:
-            image_size = image[0].size if isinstance(image, list) else image.size
-            calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1])
-
-            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
-                image = self.image_processor.resize(image, calculated_height, calculated_width)
-                prompt_image = self.image_processor.resize(image, calculated_height // 2, calculated_width // 2)
-                image = self.image_processor.preprocess(image, calculated_height, calculated_width)
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
         self.check_inputs(
             prompt,
@@ -621,7 +593,8 @@ class LongCatImageEditPipeline(
             prompt=prompt, image=prompt_image, prompt_embeds=prompt_embeds, num_images_per_prompt=num_images_per_prompt
         )
 
-        if guidance_scale > 1:
+        do_true_cfg = guidance_scale > 1 and negative_prompt is not None
+        if do_true_cfg:
             (negative_prompt_embeds, negative_text_ids) = self.encode_prompt(
                 prompt=negative_prompt,
                 image=prompt_image,
@@ -680,7 +653,6 @@ class LongCatImageEditPipeline(
                 latent_model_input = torch.cat([latents, image_latents], dim=1)
 
             timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
-            do_true_cfg = guidance_scale > 1
             positive_kwargs = {
                 "hidden_states": latent_model_input,
                 "timestep": timestep / 1000,
diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py
index bf57d3f9e..76a26b5b0 100644
--- a/vllm_omni/diffusion/models/ltx2/__init__.py
+++ b/vllm_omni/diffusion/models/ltx2/__init__.py
@@ -10,10 +10,8 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     get_ltx2_post_process_func,
     load_transformer_config,
 )
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import (
-    LTX23ImageToVideoPipeline,
-    LTX23Pipeline,
-)
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import (
     LTX2I2VDMD2Pipeline,
     LTX2ImageToVideoPipeline,
diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py
index 9183065dc..45be23296 100644
--- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py
+++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py
@@ -19,6 +19,7 @@ import json
 import os
 from collections.abc import Iterable
 from contextlib import nullcontext
+from dataclasses import dataclass
 from typing import Any, ClassVar
 
 import numpy as np
@@ -68,6 +69,26 @@ def _get_audio_latents_from_sampling(sampling: Any) -> torch.Tensor | None:
     return sampling.extra_args.get("audio_latents")
 
 
+def _is_output_rank() -> bool:
+    return not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+
+
+def _vae_decode_needs_all_ranks(vae: Any) -> bool:
+    if not torch.distributed.is_initialized():
+        return False
+    is_distributed_enabled = getattr(vae, "is_distributed_enabled", None)
+    if not callable(is_distributed_enabled):
+        return False
+    try:
+        return bool(is_distributed_enabled())
+    except Exception:
+        return False
+
+
+def _should_decode_video_on_rank(vae: Any) -> bool:
+    return _is_output_rank() or _vae_decode_needs_all_ranks(vae)
+
+
 # Try to import LTX2VocoderWithBWE (diffusers >= 0.38.0)
 try:
     from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
@@ -75,6 +96,87 @@ except ImportError:
     LTX2VocoderWithBWE = None
 
 
+@dataclass
+class _LTX23RequestInputs:
+    prompt: str | list[str] | None
+    negative_prompt: str | list[str] | None
+    height: int
+    width: int
+    num_frames: int
+    frame_rate: float
+    num_inference_steps: int
+    guidance_scale: float
+    num_videos_per_prompt: int
+    generator: torch.Generator | list[torch.Generator] | None
+    latents: torch.Tensor | None
+    audio_latents: torch.Tensor | None
+    prompt_embeds: torch.Tensor | None
+    negative_prompt_embeds: torch.Tensor | None
+    prompt_attention_mask: torch.Tensor | None
+    negative_prompt_attention_mask: torch.Tensor | None
+    decode_timestep: float | list[float]
+    decode_noise_scale: float | list[float] | None
+    output_type: str
+    max_sequence_length: int
+
+
+@dataclass
+class _LTX23PromptContext:
+    batch_size: int
+    connector_prompt_embeds: torch.Tensor
+    connector_audio_prompt_embeds: torch.Tensor
+    connector_attention_mask: torch.Tensor
+    positive_connector_prompt_embeds: torch.Tensor
+    positive_connector_audio_prompt_embeds: torch.Tensor
+    positive_connector_attention_mask: torch.Tensor
+    negative_connector_prompt_embeds: torch.Tensor | None
+    negative_connector_audio_prompt_embeds: torch.Tensor | None
+    negative_connector_attention_mask: torch.Tensor | None
+
+
+@dataclass
+class _LTX23ForwardContext:
+    req: DiffusionRequestBatch
+    request_inputs: _LTX23RequestInputs
+    prompt_context: _LTX23PromptContext
+    device: torch.device
+    cfg_parallel_ready: bool
+    attention_kwargs: dict[str, Any] | None
+    latent_num_frames: int
+    latent_height: int
+    latent_width: int
+    latent_mel_bins: int
+    original_audio_num_frames: int
+    padded_audio_num_frames: int
+    timesteps: torch.Tensor
+    audio_scheduler: Any
+    video_audio_scheduler: Any
+
+    @property
+    def batch_size(self) -> int:
+        return self.prompt_context.batch_size
+
+    @property
+    def num_videos_per_prompt(self) -> int:
+        return self.request_inputs.num_videos_per_prompt
+
+
+@dataclass
+class _LTX23DenoiseContext:
+    latents: torch.Tensor
+    audio_latents: torch.Tensor
+    video_coords: torch.Tensor
+    audio_coords: torch.Tensor
+    conditioning_mask: torch.Tensor | None = None
+    conditioning_mask_for_model: torch.Tensor | None = None
+
+
+def _repeat_prompt_tensor_for_outputs(tensor: torch.Tensor, num_videos_per_prompt: int) -> torch.Tensor:
+    if num_videos_per_prompt == 1:
+        return tensor
+    return tensor.repeat_interleave(num_videos_per_prompt, dim=0)
+
+
 def _detect_vocoder_output_sample_rate(model: str) -> int | None:
     """Detect the vocoder output sample rate from vocoder/config.json.
 
@@ -121,6 +223,58 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
+def _expand_per_prompt_decode_value(
+    value: float | list[float],
+    *,
+    prompt_batch_size: int,
+    effective_batch_size: int,
+    field_name: str,
+) -> list[float]:
+    if not isinstance(value, list):
+        return [value] * effective_batch_size
+    if len(value) == 1:
+        return value * effective_batch_size
+    if len(value) == effective_batch_size:
+        return value
+    if prompt_batch_size > 0 and len(value) == prompt_batch_size and effective_batch_size % prompt_batch_size == 0:
+        repeats = effective_batch_size // prompt_batch_size
+        return [item for item in value for _ in range(repeats)]
+    raise ValueError(
+        f"`{field_name}` must have length 1, prompt batch size ({prompt_batch_size}), or effective batch size"
+        f" ({effective_batch_size}); got {len(value)}."
+    )
+
+
+def _prepare_decode_timestep_conditioning(
+    *,
+    decode_timestep: float | list[float],
+    decode_noise_scale: float | list[float] | None,
+    prompt_batch_size: int,
+    effective_batch_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    decode_timestep_values = _expand_per_prompt_decode_value(
+        decode_timestep,
+        prompt_batch_size=prompt_batch_size,
+        effective_batch_size=effective_batch_size,
+        field_name="decode_timestep",
+    )
+    if decode_noise_scale is None:
+        decode_noise_scale_values = decode_timestep_values
+    else:
+        decode_noise_scale_values = _expand_per_prompt_decode_value(
+            decode_noise_scale,
+            prompt_batch_size=prompt_batch_size,
+            effective_batch_size=effective_batch_size,
+            field_name="decode_noise_scale",
+        )
+    return (
+        torch.tensor(decode_timestep_values, device=device, dtype=dtype),
+        torch.tensor(decode_noise_scale_values, device=device, dtype=dtype)[:, None, None, None, None],
+    )
+
+
 class LTX23Pipeline(
     nn.Module,
     CFGParallelMixin,
@@ -353,12 +507,9 @@ class LTX23Pipeline(
         # [49 x (B, seq, 3840)] -> [B, seq, 3840, 49] -> [B, seq, 188160]
         prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype)
 
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
-
         prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
-        prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1)
+        prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt)
+        prompt_attention_mask = _repeat_prompt_tensor_for_outputs(prompt_attention_mask, num_videos_per_prompt)
 
         return prompt_embeds, prompt_attention_mask
 
@@ -384,6 +535,8 @@ class LTX23Pipeline(
         else:
             batch_size = prompt_embeds.shape[0]
 
+        negative_prompt_embeds_provided = negative_prompt_embeds is not None
+
         if prompt_embeds is None:
             prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
                 prompt=prompt,
@@ -392,6 +545,9 @@ class LTX23Pipeline(
                 device=device,
                 dtype=dtype,
             )
+        elif num_videos_per_prompt > 1:
+            prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt)
+            prompt_attention_mask = _repeat_prompt_tensor_for_outputs(prompt_attention_mask, num_videos_per_prompt)
 
         if do_classifier_free_guidance and negative_prompt_embeds is None:
             negative_prompt = negative_prompt or ""
@@ -416,6 +572,12 @@ class LTX23Pipeline(
                 device=device,
                 dtype=dtype,
             )
+        elif do_classifier_free_guidance and negative_prompt_embeds_provided and num_videos_per_prompt > 1:
+            negative_prompt_embeds = _repeat_prompt_tensor_for_outputs(negative_prompt_embeds, num_videos_per_prompt)
+            negative_prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
+                negative_prompt_attention_mask,
+                num_videos_per_prompt,
+            )
 
         return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask
 
@@ -542,6 +704,67 @@ class LTX23Pipeline(
             return audio_latent_length
         return provided_latent_length
 
+    def _decode_output(
+        self,
+        *,
+        latents: torch.Tensor,
+        audio_latents: torch.Tensor,
+        output_type: str,
+        connector_prompt_embeds: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None,
+        device: torch.device,
+        decode_timestep: float | list[float],
+        decode_noise_scale: float | list[float] | None,
+        prompt_batch_size: int,
+    ) -> DiffusionOutput:
+        if output_type == "latent":
+            return DiffusionOutput(
+                output=(latents, audio_latents),
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            )
+
+        latents = latents.to(connector_prompt_embeds.dtype)
+        if not self.vae.config.timestep_conditioning:
+            timestep_decode = None
+        else:
+            noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
+            timestep_decode, decode_noise_scale_t = _prepare_decode_timestep_conditioning(
+                decode_timestep=decode_timestep,
+                decode_noise_scale=decode_noise_scale,
+                prompt_batch_size=prompt_batch_size,
+                effective_batch_size=latents.shape[0],
+                device=device,
+                dtype=latents.dtype,
+            )
+            latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise
+
+        if _should_decode_video_on_rank(self.vae):
+            latents = latents.to(self.vae.dtype)
+            video = self.vae.decode(latents, timestep_decode, return_dict=False)[0]
+        else:
+            video = torch.empty(0, device=latents.device, dtype=latents.dtype)
+
+        if not _is_output_rank():
+            return DiffusionOutput(
+                output=(
+                    torch.empty(0, device=video.device, dtype=video.dtype),
+                    torch.empty(0, device=audio_latents.device, dtype=audio_latents.dtype),
+                ),
+                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+            )
+
+        if video.numel() > 0:
+            video = self.video_processor.postprocess_video(video, output_type=output_type)
+
+        audio_latents = audio_latents.to(self.audio_vae.dtype)
+        generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
+        audio = self.vocoder(generated_mel_spectrograms)
+
+        return DiffusionOutput(
+            output=(video, audio),
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
     # ------------------------------------------------------------------
     # Latent preparation
     # ------------------------------------------------------------------
@@ -830,6 +1053,713 @@ class LTX23Pipeline(
             torch.cuda.current_stream(device).synchronize()
         return latents
 
+    def _resolve_request_inputs(
+        self,
+        req: DiffusionRequestBatch,
+        *,
+        prompt: str | list[str] | None,
+        negative_prompt: str | list[str] | None,
+        height: int | None,
+        width: int | None,
+        num_frames: int | None,
+        frame_rate: float | None,
+        num_inference_steps: int | None,
+        timesteps: list[int] | None,
+        guidance_scale: float,
+        num_videos_per_prompt: int | None,
+        generator: torch.Generator | list[torch.Generator] | None,
+        latents: torch.Tensor | None,
+        audio_latents: torch.Tensor | None,
+        prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds: torch.Tensor | None,
+        prompt_attention_mask: torch.Tensor | None,
+        negative_prompt_attention_mask: torch.Tensor | None,
+        decode_timestep: float | list[float],
+        decode_noise_scale: float | list[float] | None,
+        output_type: str,
+        max_sequence_length: int | None,
+    ) -> _LTX23RequestInputs:
+        sampling_params_list = req.sampling_params_list
+        common_sampling_params = sampling_params_list[0]
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
+            negative_prompt = None
+        elif req.prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+
+        height = common_sampling_params.height or height or 512
+        width = common_sampling_params.width or width or 768
+        num_frames = common_sampling_params.num_frames or num_frames or 121
+        frame_rate = common_sampling_params.resolved_frame_rate or frame_rate or 24.0
+        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps or 40
+        if timesteps is None:
+            num_inference_steps = max(int(num_inference_steps), 2)
+        elif len(timesteps) < 2:
+            raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.")
+
+        num_videos_per_prompt = (
+            common_sampling_params.num_outputs_per_prompt
+            if common_sampling_params.num_outputs_per_prompt > 0
+            else num_videos_per_prompt or 1
+        )
+        max_sequence_length = (
+            common_sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length
+        )
+
+        if common_sampling_params.guidance_scale_provided:
+            guidance_scale = common_sampling_params.guidance_scale
+
+        if generator is None:
+            generator = req.collate_request_generators(num_videos_per_prompt, generator)
+
+        latents = req.collate_request_tensors("latents", latents)
+        audio_latents = DiffusionRequestBatch.collate_tensors(
+            [_get_audio_latents_from_sampling(sampling) for sampling in sampling_params_list],
+            "audio_latents",
+            audio_latents,
+        )
+
+        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
+            req.prompts,
+            {
+                "prompt_embeds": prompt_embeds,
+                "negative_prompt_embeds": negative_prompt_embeds,
+                "prompt_attention_mask": prompt_attention_mask,
+                "negative_prompt_attention_mask": negative_prompt_attention_mask,
+            },
+            field_aliases={
+                "prompt_attention_mask": ("prompt_attention_mask", "attention_mask"),
+                "negative_prompt_attention_mask": (
+                    "negative_prompt_attention_mask",
+                    "negative_attention_mask",
+                ),
+            },
+        )
+        prompt_embeds = prompt_fields["prompt_embeds"]
+        negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
+        prompt_attention_mask = prompt_fields["prompt_attention_mask"]
+        negative_prompt_attention_mask = prompt_fields["negative_prompt_attention_mask"]
+        if prompt_embeds is not None:
+            prompt = None
+        if negative_prompt_embeds is not None:
+            negative_prompt = None
+
+        if common_sampling_params.decode_timestep is not None:
+            decode_timestep = common_sampling_params.decode_timestep
+        if common_sampling_params.decode_noise_scale is not None:
+            decode_noise_scale = common_sampling_params.decode_noise_scale
+        if common_sampling_params.output_type is not None:
+            output_type = common_sampling_params.output_type
+
+        return _LTX23RequestInputs(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=int(height),
+            width=int(width),
+            num_frames=int(num_frames),
+            frame_rate=float(frame_rate),
+            num_inference_steps=int(num_inference_steps),
+            guidance_scale=guidance_scale,
+            num_videos_per_prompt=int(num_videos_per_prompt),
+            generator=generator,
+            latents=latents,
+            audio_latents=audio_latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            decode_timestep=decode_timestep,
+            decode_noise_scale=decode_noise_scale,
+            output_type=output_type,
+            max_sequence_length=int(max_sequence_length),
+        )
+
+    def _prepare_prompt_context(
+        self,
+        *,
+        prompt: str | list[str] | None,
+        negative_prompt: str | list[str] | None,
+        prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds: torch.Tensor | None,
+        prompt_attention_mask: torch.Tensor | None,
+        negative_prompt_attention_mask: torch.Tensor | None,
+        num_videos_per_prompt: int,
+        max_sequence_length: int,
+    ) -> _LTX23PromptContext:
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = (
+            self.encode_prompt(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                do_classifier_free_guidance=self.do_classifier_free_guidance,
+                num_videos_per_prompt=num_videos_per_prompt,
+                prompt_embeds=prompt_embeds,
+                negative_prompt_embeds=negative_prompt_embeds,
+                prompt_attention_mask=prompt_attention_mask,
+                negative_prompt_attention_mask=negative_prompt_attention_mask,
+                max_sequence_length=max_sequence_length,
+                device=self.device,
+            )
+        )
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
+
+        connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
+            prompt_embeds,
+            prompt_attention_mask,
+            padding_side=getattr(self.tokenizer, "padding_side", "left"),
+        )
+
+        positive_connector_prompt_embeds = connector_prompt_embeds
+        positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds
+        positive_connector_attention_mask = connector_attention_mask
+        negative_connector_prompt_embeds = None
+        negative_connector_audio_prompt_embeds = None
+        negative_connector_attention_mask = None
+        if self.do_classifier_free_guidance:
+            split_batch = batch_size * num_videos_per_prompt
+            negative_connector_prompt_embeds = connector_prompt_embeds[:split_batch]
+            positive_connector_prompt_embeds = connector_prompt_embeds[split_batch:]
+            negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:split_batch]
+            positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds[split_batch:]
+            negative_connector_attention_mask = connector_attention_mask[:split_batch]
+            positive_connector_attention_mask = connector_attention_mask[split_batch:]
+
+        return _LTX23PromptContext(
+            batch_size=batch_size,
+            connector_prompt_embeds=connector_prompt_embeds,
+            connector_audio_prompt_embeds=connector_audio_prompt_embeds,
+            connector_attention_mask=connector_attention_mask,
+            positive_connector_prompt_embeds=positive_connector_prompt_embeds,
+            positive_connector_audio_prompt_embeds=positive_connector_audio_prompt_embeds,
+            positive_connector_attention_mask=positive_connector_attention_mask,
+            negative_connector_prompt_embeds=negative_connector_prompt_embeds,
+            negative_connector_audio_prompt_embeds=negative_connector_audio_prompt_embeds,
+            negative_connector_attention_mask=negative_connector_attention_mask,
+        )
+
+    def _setup_forward_runtime(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        attention_kwargs: dict[str, Any] | None,
+    ) -> bool:
+        self._guidance_scale = request_inputs.guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._interrupt = False
+        self._current_timestep = None
+        cfg_world_size = get_classifier_free_guidance_world_size()
+        if self.do_classifier_free_guidance and cfg_world_size not in (1, 2):
+            raise ValueError(
+                f"LTX23Pipeline supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}."
+            )
+        return self.do_classifier_free_guidance and cfg_world_size > 1
+
+    def _check_forward_inputs(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        image: Any | None = None,
+    ) -> None:
+        self.check_inputs(
+            prompt=request_inputs.prompt,
+            height=request_inputs.height,
+            width=request_inputs.width,
+            prompt_embeds=request_inputs.prompt_embeds,
+            negative_prompt_embeds=request_inputs.negative_prompt_embeds,
+            prompt_attention_mask=request_inputs.prompt_attention_mask,
+            negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
+        )
+
+    def _resolve_video_latent_dimensions(
+        self,
+        request_inputs: _LTX23RequestInputs,
+    ) -> tuple[int, int, int]:
+        latent_num_frames = (request_inputs.num_frames - 1) // self.vae_temporal_compression_ratio + 1
+        latent_height = request_inputs.height // self.vae_spatial_compression_ratio
+        latent_width = request_inputs.width // self.vae_spatial_compression_ratio
+        if request_inputs.latents is not None and request_inputs.latents.ndim == 5:
+            _, _, latent_num_frames, latent_height, latent_width = request_inputs.latents.shape
+        return latent_num_frames, latent_height, latent_width
+
+    def _prepare_video_latents_stage(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        prompt_context: _LTX23PromptContext,
+        *,
+        device: torch.device,
+        noise_scale: float,
+        image: Any | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        num_channels_latents = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            prompt_context.batch_size * request_inputs.num_videos_per_prompt,
+            num_channels_latents,
+            request_inputs.height,
+            request_inputs.width,
+            request_inputs.num_frames,
+            noise_scale,
+            torch.float32,
+            device,
+            request_inputs.generator,
+            request_inputs.latents,
+        )
+        return latents, None
+
+    def _prepare_audio_latents_stage(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        prompt_context: _LTX23PromptContext,
+        *,
+        device: torch.device,
+        noise_scale: float,
+    ) -> tuple[torch.Tensor, int, int, int]:
+        duration_s = request_inputs.num_frames / request_inputs.frame_rate
+        audio_latents_per_second = (
+            self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio)
+        )
+        audio_num_frames = round(duration_s * audio_latents_per_second)
+        audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, request_inputs.audio_latents)
+
+        num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64
+        latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
+        num_channels_latents_audio = self.audio_vae.config.latent_channels if self.audio_vae is not None else 8
+        audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents(
+            prompt_context.batch_size * request_inputs.num_videos_per_prompt,
+            num_channels_latents=num_channels_latents_audio,
+            audio_latent_length=audio_num_frames,
+            num_mel_bins=num_mel_bins,
+            noise_scale=noise_scale,
+            dtype=torch.float32,
+            device=device,
+            generator=request_inputs.generator,
+            latents=request_inputs.audio_latents,
+        )
+        return audio_latents, original_audio_num_frames, padded_audio_num_frames, latent_mel_bins
+
+    def _make_video_audio_scheduler(
+        self,
+        audio_scheduler: Any,
+        latent_num_frames: int,
+        latent_height: int,
+        latent_width: int,
+    ) -> Any:
+        return _VideoAudioScheduler(self.scheduler, audio_scheduler)
+
+    def _prepare_scheduler_stage(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        *,
+        device: torch.device,
+        sigmas: list[float] | None,
+        timesteps: list[int] | None,
+        latent_num_frames: int,
+        latent_height: int,
+        latent_width: int,
+    ) -> tuple[Any, Any, torch.Tensor]:
+        sigmas = (
+            np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps)
+            if sigmas is None
+            else sigmas
+        )
+        # Use max_image_seq_len (not actual video_sequence_length) for mu calculation,
+        # matching diffusers' LTX2Pipeline which hardcodes this value.
+        mu = calculate_shift(
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_image_seq_len", 1024),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.95),
+            self.scheduler.config.get("max_shift", 2.05),
+        )
+        audio_scheduler = copy.deepcopy(self.scheduler)
+        video_audio_scheduler = self._make_video_audio_scheduler(
+            audio_scheduler,
+            latent_num_frames,
+            latent_height,
+            latent_width,
+        )
+        _ = retrieve_timesteps(
+            audio_scheduler,
+            request_inputs.num_inference_steps,
+            device,
+            timesteps,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        timesteps, _ = retrieve_timesteps(
+            self.scheduler,
+            request_inputs.num_inference_steps,
+            device,
+            timesteps,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        self._num_timesteps = len(timesteps)
+        return audio_scheduler, video_audio_scheduler, timesteps
+
+    def _prepare_rope_coords_stage(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        latents: torch.Tensor,
+        audio_latents: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        request_inputs = forward_ctx.request_inputs
+        video_coords = self.transformer.rope.prepare_video_coords(
+            latents.shape[0],
+            forward_ctx.latent_num_frames,
+            forward_ctx.latent_height,
+            forward_ctx.latent_width,
+            latents.device,
+            fps=request_inputs.frame_rate,
+        )
+        audio_coords = self.transformer.audio_rope.prepare_audio_coords(
+            audio_latents.shape[0],
+            forward_ctx.padded_audio_num_frames,
+            audio_latents.device,
+        )
+        return video_coords, audio_coords
+
+    def _prepare_denoise_context_for_cfg(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> _LTX23DenoiseContext:
+        if self.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready:
+            denoise_ctx.video_coords = denoise_ctx.video_coords.repeat(
+                (2,) + (1,) * (denoise_ctx.video_coords.ndim - 1)
+            )
+            denoise_ctx.audio_coords = denoise_ctx.audio_coords.repeat(
+                (2,) + (1,) * (denoise_ctx.audio_coords.ndim - 1)
+            )
+        return denoise_ctx
+
+    def _denoise_timestep_kwargs(
+        self,
+        ts: torch.Tensor,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> dict[str, torch.Tensor]:
+        return {"timestep": ts, "sigma": ts}
+
+    def _build_transformer_kwargs(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+        *,
+        hidden_states: torch.Tensor,
+        audio_hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        audio_encoder_hidden_states: torch.Tensor,
+        encoder_attention_mask: torch.Tensor,
+        audio_encoder_attention_mask: torch.Tensor,
+        ts: torch.Tensor,
+    ) -> dict[str, Any]:
+        request_inputs = forward_ctx.request_inputs
+        return {
+            "hidden_states": hidden_states,
+            "audio_hidden_states": audio_hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "audio_encoder_hidden_states": audio_encoder_hidden_states,
+            **self._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx),
+            "encoder_attention_mask": encoder_attention_mask,
+            "audio_encoder_attention_mask": audio_encoder_attention_mask,
+            "num_frames": forward_ctx.latent_num_frames,
+            "height": forward_ctx.latent_height,
+            "width": forward_ctx.latent_width,
+            "fps": request_inputs.frame_rate,
+            "audio_num_frames": forward_ctx.padded_audio_num_frames,
+            "video_coords": denoise_ctx.video_coords,
+            "audio_coords": denoise_ctx.audio_coords,
+            "attention_kwargs": forward_ctx.attention_kwargs,
+            "return_dict": False,
+        }
+
+    def _step_denoised_latents(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+        noise_pred_video: torch.Tensor,
+        noise_pred_audio: torch.Tensor,
+        t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if forward_ctx.cfg_parallel_ready:
+            latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+                (noise_pred_video, noise_pred_audio),
+                (t, t),
+                (denoise_ctx.latents, denoise_ctx.audio_latents),
+                do_true_cfg=self.do_classifier_free_guidance,
+                per_request_scheduler=forward_ctx.video_audio_scheduler,
+            )
+            return self._synchronize_cfg_parallel_step_output(
+                (latents, audio_latents),
+                do_true_cfg=self.do_classifier_free_guidance,
+            )
+
+        latents = self.scheduler.step(noise_pred_video, t, denoise_ctx.latents, return_dict=False)[0]
+        audio_latents = forward_ctx.audio_scheduler.step(
+            noise_pred_audio,
+            t,
+            denoise_ctx.audio_latents,
+            return_dict=False,
+        )[0]
+        return latents, audio_latents
+
+    def _denoise_loop(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        request_inputs = forward_ctx.request_inputs
+        prompt_context = forward_ctx.prompt_context
+        guidance_scale = request_inputs.guidance_scale
+        audio_scheduler = forward_ctx.audio_scheduler
+
+        with self.progress_bar(total=len(forward_ctx.timesteps)) as pbar:
+            for i, t in enumerate(forward_ctx.timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+
+                if forward_ctx.cfg_parallel_ready:
+                    latent_model_input = denoise_ctx.latents.to(prompt_context.positive_connector_prompt_embeds.dtype)
+                    audio_latent_model_input = denoise_ctx.audio_latents.to(
+                        prompt_context.positive_connector_prompt_embeds.dtype
+                    )
+                    ts = t.expand(latent_model_input.shape[0])
+                    positive_kwargs = self._build_transformer_kwargs(
+                        forward_ctx,
+                        denoise_ctx,
+                        hidden_states=latent_model_input,
+                        audio_hidden_states=audio_latent_model_input,
+                        encoder_hidden_states=prompt_context.positive_connector_prompt_embeds,
+                        audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds,
+                        encoder_attention_mask=prompt_context.positive_connector_attention_mask,
+                        audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask,
+                        ts=ts,
+                    )
+                    negative_kwargs = {
+                        **positive_kwargs,
+                        "encoder_hidden_states": prompt_context.negative_connector_prompt_embeds,
+                        "audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds,
+                        "encoder_attention_mask": prompt_context.negative_connector_attention_mask,
+                        "audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask,
+                    }
+                    noise_pred_video, noise_pred_audio = self.predict_noise_with_parallel_cfg(
+                        true_cfg_scale=guidance_scale,
+                        positive_kwargs=positive_kwargs,
+                        negative_kwargs=negative_kwargs,
+                        cfg_normalize=False,
+                        video_latents=denoise_ctx.latents,
+                        audio_latents=denoise_ctx.audio_latents,
+                        video_sigma=self.scheduler.sigmas[i],
+                        audio_sigma=audio_scheduler.sigmas[i],
+                    )
+                else:
+                    latent_model_input = (
+                        torch.cat([denoise_ctx.latents] * 2)
+                        if self.do_classifier_free_guidance
+                        else denoise_ctx.latents
+                    )
+                    latent_model_input = latent_model_input.to(prompt_context.connector_prompt_embeds.dtype)
+                    audio_latent_model_input = (
+                        torch.cat([denoise_ctx.audio_latents] * 2)
+                        if self.do_classifier_free_guidance
+                        else denoise_ctx.audio_latents
+                    )
+                    audio_latent_model_input = audio_latent_model_input.to(prompt_context.connector_prompt_embeds.dtype)
+                    ts = t.expand(latent_model_input.shape[0])
+
+                    transformer_kwargs = self._build_transformer_kwargs(
+                        forward_ctx,
+                        denoise_ctx,
+                        hidden_states=latent_model_input,
+                        audio_hidden_states=audio_latent_model_input,
+                        encoder_hidden_states=prompt_context.connector_prompt_embeds,
+                        audio_encoder_hidden_states=prompt_context.connector_audio_prompt_embeds,
+                        encoder_attention_mask=prompt_context.connector_attention_mask,
+                        audio_encoder_attention_mask=prompt_context.connector_attention_mask,
+                        ts=ts,
+                    )
+                    with self._transformer_cache_context("cond_uncond"):
+                        noise_pred_video, noise_pred_audio = self.transformer(**transformer_kwargs)
+
+                    noise_pred_video = noise_pred_video.float()
+                    noise_pred_audio = noise_pred_audio.float()
+
+                    if self.do_classifier_free_guidance:
+                        noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2)
+                        noise_pred_video = self._combine_x0_space_cfg(
+                            denoise_ctx.latents,
+                            noise_pred_video_cond,
+                            noise_pred_video_uncond,
+                            self.scheduler.sigmas[i],
+                            guidance_scale,
+                        )
+
+                        noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2)
+                        noise_pred_audio = self._combine_x0_space_cfg(
+                            denoise_ctx.audio_latents,
+                            noise_pred_audio_cond,
+                            noise_pred_audio_uncond,
+                            audio_scheduler.sigmas[i],
+                            guidance_scale,
+                        )
+
+                denoise_ctx.latents, denoise_ctx.audio_latents = self._step_denoised_latents(
+                    forward_ctx,
+                    denoise_ctx,
+                    noise_pred_video,
+                    noise_pred_audio,
+                    t,
+                )
+                pbar.update()
+
+        return denoise_ctx.latents, denoise_ctx.audio_latents
+
+    def _unpack_and_denormalize_stage(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        latents: torch.Tensor,
+        audio_latents: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        latents = self._unpack_latents(
+            latents,
+            forward_ctx.latent_num_frames,
+            forward_ctx.latent_height,
+            forward_ctx.latent_width,
+            self.transformer_spatial_patch_size,
+            self.transformer_temporal_patch_size,
+        )
+        latents = self._denormalize_latents(
+            latents,
+            self.vae.latents_mean,
+            self.vae.latents_std,
+            self.vae.config.scaling_factor,
+        )
+
+        audio_latents = self._unpad_audio_latents(audio_latents, forward_ctx.original_audio_num_frames)
+        audio_latents = self._denormalize_audio_latents(
+            audio_latents,
+            self.audio_vae.latents_mean,
+            self.audio_vae.latents_std,
+        )
+        audio_latents = self._unpack_audio_latents(
+            audio_latents,
+            forward_ctx.original_audio_num_frames,
+            num_mel_bins=forward_ctx.latent_mel_bins,
+        )
+        return latents, audio_latents
+
+    def _decode_and_split(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        latents: torch.Tensor,
+        audio_latents: torch.Tensor,
+    ) -> list[DiffusionOutput]:
+        request_inputs = forward_ctx.request_inputs
+        return split_diffusion_output_by_request(
+            self._decode_output(
+                latents=latents,
+                audio_latents=audio_latents,
+                output_type=request_inputs.output_type,
+                connector_prompt_embeds=forward_ctx.prompt_context.connector_prompt_embeds,
+                generator=request_inputs.generator,
+                device=forward_ctx.device,
+                decode_timestep=request_inputs.decode_timestep,
+                decode_noise_scale=request_inputs.decode_noise_scale,
+                prompt_batch_size=forward_ctx.batch_size,
+            ),
+            forward_ctx.req,
+            num_outputs_per_prompt=forward_ctx.num_videos_per_prompt,
+        )
+
+    def _forward_impl(
+        self,
+        req: DiffusionRequestBatch,
+        request_inputs: _LTX23RequestInputs,
+        *,
+        noise_scale: float,
+        sigmas: list[float] | None,
+        timesteps: list[int] | None,
+        attention_kwargs: dict[str, Any] | None,
+        image: Any | None = None,
+    ) -> list[DiffusionOutput]:
+        self._check_forward_inputs(request_inputs, image=image)
+        cfg_parallel_ready = self._setup_forward_runtime(request_inputs, attention_kwargs)
+        device = self.device
+        prompt_context = self._prepare_prompt_context(
+            prompt=request_inputs.prompt,
+            negative_prompt=request_inputs.negative_prompt,
+            prompt_embeds=request_inputs.prompt_embeds,
+            negative_prompt_embeds=request_inputs.negative_prompt_embeds,
+            prompt_attention_mask=request_inputs.prompt_attention_mask,
+            negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
+            num_videos_per_prompt=request_inputs.num_videos_per_prompt,
+            max_sequence_length=request_inputs.max_sequence_length,
+        )
+
+        latent_num_frames, latent_height, latent_width = self._resolve_video_latent_dimensions(request_inputs)
+        latents, conditioning_mask = self._prepare_video_latents_stage(
+            request_inputs,
+            prompt_context,
+            device=device,
+            noise_scale=noise_scale,
+            image=image,
+        )
+        audio_latents, original_audio_num_frames, padded_audio_num_frames, latent_mel_bins = (
+            self._prepare_audio_latents_stage(
+                request_inputs,
+                prompt_context,
+                device=device,
+                noise_scale=noise_scale,
+            )
+        )
+        audio_scheduler, video_audio_scheduler, timesteps_tensor = self._prepare_scheduler_stage(
+            request_inputs,
+            device=device,
+            sigmas=sigmas,
+            timesteps=timesteps,
+            latent_num_frames=latent_num_frames,
+            latent_height=latent_height,
+            latent_width=latent_width,
+        )
+        forward_ctx = _LTX23ForwardContext(
+            req=req,
+            request_inputs=request_inputs,
+            prompt_context=prompt_context,
+            device=device,
+            cfg_parallel_ready=cfg_parallel_ready,
+            attention_kwargs=attention_kwargs,
+            latent_num_frames=latent_num_frames,
+            latent_height=latent_height,
+            latent_width=latent_width,
+            latent_mel_bins=latent_mel_bins,
+            original_audio_num_frames=original_audio_num_frames,
+            padded_audio_num_frames=padded_audio_num_frames,
+            timesteps=timesteps_tensor,
+            audio_scheduler=audio_scheduler,
+            video_audio_scheduler=video_audio_scheduler,
+        )
+        video_coords, audio_coords = self._prepare_rope_coords_stage(forward_ctx, latents, audio_latents)
+        denoise_ctx = _LTX23DenoiseContext(
+            latents=latents,
+            audio_latents=audio_latents,
+            video_coords=video_coords,
+            audio_coords=audio_coords,
+            conditioning_mask=conditioning_mask,
+        )
+        denoise_ctx = self._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
+        latents, audio_latents = self._denoise_loop(forward_ctx, denoise_ctx)
+        latents, audio_latents = self._unpack_and_denormalize_stage(forward_ctx, latents, audio_latents)
+        return self._decode_and_split(forward_ctx, latents, audio_latents)
+
     # ------------------------------------------------------------------
     # Forward pass
     # ------------------------------------------------------------------
@@ -864,446 +1794,39 @@ class LTX23Pipeline(
         attention_kwargs: dict[str, Any] | None = None,
         max_sequence_length: int | None = None,
     ) -> list[DiffusionOutput]:
-        # ---- Extract from request ----
-        sampling_params_list = req.sampling_params_list
-        common_sampling_params = sampling_params_list[0]
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
-        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
-            negative_prompt = None
-        elif req.prompts:
-            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
-
-        height = common_sampling_params.height or height or 512
-        width = common_sampling_params.width or width or 768
-        num_frames = common_sampling_params.num_frames or num_frames or 121
-        frame_rate = common_sampling_params.resolved_frame_rate or frame_rate or 24.0
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps or 40
-        # Enforce minimum of 2 timesteps for flow matching scheduler
-        if timesteps is None:
-            num_inference_steps = max(int(num_inference_steps), 2)
-        elif len(timesteps) < 2:
-            raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.")
-        num_videos_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_videos_per_prompt or 1
-        )
-        max_sequence_length = (
-            common_sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length
-        )
-
-        if common_sampling_params.guidance_scale_provided:
-            guidance_scale = common_sampling_params.guidance_scale
-
-        if generator is None:
-            generator = req.collate_request_generators(num_videos_per_prompt, generator)
-
-        latents = req.collate_request_tensors("latents", latents)
-        audio_latents = DiffusionRequestBatch.collate_tensors(
-            [_get_audio_latents_from_sampling(sampling) for sampling in sampling_params_list],
-            "audio_latents",
-            audio_latents,
-        )
-
-        # Override with pre-computed embeddings if provided in request
-        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
-            req.prompts,
-            {
-                "prompt_embeds": prompt_embeds,
-                "negative_prompt_embeds": negative_prompt_embeds,
-                "prompt_attention_mask": prompt_attention_mask,
-                "negative_prompt_attention_mask": negative_prompt_attention_mask,
-            },
-            field_aliases={
-                "prompt_attention_mask": ("prompt_attention_mask", "attention_mask"),
-                "negative_prompt_attention_mask": (
-                    "negative_prompt_attention_mask",
-                    "negative_attention_mask",
-                ),
-            },
-        )
-        prompt_embeds = prompt_fields["prompt_embeds"]
-        negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
-        prompt_attention_mask = prompt_fields["prompt_attention_mask"]
-        negative_prompt_attention_mask = prompt_fields["negative_prompt_attention_mask"]
-        if prompt_embeds is not None:
-            prompt = None
-        if negative_prompt_embeds is not None:
-            negative_prompt = None
-
-        if common_sampling_params.decode_timestep is not None:
-            decode_timestep = common_sampling_params.decode_timestep
-        if common_sampling_params.decode_noise_scale is not None:
-            decode_noise_scale = common_sampling_params.decode_noise_scale
-        if common_sampling_params.output_type is not None:
-            output_type = common_sampling_params.output_type
-
-        self.check_inputs(
-            prompt=prompt,
-            height=height,
-            width=width,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._interrupt = False
-        self._current_timestep = None
-        cfg_world_size = get_classifier_free_guidance_world_size()
-        if self.do_classifier_free_guidance and cfg_world_size not in (1, 2):
-            raise ValueError(
-                f"LTX23Pipeline supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}."
-            )
-        cfg_parallel_ready = self.do_classifier_free_guidance and cfg_world_size > 1
-
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self.device
-
-        # ---- Encode prompts ----
-        (
-            prompt_embeds,
-            prompt_attention_mask,
-            negative_prompt_embeds,
-            negative_prompt_attention_mask,
-        ) = self.encode_prompt(
+        request_inputs = self._resolve_request_inputs(
+            req,
             prompt=prompt,
             negative_prompt=negative_prompt,
-            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_inference_steps=num_inference_steps,
+            timesteps=timesteps,
+            guidance_scale=guidance_scale,
             num_videos_per_prompt=num_videos_per_prompt,
+            generator=generator,
+            latents=latents,
+            audio_latents=audio_latents,
             prompt_embeds=prompt_embeds,
             negative_prompt_embeds=negative_prompt_embeds,
             prompt_attention_mask=prompt_attention_mask,
             negative_prompt_attention_mask=negative_prompt_attention_mask,
+            decode_timestep=decode_timestep,
+            decode_noise_scale=decode_noise_scale,
+            output_type=output_type,
             max_sequence_length=max_sequence_length,
-            device=device,
         )
-
-        # ---- Connectors (LTX-2.3: padding_side API) ----
-        # Concatenate negative + positive embeddings BEFORE connector call,
-        # matching diffusers which calls connectors once with batch=2.
-        # This ensures batch-dependent operations produce identical results.
-        if self.do_classifier_free_guidance:
-            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
-            prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
-
-        tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left")
-        connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors(
-            prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side
-        )
-
-        positive_connector_prompt_embeds = connector_prompt_embeds
-        positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds
-        positive_connector_attention_mask = connector_attention_mask
-        negative_connector_prompt_embeds = None
-        negative_connector_audio_prompt_embeds = None
-        negative_connector_attention_mask = None
-        if self.do_classifier_free_guidance:
-            split_batch = batch_size * num_videos_per_prompt
-            negative_connector_prompt_embeds = connector_prompt_embeds[:split_batch]
-            positive_connector_prompt_embeds = connector_prompt_embeds[split_batch:]
-            negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:split_batch]
-            positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds[split_batch:]
-            negative_connector_attention_mask = connector_attention_mask[:split_batch]
-            positive_connector_attention_mask = connector_attention_mask[split_batch:]
-
-        # ---- Prepare latents ----
-        latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
-        latent_height = height // self.vae_spatial_compression_ratio
-        latent_width = width // self.vae_spatial_compression_ratio
-        if latents is not None and latents.ndim == 5:
-            _, _, latent_num_frames, latent_height, latent_width = latents.shape
-
-        num_channels_latents = self.transformer.config.in_channels
-        latents = self.prepare_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            num_frames,
-            noise_scale,
-            torch.float32,
-            device,
-            generator,
-            latents,
-        )
-
-        duration_s = num_frames / frame_rate
-        audio_latents_per_second = (
-            self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio)
-        )
-        audio_num_frames = round(duration_s * audio_latents_per_second)
-        audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, audio_latents)
-
-        num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64
-        latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio
-        num_channels_latents_audio = self.audio_vae.config.latent_channels if self.audio_vae is not None else 8
-        audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents=num_channels_latents_audio,
-            audio_latent_length=audio_num_frames,
-            num_mel_bins=num_mel_bins,
-            noise_scale=noise_scale,
-            dtype=torch.float32,
-            device=device,
-            generator=generator,
-            latents=audio_latents,
-        )
-
-        # ---- Scheduler setup ----
-        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
-        # Use max_image_seq_len (not actual video_sequence_length) for mu calculation,
-        # matching diffusers' LTX2Pipeline which hardcodes this value.
-        mu = calculate_shift(
-            self.scheduler.config.get("max_image_seq_len", 4096),
-            self.scheduler.config.get("base_image_seq_len", 1024),
-            self.scheduler.config.get("max_image_seq_len", 4096),
-            self.scheduler.config.get("base_shift", 0.95),
-            self.scheduler.config.get("max_shift", 2.05),
-        )
-        audio_scheduler = copy.deepcopy(self.scheduler)
-        video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_scheduler)
-        _ = retrieve_timesteps(audio_scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu)
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            device,
-            timesteps,
-            sigmas=sigmas,
-            mu=mu,
-        )
-        self._num_timesteps = len(timesteps)
-
-        # ---- RoPE coordinates ----
-        video_coords = self.transformer.rope.prepare_video_coords(
-            latents.shape[0],
-            latent_num_frames,
-            latent_height,
-            latent_width,
-            latents.device,
-            fps=frame_rate,
-        )
-        audio_coords = self.transformer.audio_rope.prepare_audio_coords(
-            audio_latents.shape[0],
-            padded_audio_num_frames,
-            audio_latents.device,
-        )
-
-        # ---- CFG: duplicate coords for single-rank batch=2 CFG ----
-        # Connector outputs are already batch=2 (neg+pos concatenated before connector call)
-        if self.do_classifier_free_guidance and not cfg_parallel_ready:
-            video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1))
-            audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1))
-
-        # ---- Denoising loop ----
-        # Uses x0-space CFG (delta formulation) matching diffusers' LTX2Pipeline.
-        # The velocity predictions are converted to x0, guidance is applied in x0
-        # space, then converted back to velocity for the scheduler step.
-        with self.progress_bar(total=len(timesteps)) as pbar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                self._current_timestep = t
-
-                if cfg_parallel_ready:
-                    latent_model_input = latents.to(positive_connector_prompt_embeds.dtype)
-                    audio_latent_model_input = audio_latents.to(positive_connector_prompt_embeds.dtype)
-                    ts = t.expand(latent_model_input.shape[0])
-                    positive_kwargs = {
-                        "hidden_states": latent_model_input,
-                        "audio_hidden_states": audio_latent_model_input,
-                        "encoder_hidden_states": positive_connector_prompt_embeds,
-                        "audio_encoder_hidden_states": positive_connector_audio_prompt_embeds,
-                        "timestep": ts,
-                        "sigma": ts,
-                        "encoder_attention_mask": positive_connector_attention_mask,
-                        "audio_encoder_attention_mask": positive_connector_attention_mask,
-                        "num_frames": latent_num_frames,
-                        "height": latent_height,
-                        "width": latent_width,
-                        "fps": frame_rate,
-                        "audio_num_frames": padded_audio_num_frames,
-                        "video_coords": video_coords,
-                        "audio_coords": audio_coords,
-                        "attention_kwargs": attention_kwargs,
-                        "return_dict": False,
-                    }
-                    negative_kwargs = {
-                        **positive_kwargs,
-                        "encoder_hidden_states": negative_connector_prompt_embeds,
-                        "audio_encoder_hidden_states": negative_connector_audio_prompt_embeds,
-                        "encoder_attention_mask": negative_connector_attention_mask,
-                        "audio_encoder_attention_mask": negative_connector_attention_mask,
-                    }
-                    noise_pred_video, noise_pred_audio = self.predict_noise_with_parallel_cfg(
-                        true_cfg_scale=guidance_scale,
-                        positive_kwargs=positive_kwargs,
-                        negative_kwargs=negative_kwargs,
-                        cfg_normalize=False,
-                        video_latents=latents,
-                        audio_latents=audio_latents,
-                        video_sigma=self.scheduler.sigmas[i],
-                        audio_sigma=audio_scheduler.sigmas[i],
-                    )
-
-                    latents, audio_latents = self.scheduler_step_maybe_with_cfg(
-                        (noise_pred_video, noise_pred_audio),
-                        (t, t),
-                        (latents, audio_latents),
-                        do_true_cfg=self.do_classifier_free_guidance,
-                        per_request_scheduler=video_audio_scheduler,
-                    )
-                    latents, audio_latents = self._synchronize_cfg_parallel_step_output(
-                        (latents, audio_latents),
-                        do_true_cfg=self.do_classifier_free_guidance,
-                    )
-                else:
-                    latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
-                    latent_model_input = latent_model_input.to(connector_prompt_embeds.dtype)
-                    audio_latent_model_input = (
-                        torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents
-                    )
-                    audio_latent_model_input = audio_latent_model_input.to(connector_prompt_embeds.dtype)
-                    ts = t.expand(latent_model_input.shape[0])
-
-                    with self._transformer_cache_context("cond_uncond"):
-                        noise_pred_video, noise_pred_audio = self.transformer(
-                            hidden_states=latent_model_input,
-                            audio_hidden_states=audio_latent_model_input,
-                            encoder_hidden_states=connector_prompt_embeds,
-                            audio_encoder_hidden_states=connector_audio_prompt_embeds,
-                            timestep=ts,
-                            sigma=ts,  # LTX-2.3: sigma for prompt_adaln
-                            encoder_attention_mask=connector_attention_mask,
-                            audio_encoder_attention_mask=connector_attention_mask,
-                            num_frames=latent_num_frames,
-                            height=latent_height,
-                            width=latent_width,
-                            fps=frame_rate,
-                            audio_num_frames=padded_audio_num_frames,
-                            video_coords=video_coords,
-                            audio_coords=audio_coords,
-                            attention_kwargs=attention_kwargs,
-                            return_dict=False,
-                        )
-
-                    noise_pred_video = noise_pred_video.float()
-                    noise_pred_audio = noise_pred_audio.float()
-
-                    if self.do_classifier_free_guidance:
-                        noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2)
-                        noise_pred_video = self._combine_x0_space_cfg(
-                            latents,
-                            noise_pred_video_cond,
-                            noise_pred_video_uncond,
-                            self.scheduler.sigmas[i],
-                            guidance_scale,
-                        )
-
-                        noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2)
-                        noise_pred_audio = self._combine_x0_space_cfg(
-                            audio_latents,
-                            noise_pred_audio_cond,
-                            noise_pred_audio_uncond,
-                            audio_scheduler.sigmas[i],
-                            guidance_scale,
-                        )
-
-                    latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0]
-                    audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0]
-
-                pbar.update()
-
-        # ---- Unpack and denormalize ----
-        latents = self._unpack_latents(
-            latents,
-            latent_num_frames,
-            latent_height,
-            latent_width,
-            self.transformer_spatial_patch_size,
-            self.transformer_temporal_patch_size,
-        )
-        latents = self._denormalize_latents(
-            latents,
-            self.vae.latents_mean,
-            self.vae.latents_std,
-            self.vae.config.scaling_factor,
-        )
-
-        audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames)
-        audio_latents = self._denormalize_audio_latents(
-            audio_latents,
-            self.audio_vae.latents_mean,
-            self.audio_vae.latents_std,
-        )
-        audio_latents = self._unpack_audio_latents(
-            audio_latents,
-            original_audio_num_frames,
-            num_mel_bins=latent_mel_bins,
-        )
-
-        # ---- Decode ----
-        if output_type == "latent":
-            video = latents
-            audio = audio_latents
-        else:
-            latents = latents.to(connector_prompt_embeds.dtype)
-
-            if not self.vae.config.timestep_conditioning:
-                timestep_decode = None
-            else:
-                noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
-                if not isinstance(decode_timestep, list):
-                    decode_timestep = [decode_timestep] * batch_size
-                if decode_noise_scale is None:
-                    decode_noise_scale = decode_timestep
-                elif not isinstance(decode_noise_scale, list):
-                    decode_noise_scale = [decode_noise_scale] * batch_size
-                timestep_decode = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
-                decode_noise_scale_t = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
-                    :, None, None, None, None
-                ]
-                latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise
-
-            latents = latents.to(self.vae.dtype)
-            video = self.vae.decode(latents, timestep_decode, return_dict=False)[0]
-            if video.numel() > 0:
-                video = self.video_processor.postprocess_video(video, output_type=output_type)
-
-            audio_latents = audio_latents.to(self.audio_vae.dtype)
-            generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
-
-            audio = self.vocoder(generated_mel_spectrograms)
-
-        return split_diffusion_output_by_request(
-            DiffusionOutput(
-                output=(video, audio),
-                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
-            ),
+        return self._forward_impl(
             req,
-            num_outputs_per_prompt=num_videos_per_prompt,
+            request_inputs,
+            noise_scale=noise_scale,
+            sigmas=sigmas,
+            timesteps=timesteps,
+            attention_kwargs=attention_kwargs,
         )
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)
         return loader.load_weights(weights)
-
-
-class LTX23ImageToVideoPipeline(nn.Module):
-    """LTX-2.3 image-to-video pipeline placeholder."""
-
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        super().__init__()
-        raise NotImplementedError(
-            "LTX23ImageToVideoPipeline is not yet implemented. "
-            "Use LTX23Pipeline for single-stage text-to-video generation."
-        )
diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py
index d30ef5464..bcccd7596 100644
--- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py
+++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py
@@ -1,16 +1,452 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Re-exports for LTX-2.3 I2V pipeline variants.
+"""LTX-2.3 image-to-video pipeline."""
 
-The registry loads pipeline classes by (mod_folder, mod_relname, cls_name).
-This module exposes the I2V class names so the registry can find them.
-"""
+from __future__ import annotations
+
+from typing import Any
+
+import PIL.Image
+import torch
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 from .pipeline_ltx2_3 import (
-    LTX23ImageToVideoPipeline,
-    get_ltx2_post_process_func,  # noqa: F401 - loaded by registry via getattr
+    LTX23Pipeline,
+    _LTX23DenoiseContext,
+    _LTX23ForwardContext,
+    _LTX23PromptContext,
+    _LTX23RequestInputs,
+    get_ltx2_post_process_func,
 )
+from .pipeline_ltx2_image2video import LTX2ImageToVideoPipeline, _I2VVideoAudioScheduler
+
+
+class LTX23ImageToVideoPipeline(LTX23Pipeline):
+    """LTX-2.3 image-to-video pipeline.
+
+    This keeps the LTX-2.3 prompt connector, x0-space CFG, sigma prompt
+    modulation, and audio branch semantics from ``LTX23Pipeline`` while
+    reusing the existing LTX image-conditioning contract: the first video
+    latent frame is encoded from the input image and remains fixed during
+    denoising.
+    """
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear")
+
+    support_image_input = True
+
+    _normalize_latents = staticmethod(LTX2ImageToVideoPipeline._normalize_latents)
+    _create_noised_state = staticmethod(LTX2ImageToVideoPipeline._create_noised_state)
+
+    @staticmethod
+    def _resolve_single_prompt_image(raw_image: Any) -> Any:
+        if isinstance(raw_image, list):
+            if len(raw_image) != 1:
+                raise ValueError(
+                    "LTX-2.3 I2V prompt dictionaries support exactly one image per prompt. "
+                    "Pass one image per prompt for batched I2V requests."
+                )
+            return raw_image[0]
+        return raw_image
+
+    @staticmethod
+    def _resolve_additional_image(additional: dict[str, Any]) -> Any:
+        raw_image = additional.get("preprocessed_image")
+        if raw_image is None:
+            raw_image = additional.get("pixel_values")
+        if raw_image is None:
+            raw_image = additional.get("image")
+        return raw_image
+
+    def prepare_latents(
+        self,
+        image: torch.Tensor | None = None,
+        batch_size: int = 1,
+        num_channels_latents: int = 128,
+        height: int = 512,
+        width: int = 768,
+        num_frames: int = 121,
+        noise_scale: float = 0.0,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare I2V latents and the first-frame conditioning mask.
+
+        If caller-provided latents are used without an image, the latents must
+        already represent the full video state including the conditioning first
+        frame. Packed 3D latents are assumed to be in transformer token layout.
+        """
+        height = height // self.vae_spatial_compression_ratio
+        width = width // self.vae_spatial_compression_ratio
+        num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
+
+        shape = (batch_size, num_channels_latents, num_frames, height, width)
+        mask_shape = (batch_size, 1, num_frames, height, width)
+
+        if latents is not None:
+            if latents.ndim == 5:
+                batch_size, _, num_frames, height, width = latents.shape
+                mask_shape = (batch_size, 1, num_frames, height, width)
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+
+                latents = self._normalize_latents(
+                    latents,
+                    self.vae.latents_mean,
+                    self.vae.latents_std,
+                    self.vae.config.scaling_factor,
+                )
+                latents = self._create_noised_state(latents, noise_scale * (1 - conditioning_mask), generator)
+                latents = self._pack_latents(
+                    latents,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
+            else:
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+
+            conditioning_mask = self._pack_latents(
+                conditioning_mask,
+                self.transformer_spatial_patch_size,
+                self.transformer_temporal_patch_size,
+            ).squeeze(-1)
+            if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
+                raise ValueError(
+                    "Provided `latents` tensor has shape"
+                    f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
+                )
+            return latents.to(device=device, dtype=dtype), conditioning_mask
+
+        if image is None:
+            raise ValueError("`image` must be provided when `latents` is None.")
+
+        image_batch_size = image.shape[0]
+        if image_batch_size == 0:
+            raise ValueError("`image` batch is empty.")
+        if batch_size % image_batch_size != 0:
+            raise ValueError(
+                f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) "
+                "for image-to-video outputs."
+            )
+        num_videos_per_prompt = batch_size // image_batch_size
+
+        if isinstance(generator, list):
+            if len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective"
+                    f" batch size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+            image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)]
+            init_latents = [
+                retrieve_latents(self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), image_generators[i], "argmax")
+                for i in range(image_batch_size)
+            ]
+        else:
+            init_latents = [
+                retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator, "argmax") for img in image
+            ]
+
+        init_latents = torch.cat(init_latents, dim=0).to(dtype)
+        if num_videos_per_prompt > 1:
+            init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
+        init_latents = self._normalize_latents(
+            init_latents,
+            self.vae.latents_mean,
+            self.vae.latents_std,
+        )
+        init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
+
+        conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype)
+        conditioning_mask[:, :, 0] = 1.0
+
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
+
+        conditioning_mask = self._pack_latents(
+            conditioning_mask,
+            self.transformer_spatial_patch_size,
+            self.transformer_temporal_patch_size,
+        ).squeeze(-1)
+        latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size)
+
+        return latents, conditioning_mask
+
+    def check_inputs(
+        self,
+        image,
+        height,
+        width,
+        prompt,
+        latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if image is None and latents is None:
+            raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.")
+        super().check_inputs(
+            prompt=prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+        )
+
+    _step_video_latents_i2v = LTX2ImageToVideoPipeline._step_video_latents_i2v
+
+    def _resolve_request_image(
+        self,
+        req: DiffusionRequestBatch,
+        image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None,
+        request_inputs: _LTX23RequestInputs,
+    ) -> PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None:
+        if image is not None or not req.prompts:
+            return image
+
+        raw_images = []
+        for prompt_item in req.prompts:
+            if isinstance(prompt_item, str):
+                raw_image = None
+            else:
+                multi_modal_data = prompt_item.get("multi_modal_data") or {}
+                raw_image = multi_modal_data.get("image")
+                if raw_image is None:
+                    additional = prompt_item.get("additional_information") or {}
+                    raw_image = self._resolve_additional_image(additional)
+            raw_image = self._resolve_single_prompt_image(raw_image)
+            if isinstance(raw_image, str):
+                raw_image = PIL.Image.open(raw_image).convert("RGB")
+            raw_images.append(raw_image)
+
+        if any(img is None for img in raw_images) and request_inputs.latents is None:
+            raise ValueError("Image is required for LTX-2.3 I2V generation.")
+        if len(raw_images) == 1:
+            return raw_images[0]
+        if raw_images:
+            return raw_images
+        return image
+
+    def _check_forward_inputs(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        image: Any | None = None,
+    ) -> None:
+        self.check_inputs(
+            image=image,
+            height=request_inputs.height,
+            width=request_inputs.width,
+            prompt=request_inputs.prompt,
+            latents=request_inputs.latents,
+            prompt_embeds=request_inputs.prompt_embeds,
+            negative_prompt_embeds=request_inputs.negative_prompt_embeds,
+            prompt_attention_mask=request_inputs.prompt_attention_mask,
+            negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
+        )
+
+    def _resolve_video_latent_dimensions(
+        self,
+        request_inputs: _LTX23RequestInputs,
+    ) -> tuple[int, int, int]:
+        latent_num_frames = (request_inputs.num_frames - 1) // self.vae_temporal_compression_ratio + 1
+        latent_height = request_inputs.height // self.vae_spatial_compression_ratio
+        latent_width = request_inputs.width // self.vae_spatial_compression_ratio
+        if request_inputs.latents is not None:
+            if request_inputs.latents.ndim == 5:
+                _, _, latent_num_frames, latent_height, latent_width = request_inputs.latents.shape
+            elif request_inputs.latents.ndim != 3:
+                raise ValueError(
+                    f"Provided `latents` tensor has shape {request_inputs.latents.shape}, but the expected shape is "
+                    "either [batch_size, seq_len, num_features] or "
+                    "[batch_size, latent_dim, latent_frames, latent_height, latent_width]."
+                )
+        return latent_num_frames, latent_height, latent_width
+
+    def _prepare_video_latents_stage(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        prompt_context: _LTX23PromptContext,
+        *,
+        device: torch.device,
+        noise_scale: float,
+        image: Any | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if request_inputs.latents is None:
+            if isinstance(image, torch.Tensor):
+                if image.ndim == 3:
+                    image = image.unsqueeze(0)
+            elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor):
+                image = torch.stack(image, dim=0)
+            else:
+                image = self.video_processor.preprocess(
+                    image,
+                    height=request_inputs.height,
+                    width=request_inputs.width,
+                )
+            image = image.to(device=device, dtype=prompt_context.positive_connector_prompt_embeds.dtype)
+
+        num_channels_latents = self.transformer.config.in_channels
+        return self.prepare_latents(
+            image,
+            prompt_context.batch_size * request_inputs.num_videos_per_prompt,
+            num_channels_latents,
+            request_inputs.height,
+            request_inputs.width,
+            request_inputs.num_frames,
+            noise_scale,
+            torch.float32,
+            device,
+            request_inputs.generator,
+            request_inputs.latents,
+        )
+
+    def _make_video_audio_scheduler(
+        self,
+        audio_scheduler: Any,
+        latent_num_frames: int,
+        latent_height: int,
+        latent_width: int,
+    ) -> Any:
+        return _I2VVideoAudioScheduler(
+            self,
+            audio_scheduler,
+            latent_num_frames,
+            latent_height,
+            latent_width,
+        )
+
+    def _prepare_denoise_context_for_cfg(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> _LTX23DenoiseContext:
+        denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
+        if denoise_ctx.conditioning_mask is None:
+            raise ValueError("I2V denoising requires a conditioning mask.")
+        if self.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready:
+            denoise_ctx.conditioning_mask_for_model = torch.cat(
+                [denoise_ctx.conditioning_mask, denoise_ctx.conditioning_mask]
+            )
+        else:
+            denoise_ctx.conditioning_mask_for_model = denoise_ctx.conditioning_mask
+        return denoise_ctx
+
+    def _denoise_timestep_kwargs(
+        self,
+        ts: torch.Tensor,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> dict[str, torch.Tensor]:
+        conditioning_mask = (
+            denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model
+        )
+        if conditioning_mask is None:
+            raise ValueError("I2V denoising requires a conditioning mask.")
+        return {
+            "timestep": ts.unsqueeze(-1) * (1 - conditioning_mask),
+            "audio_timestep": ts,
+            "sigma": ts,
+        }
+
+    def _step_denoised_latents(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+        noise_pred_video: torch.Tensor,
+        noise_pred_audio: torch.Tensor,
+        t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+            (noise_pred_video, noise_pred_audio),
+            (t, t),
+            (denoise_ctx.latents, denoise_ctx.audio_latents),
+            do_true_cfg=self.do_classifier_free_guidance,
+            per_request_scheduler=forward_ctx.video_audio_scheduler,
+        )
+        return self._synchronize_cfg_parallel_step_output(
+            (latents, audio_latents),
+            do_true_cfg=self.do_classifier_free_guidance,
+        )
+
+    @torch.no_grad()
+    def forward(
+        self,
+        req: DiffusionRequestBatch,
+        image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None = None,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_frames: int | None = None,
+        frame_rate: float | None = None,
+        num_inference_steps: int | None = None,
+        sigmas: list[float] | None = None,
+        timesteps: list[int] | None = None,
+        guidance_scale: float = 4.0,
+        noise_scale: float = 0.0,
+        num_videos_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        audio_latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        negative_prompt_attention_mask: torch.Tensor | None = None,
+        decode_timestep: float | list[float] = 0.0,
+        decode_noise_scale: float | list[float] | None = None,
+        output_type: str = "np",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        max_sequence_length: int | None = None,
+    ) -> list[DiffusionOutput]:
+        request_inputs = self._resolve_request_inputs(
+            req,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_inference_steps=num_inference_steps,
+            timesteps=timesteps,
+            guidance_scale=guidance_scale,
+            num_videos_per_prompt=num_videos_per_prompt,
+            generator=generator,
+            latents=latents,
+            audio_latents=audio_latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            decode_timestep=decode_timestep,
+            decode_noise_scale=decode_noise_scale,
+            output_type=output_type,
+            max_sequence_length=max_sequence_length,
+        )
+        image = self._resolve_request_image(req, image, request_inputs)
+        return self._forward_impl(
+            req,
+            request_inputs,
+            noise_scale=noise_scale,
+            sigmas=sigmas,
+            timesteps=timesteps,
+            attention_kwargs=attention_kwargs,
+            image=image,
+        )
+
 
 __all__ = [
     "LTX23ImageToVideoPipeline",
diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py
index 0c312aa7e..614c8268c 100644
--- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py
+++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py
@@ -9,6 +9,7 @@ from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
 from vllm.config import VllmConfig
+from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
@@ -19,6 +20,8 @@ from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DMode
 from .rope_real import RotaryPosEmbedReal
 from .schedulers import FlowMatchEulerDiscreteScheduler
 
+logger = init_logger(__name__)
+
 
 class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     """
@@ -60,14 +63,14 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)
 
         # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
-        llm_hidden_size = int(
-            getattr(
-                getattr(self.config.llm_config, "text_config", None),
-                "hidden_size",
-                0,
-            )
-            or 0
-        )
+        llm_hidden_size = 0
+        text_config = self.config.get_text_config()
+        if text_config is None:
+            logger.warning("No text config; failed to infer llm_hidden_size.")
+        elif not hasattr(text_config, "hidden_size"):
+            logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.")
+        else:
+            llm_hidden_size = int(text_config.hidden_size or 0)
         if llm_hidden_size <= 0:
             raise ValueError(
                 "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"
diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py
index a55a79987..c67f8cb24 100644
--- a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py
+++ b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py
@@ -24,7 +24,6 @@ from __future__ import annotations
 
 import logging
 import os
-from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
 
@@ -49,7 +48,9 @@ from vllm_omni.diffusion.models.ming_flash_omni.ming_zimage_transformer import (
     MingZImageTransformer2DModel,
 )
 from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -58,38 +59,6 @@ from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenCon
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class _ZPipelineSamplingParams:
-    """Typed shim satisfying the attributes ZImagePipeline.forward reads
-    from ``req.sampling_params``.  Unlike SimpleNamespace this fails loudly
-    at construction if a required field is missing."""
-
-    height: int
-    width: int
-    num_inference_steps: int
-    guidance_scale: float
-    generator: torch.Generator | None = None
-    strength: float | None = None
-    sigmas: list[float] | None = None
-    max_sequence_length: int = 512
-    guidance_rescale: float | None = None
-    num_outputs_per_prompt: int = 1
-
-
-@dataclass
-class _ZPipelineRequest:
-    """Typed shim satisfying the attributes ZImagePipeline.forward reads
-    from ``req``."""
-
-    request_id: str
-    sampling_params: _ZPipelineSamplingParams
-    prompts: list[dict[str, Any]] = field(default_factory=lambda: [{"prompt": "", "negative_prompt": ""}])
-
-    @property
-    def num_reqs(self) -> int:
-        return 1
-
-
 class MingImagePipeline(ZImagePipeline):
     """Ming-flash-omni-2.0 text-to-image diffusion pipeline.
 
@@ -123,6 +92,10 @@ class MingImagePipeline(ZImagePipeline):
         self.device = self._execution_device  # Ming convention alias
         self._dtype = dtype
 
+        # Request-scoped conditioning handed to the inherited encode_prompt override
+        self._pending_prompt_embeds: list[torch.Tensor] | None = None
+        self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None
+
         # Ming's per-checkpoint image-gen configuration. We cannot rely on
         # ``od_config.hf_config.image_gen_config`` because the diffusion
         # stage is started with ``hf_config_name: thinker_config`` (the
@@ -273,6 +246,21 @@ class MingImagePipeline(ZImagePipeline):
         latent = self.vae.encode(ref).latent_dist.mode()
         return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
 
+    def encode_prompt(self, *args, **kwargs):  # noqa: ARG002
+        """Return Ming's precomputed conditioning instead of encoding text.
+
+        NOTE: Ming has no Z-Image text_encoder; its conditioning (cap_feats, optionally ByT5-augmented)
+        is computed in forward and stashed on `self._pending_*` immediately before
+        the `super().forward` call, so we simply hand it back here.
+        """
+        if self._pending_prompt_embeds is None:
+            raise RuntimeError(
+                "MingImagePipeline.encode_prompt called without pending "
+                "conditioning; it must run within MingImagePipeline.forward."
+            )
+
+        return self._pending_prompt_embeds, self._pending_negative_prompt_embeds
+
     # ------------------------------------------------------------------
     # Forward
     # ------------------------------------------------------------------
@@ -416,17 +404,27 @@ class MingImagePipeline(ZImagePipeline):
             negative_prompt_embeds = [negative_cap_feats[i] for i in range(negative_cap_feats.shape[0])]
         else:
             negative_prompt_embeds = [self.condition_encoder.zero_negative(e) for e in prompt_embeds]
+        self._pending_prompt_embeds = prompt_embeds
+        self._pending_negative_prompt_embeds = negative_prompt_embeds
 
-        z_sp = _ZPipelineSamplingParams(
+        # Build a real single-request DiffusionRequestBatch;
+        # the prompt text itself is a neutral placeholder here.
+        z_sp = OmniDiffusionSamplingParams(
             height=height,
             width=width,
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             generator=generator,
+            output_type="pt",
         )
-        z_req = _ZPipelineRequest(
-            request_id=req.request_id or "ming-imagegen",
-            sampling_params=z_sp,
+        z_req = DiffusionRequestBatch(
+            requests=[
+                OmniDiffusionRequest(
+                    prompt={"prompt": ""},
+                    sampling_params=z_sp,
+                    request_id=req.request_id or "ming-imagegen",
+                )
+            ]
         )
 
         # Reference image (img2img) → VAE-encoded latent published on the
@@ -446,29 +444,14 @@ class MingImagePipeline(ZImagePipeline):
             None if ref_latent is None else tuple(ref_latent.shape),
         )
         try:
-            outputs = super().forward(
-                z_req,
-                prompt=None,
-                height=height,
-                width=width,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale,
-                prompt_embeds=prompt_embeds,
-                negative_prompt_embeds=negative_prompt_embeds,
-                generator=generator,
-                output_type="pt",
-                return_dict=True,
-            )
+            outputs: DiffusionOutput = super().forward(z_req)
         finally:
             set_forward_context_ref_latent(None)
+            # Drop request-scoped conditioning so we don't retain GPU tensors.
+            self._pending_prompt_embeds = None
+            self._pending_negative_prompt_embeds = None
 
-        output = outputs[0]
-        if hasattr(output, "output") and output.output is not None:
-            raw = output.output
-        elif hasattr(output, "images"):
-            raw = output.images
-        else:
-            raw = output
+        raw = outputs.output
         if not isinstance(raw, torch.Tensor):
             raise RuntimeError(f"ZImagePipeline returned non-tensor output: {type(raw).__name__}")
         if logger.isEnabledFor(logging.DEBUG):
@@ -478,7 +461,7 @@ class MingImagePipeline(ZImagePipeline):
                 raw.float().min().item(),
                 raw.float().max().item(),
             )
-        return DiffusionOutput(output=raw)
+        return outputs
 
 
 # ----------------------------------------------------------------------
diff --git a/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py b/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py
index 3f03563a1..8c2c9db88 100644
--- a/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py
+++ b/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py
@@ -22,6 +22,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
@@ -808,6 +812,36 @@ class OmniGen2TransformerBlock(nn.Module):
         return hidden_states
 
 
+class _OmniGen2SPInputBoundary(nn.Module):
+    """SP input-sharding boundary marker for ``OmniGen2Transformer2DModel._sp_plan``.
+
+    Resets ``sp_original_seq_len`` on every call: OmniGen2 runs the transformer
+    once per CFG branch and the unconditional / conditional branches have
+    different joint sequence lengths (text prefix differs). The SP framework's
+    auto-pad records the length only on the first call, so without this reset
+    a stale length from an earlier branch truncates the gather of later ones.
+    """
+
+    def forward(self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            is_forward_context_available,
+        )
+
+        if is_forward_context_available():
+            ctx = get_forward_context()
+            ctx.sp_original_seq_len = None
+            ctx.sp_padding_size = 0
+        return hidden_states, rotary_emb
+
+
+class _OmniGen2SPOutputBoundary(nn.Module):
+    """SP output-gathering boundary marker for ``OmniGen2Transformer2DModel._sp_plan``."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states
+
+
 class OmniGen2Transformer2DModel(nn.Module):
     """
     OmniGen2 Transformer 2D Model.
@@ -836,6 +870,17 @@ class OmniGen2Transformer2DModel(nn.Module):
         timestep_scale: Scale factor for timestep embeddings
     """
 
+    # ``attention_mask`` is intentionally omitted from the plan: OmniGen2Attention
+    # accepts but never reads it. ``auto_pad=True`` requires Ulysses SP — Ring SP
+    # cannot use attention masks.
+    _sp_plan = {
+        "sp_input_boundary": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+        },
+        "sp_output_boundary": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
     def __init__(
         self,
         patch_size: int = 2,
@@ -983,6 +1028,9 @@ class OmniGen2Transformer2DModel(nn.Module):
         # Add learnable embeddings to distinguish different images
         self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size))  # support max 5 ref images
 
+        self.sp_input_boundary = _OmniGen2SPInputBoundary()
+        self.sp_output_boundary = _OmniGen2SPOutputBoundary()
+
     def img_patch_embed_and_refine(
         self,
         hidden_states,
@@ -1251,9 +1299,13 @@ class OmniGen2Transformer2DModel(nn.Module):
 
         hidden_states = joint_hidden_states
 
+        hidden_states, rotary_emb = self.sp_input_boundary(hidden_states, rotary_emb)
+
         for layer in self.layers:
             hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
 
+        hidden_states = self.sp_output_boundary(hidden_states)
+
         # 4. Output norm & projection
         hidden_states = self.norm_out(hidden_states, temb)
 
diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
index 371af3f34..0baad965b 100644
--- a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
+++ b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py
@@ -18,7 +18,7 @@
 import inspect
 import json
 import os
-from collections.abc import Callable, Iterable
+from collections.abc import Iterable
 from typing import Any, ClassVar
 
 import numpy as np
@@ -551,117 +551,33 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        guidance_scale: float = 5.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        num_images_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.FloatTensor | None = None,
-        negative_prompt_embeds: torch.FloatTensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 256,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            negative_prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                not greater than `1`).
-            guidance_scale (`float`, *optional*, defaults to 1.0):
-                True classifier-free guidance (guidance scale) is enabled when `guidance_scale` > 1 and
-                `negative_prompt` is provided.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`list[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.FloatTensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
-            joint_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`list`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.ovis_image.OvisImagePipelineOutput`] or `tuple`:
-            [`~pipelines.ovis_image.OvisImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
-            returning a tuple, the first element is a list with the generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
         height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        generator = req.sampling_params.generator or generator
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        generator = req.sampling_params.generator
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        max_sequence_length = req.sampling_params.max_sequence_length or 256
+
+        joint_attention_kwargs: dict[str, Any] | None = None
+        output_type = req.sampling_params.output_type or "pil"
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # Steps:
         # 1. Check Inputs
@@ -691,12 +607,7 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = len(prompt)
 
         do_classifier_free_guidance = guidance_scale > 1.0
         prompt_embeds, text_ids = self.encode_prompt(
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
index 43cd4bdc4..986c2ef86 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
@@ -73,8 +73,19 @@ def get_qwen_image_post_process_func(
     image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
 
     def post_process_func(
-        images: torch.Tensor,
+        images: torch.Tensor | dict[str, Any],
     ):
+        if isinstance(images, dict) and isinstance(images.get("payload"), dict):
+            payload = dict(images["payload"])
+            image_payload = payload.get("image")
+            if image_payload is None:
+                raise ValueError("Qwen-Image postprocess expected payload['image'] in output envelope.")
+            payload["image"] = image_processor.postprocess(image_payload)
+            metadata = images.get("metadata") or {}
+            return {
+                "payload": payload,
+                "metadata": metadata if isinstance(metadata, dict) else {},
+            }
         return image_processor.postprocess(images)
 
     return post_process_func
@@ -978,58 +989,34 @@ class QwenImagePipeline(
 
         height = state.sampling.height or self.default_sample_size * self.vae_scale_factor
         width = state.sampling.width or self.default_sample_size * self.vae_scale_factor
-        output_type = kwargs.get("output_type", "pil")
+        output_type = kwargs.get("output_type") or state.sampling.output_type or "pil"
 
         return self._decode_latents(state.latents, height, width, output_type)
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | 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 = 1024,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
         extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
-        prompt = extracted_prompt or prompt
+        prompt = extracted_prompt
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
         height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
-        sigmas = common_sampling_params.sigmas or sigmas
-        max_sequence_length = common_sampling_params.max_sequence_length or max_sequence_length
+        num_inference_steps = common_sampling_params.num_inference_steps or 50
+        sigmas = common_sampling_params.sigmas
+        max_sequence_length = common_sampling_params.max_sequence_length or 1024
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
-                "prompt_embeds": prompt_embeds,
-                "prompt_embeds_mask": prompt_embeds_mask,
-                "negative_prompt_embeds": negative_prompt_embeds,
-                "negative_prompt_embeds_mask": negative_prompt_embeds_mask,
+                "prompt_embeds": None,
+                "prompt_embeds_mask": None,
+                "negative_prompt_embeds": None,
+                "negative_prompt_embeds_mask": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
@@ -1040,9 +1027,16 @@ class QwenImagePipeline(
             prompt = None
         if negative_prompt_embeds is not None:
             negative_prompt = None
-        true_cfg_scale = common_sampling_params.true_cfg_scale or true_cfg_scale
+        true_cfg_scale = common_sampling_params.true_cfg_scale or 4.0
         if common_sampling_params.guidance_scale_provided:
             guidance_scale = common_sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
+
+        latents = req.collate_request_tensors("latents", None)
+        output_type = common_sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         ctx = self._prepare_generation_context(
             prompt=prompt,
@@ -1087,6 +1081,7 @@ class QwenImagePipeline(
         )
 
         self._current_timestep = None
+
         result = self._decode_latents(latents, height, width, output_type)
         return split_diffusion_output_by_request(
             result,
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
index e594d5ad8..3d6e937c7 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py
@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -677,30 +677,7 @@ class QwenImageEditPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        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 = 1024,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image editing."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
@@ -731,32 +708,28 @@ class QwenImageEditPipeline(
             height = req.sampling_params.height
             width = req.sampling_params.width
         else:
-            # fallback to run pre-processing in pipeline (debug only)
-            image_size = image[0].size if isinstance(image, list) else image.size
-            calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
-            height = height or calculated_height
-            width = width or calculated_width
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
-                image = self.image_processor.resize(image, calculated_height, calculated_width)
-                prompt_image = image
-                image = self.image_processor.preprocess(image, calculated_height, calculated_width)
-                image = image.unsqueeze(2)
-
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. check inputs
         # 2. encode prompts
@@ -783,16 +756,9 @@ class QwenImageEditPipeline(
         self._current_timestep = None
         self._interrupt = False
 
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
         self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
index bf63805b3..a11637163 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py
@@ -5,7 +5,7 @@ import json
 import logging
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -623,30 +623,7 @@ class QwenImageEditPlusPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        image: PIL.Image.Image | list[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 = 1024,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image editing with support for multiple images."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
@@ -674,55 +651,35 @@ class QwenImageEditPlusPipeline(
         ):
             condition_images = additional_information.get("condition_images")
             vae_images = additional_information.get("vae_images")
-            condition_image_sizes = additional_information.get("condition_image_sizes")
             vae_image_sizes = additional_information.get("vae_image_sizes")
             calculated_height = additional_information.get("calculated_height")
             calculated_width = additional_information.get("calculated_width")
-            height = req.sampling_params.height
-            width = req.sampling_params.width
+            height = req.sampling_params.height or calculated_height
+            width = req.sampling_params.width or calculated_width
+            image = None
         else:
-            # fallback to run pre-processing in pipeline (debug only)
-            if image is None:
-                raise ValueError("Image is required for QwenImageEditPlusPipeline")
+            raise RuntimeError("Missing preprocess images that should have been created by the preprocess function.")
 
-            if not isinstance(image, list):
-                image = [image]
-
-            image_size = image[0].size
-            calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1])
-            height = height or calculated_height
-            width = width or calculated_width
-
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            condition_images = []
-            vae_images = []
-            condition_image_sizes = []
-            vae_image_sizes = []
-
-            for img in image:
-                image_width, image_height = img.size
-                condition_width, condition_height = calculate_dimensions(
-                    CONDITION_IMAGE_SIZE, image_width / image_height
-                )
-                vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height)
-                condition_image_sizes.append((condition_width, condition_height))
-                vae_image_sizes.append((vae_width, vae_height))
-                condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
-                vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
-
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. check inputs
         # 2. encode prompts
@@ -749,16 +706,9 @@ class QwenImageEditPlusPipeline(
         self._current_timestep = None
         self._interrupt = False
 
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
         self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)
diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
index f7a9fc4b1..e2fe67976 100644
--- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
+++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py
@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -648,31 +648,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        true_cfg_scale: float = 4.0,
-        layers: int | None = 4,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float | None = None,
-        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,
-        max_sequence_length: int = 1024,
-        resolution: int = 640,
-        cfg_normalize: bool = False,
-        use_en_prompt: bool = False,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image layered."""
 
         # 1. Get preprocessed image from request (pre-processing is done in DiffusionEngine)
@@ -688,26 +664,27 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
         negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
-        layers = req.sampling_params.layers if req.sampling_params.layers is not None else layers
-        resolution = req.sampling_params.resolution if req.sampling_params.resolution is not None else resolution
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        cfg_normalize = (
-            req.sampling_params.cfg_normalize if req.sampling_params.cfg_normalize is not None else cfg_normalize
-        )
-        use_en_prompt = (
-            req.sampling_params.use_en_prompt if req.sampling_params.use_en_prompt is not None else use_en_prompt
-        )
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        layers = req.sampling_params.layers
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        cfg_normalize = req.sampling_params.cfg_normalize
+        use_en_prompt = req.sampling_params.use_en_prompt
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
+        guidance_scale = None
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
 
         if not isinstance(first_prompt, str) and "preprocessed_image" in (
             additional_information := first_prompt.get("additional_information", {})
@@ -720,25 +697,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
             height = req.sampling_params.height
             width = req.sampling_params.width
         else:
-            # fallback to run pre-processing in pipeline (debug only)
-            if isinstance(image, PIL.Image.Image) and image.mode != "RGBA":
-                image = image.convert("RGBA")
-            image_size = image[0].size if isinstance(image, list) else image.size
-            assert resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {resolution}"
-            calculated_width, calculated_height = calculate_dimensions(
-                resolution * resolution, image_size[0] / image_size[1]
-            )
-            height = calculated_height
-            width = calculated_width
-
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
-                image = self.image_processor.resize(image, calculated_height, calculated_width)
-                prompt_image = image
-                image = self.image_processor.preprocess(image, calculated_height, calculated_width)
-                image = image.unsqueeze(2)
-                image = image.to(dtype=self.text_encoder.dtype)
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
         # 2. check inputs
         self.check_inputs(
@@ -761,16 +720,9 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         # 3. encode prompot & negative prompt
         if prompt is None or prompt == "" or prompt == " ":
             prompt = self.get_image_caption(prompt_image, use_en_prompt=use_en_prompt, device=self.device)
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         if true_cfg_scale > 1 and not has_neg_prompt:
             logger.warning(
diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
index 07015a4f4..67e2d0db7 100644
--- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
+++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
@@ -17,7 +17,6 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -555,8 +554,8 @@ class QwenImageCrossAttention(nn.Module):
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
 
-        self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
 
@@ -593,8 +592,8 @@ class QwenImageCrossAttention(nn.Module):
             prefix=f"{prefix}.to_out.0",
         )
 
-        self.norm_added_q = RMSNorm(head_dim, eps=eps)
-        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
 
         self.attn = Attention(
             num_heads=self.query_num_heads,
@@ -1007,7 +1006,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
             quant_config=quant_config,
         )
 
-        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
 
         # Entry projections (image/text) are kept full precision —
         # small sensitive layers at the network boundary (see #2728).
diff --git a/vllm_omni/diffusion/models/schedulers/__init__.py b/vllm_omni/diffusion/models/schedulers/__init__.py
index e683ed272..11aedcb1d 100644
--- a/vllm_omni/diffusion/models/schedulers/__init__.py
+++ b/vllm_omni/diffusion/models/schedulers/__init__.py
@@ -2,11 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from vllm_omni.diffusion.models.schedulers.scheduling_dmd2_euler import DMD2EulerScheduler
+from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler,
 )
 
 __all__ = [
     "DMD2EulerScheduler",
+    "FlowMatchEulerDiscreteScheduler",
     "FlowUniPCMultistepScheduler",
 ]
diff --git a/vllm_omni/diffusion/models/schedulers/scheduling_flow_match_euler_discrete.py b/vllm_omni/diffusion/models/schedulers/scheduling_flow_match_euler_discrete.py
new file mode 100644
index 000000000..8aa7d3a00
--- /dev/null
+++ b/vllm_omni/diffusion/models/schedulers/scheduling_flow_match_euler_discrete.py
@@ -0,0 +1,660 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Originally from https://github.com/huggingface/diffusers/blob/v0.39.0/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py
+# Copyright 2025 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Flow Match Euler scheduler copied from Diffusers v0.39.0.
+
+Modified by vLLM-Omni contributors for local integration and RNG compatibility.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput, is_scipy_available, logging
+from diffusers.utils.torch_utils import randn_tensor
+
+if is_scipy_available():
+    import scipy.stats
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Euler scheduler.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
+        use_dynamic_shifting (`bool`, defaults to False):
+            Whether to apply timestep shifting on-the-fly based on the image resolution.
+        base_shift (`float`, defaults to 0.5):
+            Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent
+            with desired output.
+        max_shift (`float`, defaults to 1.15):
+            Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be
+            more exaggerated or stylized.
+        base_image_seq_len (`int`, defaults to 256):
+            The base image sequence length.
+        max_image_seq_len (`int`, defaults to 4096):
+            The maximum image sequence length.
+        invert_sigmas (`bool`, defaults to False):
+            Whether to invert the sigmas.
+        shift_terminal (`float`, defaults to None):
+            The end value of the shifted timestep schedule.
+        use_karras_sigmas (`bool`, defaults to False):
+            Whether to use Karras sigmas for step sizes in the noise schedule during sampling.
+        use_exponential_sigmas (`bool`, defaults to False):
+            Whether to use exponential sigmas for step sizes in the noise schedule during sampling.
+        use_beta_sigmas (`bool`, defaults to False):
+            Whether to use beta sigmas for step sizes in the noise schedule during sampling.
+        time_shift_type (`str`, defaults to "exponential"):
+            The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
+        stochastic_sampling (`bool`, defaults to False):
+            Whether to use stochastic sampling.
+    """
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        use_dynamic_shifting: bool = False,
+        base_shift: float | None = 0.5,
+        max_shift: float | None = 1.15,
+        base_image_seq_len: int = 256,
+        max_image_seq_len: int = 4096,
+        invert_sigmas: bool = False,
+        shift_terminal: float = None,
+        use_karras_sigmas: bool = False,
+        use_exponential_sigmas: bool = False,
+        use_beta_sigmas: bool = False,
+        time_shift_type: Literal["exponential", "linear"] = "exponential",
+        stochastic_sampling: bool = False,
+    ):
+        if self.config.use_beta_sigmas and not is_scipy_available():
+            raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
+        if (
+            sum(
+                [
+                    self.config.use_beta_sigmas,
+                    self.config.use_exponential_sigmas,
+                    self.config.use_karras_sigmas,
+                ]
+            )
+            > 1
+        ):
+            raise ValueError(
+                "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, "
+                "`config.use_karras_sigmas` can be used."
+            )
+        if time_shift_type not in {"exponential", "linear"}:
+            raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.")
+
+        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
+
+        sigmas = timesteps / num_train_timesteps
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = sigmas * num_train_timesteps
+
+        self._step_index = None
+        self._begin_index = None
+
+        self._shift = shift
+
+        self.sigmas = sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def shift(self):
+        """
+        The value used for shifting.
+        """
+        return self._shift
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`, defaults to `0`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def set_shift(self, shift: float):
+        """
+        Sets the shift value for the scheduler.
+
+        Args:
+            shift (`float`):
+                The shift value to be set.
+        """
+        self._shift = shift
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        noise: torch.FloatTensor | None = None,
+    ) -> torch.FloatTensor:
+        """
+        Forward process in flow-matching
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`torch.FloatTensor`):
+                The current timestep in the diffusion chain.
+            noise (`torch.FloatTensor`):
+                The noise tensor.
+
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            timestep = timestep.to(sample.device)
+
+        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        sample = sigma * noise + (1.0 - sigma) * sample
+
+        return sample
+
+    def _sigma_to_t(self, sigma) -> float:
+        return sigma * self.config.num_train_timesteps
+
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+        """
+        Apply time shifting to the sigmas.
+
+        Args:
+            mu (`float`):
+                The mu parameter for the time shift.
+            sigma (`float`):
+                The sigma parameter for the time shift.
+            t (`torch.Tensor`):
+                The input timesteps.
+
+        Returns:
+            `torch.Tensor`:
+                The time-shifted timesteps.
+        """
+        if self.config.time_shift_type == "exponential":
+            return self._time_shift_exponential(mu, sigma, t)
+        elif self.config.time_shift_type == "linear":
+            return self._time_shift_linear(mu, sigma, t)
+
+    def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
+        r"""
+        Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config
+        value.
+
+        Reference:
+        https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51
+
+        Args:
+            t (`torch.Tensor`):
+                A tensor of timesteps to be stretched and shifted.
+
+        Returns:
+            `torch.Tensor`:
+                A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.
+        """
+        one_minus_z = 1 - t
+        scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
+        stretched_t = 1 - (one_minus_z / scale_factor)
+        return stretched_t
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device = None,
+        sigmas: list[float] | None = None,
+        mu: float | None = None,
+        timesteps: list[float] | None = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`, *optional*):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            sigmas (`list[float]`, *optional*):
+                Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
+                automatically.
+            mu (`float`, *optional*):
+                Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
+                shifting.
+            timesteps (`list[float]`, *optional*):
+                Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
+                automatically.
+        """
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`")
+
+        if sigmas is not None and timesteps is not None:
+            if len(sigmas) != len(timesteps):
+                raise ValueError("`sigmas` and `timesteps` should have the same length")
+
+        if num_inference_steps is not None:
+            if (sigmas is not None and len(sigmas) != num_inference_steps) or (
+                timesteps is not None and len(timesteps) != num_inference_steps
+            ):
+                raise ValueError(
+                    "`sigmas` and `timesteps` should have the same length as num_inference_steps, "
+                    "if `num_inference_steps` is provided"
+                )
+        else:
+            num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps)
+
+        self.num_inference_steps = num_inference_steps
+
+        # 1. Prepare default sigmas
+        is_timesteps_provided = timesteps is not None
+
+        if is_timesteps_provided:
+            timesteps = np.array(timesteps).astype(np.float32)
+
+        if sigmas is None:
+            if timesteps is None:
+                timesteps = np.linspace(
+                    self._sigma_to_t(self.sigma_max),
+                    self._sigma_to_t(self.sigma_min),
+                    num_inference_steps,
+                )
+            sigmas = timesteps / self.config.num_train_timesteps
+        else:
+            sigmas = np.array(sigmas).astype(np.float32)
+            num_inference_steps = len(sigmas)
+
+        # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
+        #    "exponential" or "linear" type is applied
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)
+        else:
+            sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)
+
+        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
+        if self.config.shift_terminal:
+            sigmas = self.stretch_shift_to_terminal(sigmas)
+
+        # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
+        if self.config.use_karras_sigmas:
+            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+        elif self.config.use_exponential_sigmas:
+            sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+        elif self.config.use_beta_sigmas:
+            sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
+        # 5. Convert sigmas and timesteps to tensors and move to specified device
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+        if not is_timesteps_provided:
+            timesteps = sigmas * self.config.num_train_timesteps
+        else:
+            timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device)
+
+        # 6. Append the terminal sigma value.
+        #    If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
+        #    `invert_sigmas` flag can be set to `True`. This case is only required in Mochi
+        if self.config.invert_sigmas:
+            sigmas = 1.0 - sigmas
+            timesteps = sigmas * self.config.num_train_timesteps
+            sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)])
+        else:
+            sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+
+        self.timesteps = timesteps
+        self.sigmas = sigmas
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(
+        self,
+        timestep: float | torch.FloatTensor,
+        schedule_timesteps: torch.FloatTensor | None = None,
+    ) -> int:
+        """
+        Get the index for the given timestep.
+
+        Args:
+            timestep (`float` or `torch.FloatTensor`):
+                The timestep to find the index for.
+            schedule_timesteps (`torch.FloatTensor`, *optional*):
+                The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used.
+
+        Returns:
+            `int`:
+                The index of the timestep.
+        """
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep: float | torch.FloatTensor) -> None:
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        return_dict: bool = True,
+    ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            per_token_timesteps (`torch.Tensor`, *optional*):
+                The timesteps for each token in the sample.
+            return_dict (`bool`, defaults to `True`):
+                Whether or not to return a
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
+
+        Returns:
+            [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`,
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned,
+                otherwise a tuple is returned where the first element is the sample tensor.
+        """
+
+        if isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+
+        if per_token_timesteps is not None:
+            per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps
+
+            sigmas = self.sigmas[:, None, None]
+            lower_mask = sigmas < per_token_sigmas[None] - 1e-6
+            lower_sigmas = lower_mask * sigmas
+            lower_sigmas, _ = lower_sigmas.max(dim=0)
+
+            current_sigma = per_token_sigmas[..., None]
+            next_sigma = lower_sigmas[..., None]
+            dt = current_sigma - next_sigma
+        else:
+            sigma_idx = self.step_index
+            sigma = self.sigmas[sigma_idx]
+            sigma_next = self.sigmas[sigma_idx + 1]
+
+            current_sigma = sigma
+            next_sigma = sigma_next
+            dt = sigma_next - sigma
+
+        if self.config.stochastic_sampling:
+            x0 = sample - current_sigma * model_output
+            noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
+            prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
+        else:
+            prev_sample = sample + dt * model_output
+
+        # upon completion increase step index by one
+        self._step_index += 1
+        if per_token_timesteps is None:
+            # Cast sample back to model compatible dtype
+            prev_sample = prev_sample.to(model_output.dtype)
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
+        """
+        Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative
+        Models](https://huggingface.co/papers/2206.00364).
+
+        Args:
+            in_sigmas (`torch.Tensor`):
+                The input sigma values to be converted.
+            num_inference_steps (`int`):
+                The number of inference steps to generate the noise schedule for.
+
+        Returns:
+            `torch.Tensor`:
+                The converted sigma values following the Karras noise schedule.
+        """
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min ** (1 / rho)
+        max_inv_rho = sigma_max ** (1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
+    def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
+        """
+        Construct an exponential noise schedule.
+
+        Args:
+            in_sigmas (`torch.Tensor`):
+                The input sigma values to be converted.
+            num_inference_steps (`int`):
+                The number of inference steps to generate the noise schedule for.
+
+        Returns:
+            `torch.Tensor`:
+                The converted sigma values following an exponential schedule.
+        """
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
+    def _convert_to_beta(
+        self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6
+    ) -> torch.Tensor:
+        """
+        Construct a beta noise schedule as proposed in [Beta Sampling is All You
+        Need](https://huggingface.co/papers/2407.12173).
+
+        Args:
+            in_sigmas (`torch.Tensor`):
+                The input sigma values to be converted.
+            num_inference_steps (`int`):
+                The number of inference steps to generate the noise schedule for.
+            alpha (`float`, *optional*, defaults to `0.6`):
+                The alpha parameter for the beta distribution.
+            beta (`float`, *optional*, defaults to `0.6`):
+                The beta parameter for the beta distribution.
+
+        Returns:
+            `torch.Tensor`:
+                The converted sigma values following a beta distribution schedule.
+        """
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.array(
+            [
+                sigma_min + (ppf * (sigma_max - sigma_min))
+                for ppf in [
+                    scipy.stats.beta.ppf(timestep, alpha, beta)
+                    for timestep in 1 - np.linspace(0, 1, num_inference_steps)
+                ]
+            ]
+        )
+        return sigmas
+
+    def _time_shift_exponential(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def _time_shift_linear(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+        return mu / (mu + (1 / t - 1) ** sigma)
+
+    def __len__(self) -> int:
+        return self.config.num_train_timesteps
diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py
index e45614a42..b2852e524 100644
--- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py
+++ b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py
@@ -628,57 +628,39 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
 
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] = "",
-        prompt_2: str | list[str] = "",
-        prompt_3: str | list[str] = "",
-        negative_prompt: str | list[str] = "",
-        negative_prompt_2: str | list[str] = "",
-        negative_prompt_3: str | list[str] = "",
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 28,
-        sigmas: list[float] | None = None,
-        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,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        pooled_prompt_embeds: torch.Tensor | None = None,
-        negative_pooled_prompt_embeds: torch.Tensor | None = None,
-        max_sequence_length: int = 256,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
-        negative_prompt = [
-            "" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts
-        ] or negative_prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+        negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+        prompt_2 = ""
+        prompt_3 = ""
+        negative_prompt_2 = ""
+        negative_prompt_3 = ""
+        prompt_embeds = None
+        negative_prompt_embeds = None
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        sigmas = common_sampling_params.sigmas or sigmas
-        max_sequence_length = common_sampling_params.max_sequence_length or max_sequence_length
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
+        sigmas = common_sampling_params.sigmas
+        max_sequence_length = common_sampling_params.max_sequence_length or 256
+        num_inference_steps = common_sampling_params.num_inference_steps or 28
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        if generator is None:
-            generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
+        output_type = common_sampling_params.output_type or "pil"
+
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
                 "prompt_embeds": prompt_embeds,
                 "negative_prompt_embeds": negative_prompt_embeds,
-                "pooled_prompt_embeds": pooled_prompt_embeds,
-                "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds,
+                "pooled_prompt_embeds": None,
+                "negative_pooled_prompt_embeds": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
@@ -717,9 +699,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         self._current_timestep = None
         self._interrupt = False
 
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
+        if prompt is not None:
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
@@ -771,7 +751,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         )
 
         self._current_timestep = None
-        if self.output_type == "latent":
+        if output_type == "latent":
             image = latents
         else:
             # Ensure the latents are the same dtype as the VAE for decode
diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
index 9c4a052ae..67e141e6a 100644
--- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
+++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py
@@ -938,8 +938,10 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         )
         logger.info("Text generation: %d chars", len(text_output))
         return DiffusionOutput(
-            output=text_output,
-            custom_output={"text_output": text_output},
+            output={
+                "payload": {"text": text_output},
+                "metadata": {"text": {"text_output": text_output}},
+            },
         )
 
     # -----------------------------------------------------------------------
@@ -1457,10 +1459,15 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
 
         images = _to_pil(image_prediction)
         img = images[0] if images else None
-        custom = {}
+        metadata = {}
         if think_text:
-            custom["think_text"] = think_text
-        return DiffusionOutput(output=img, custom_output=custom)
+            metadata["text"] = {"think_text": think_text}
+        return DiffusionOutput(
+            output={
+                "payload": {"image": img},
+                "metadata": metadata,
+            }
+        )
 
     # -----------------------------------------------------------------------
     # Weight loading
diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py
index aed72a995..cf5f84e7a 100644
--- a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py
+++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py
@@ -164,11 +164,11 @@ class SoulXPreprocessPipeline(nn.Module, SupportAudioInput, SupportsComponentDis
                 sample_rate=self.target_sr,
             )
 
-        # Non-None placeholder so DiffusionEngine forwards custom_output to stage-1
-        # (output=None triggers an early return that drops handoff).
         return DiffusionOutput(
-            output=torch.zeros(0),
-            custom_output={SOULX_PREPROCESSED_KEY: payload},
+            output={
+                "payload": {},
+                "metadata": {SOULX_PREPROCESSED_KEY: payload},
+            },
             to_cpu=True,
         )
 
diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py
index 64fa213ab..a48b5236d 100644
--- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py
+++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py
@@ -4,6 +4,7 @@ from collections.abc import Callable
 from contextlib import contextmanager
 from typing import Any, ClassVar
 
+import numpy as np
 import torch
 import torch.nn as nn
 from vllm.logger import init_logger
@@ -30,6 +31,21 @@ _SOULX_PROFILER_TARGETS = [
 ]
 
 
+def convert_soulx_audio_output_to_numpy(audio: torch.Tensor | dict[str, Any]) -> np.ndarray | dict[str, Any]:
+    if isinstance(audio, dict) and isinstance(audio.get("payload"), dict):
+        payload = dict(audio["payload"])
+        audio_payload = payload.get("audio")
+        if isinstance(audio_payload, torch.Tensor):
+            payload["audio"] = audio_payload.detach().cpu().float().numpy()
+        return {
+            "payload": payload,
+            "metadata": audio.get("metadata") or {},
+        }
+    if isinstance(audio, torch.Tensor):
+        return audio.detach().cpu().float().numpy()
+    return audio
+
+
 class FlowMatchingAudioPipeline(
     nn.Module,
     CFGParallelMixin,
@@ -326,7 +342,7 @@ class FlowMatchingAudioPipeline(
         req: OmniDiffusionRequest,
         *,
         kind: str,
-        custom_output_key: str,
+        metadata_key: str,
         infer_batch_fn: Callable[..., tuple[torch.Tensor, int]],
         prepare_extra_args: Callable[[dict[str, Any], Any], dict[str, Any]] | None = None,
     ) -> DiffusionOutput:
@@ -353,7 +369,9 @@ class FlowMatchingAudioPipeline(
             )
 
         return DiffusionOutput(
-            output=audio,
-            custom_output={custom_output_key: pitch_shift},
+            output={
+                "payload": {"audio": audio},
+                "metadata": {"audio": {metadata_key: pitch_shift}},
+            },
             stage_durations=self._profiler_stage_durations() or {},
         )
diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py
index 829256134..85fed9ee7 100644
--- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py
+++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py
@@ -19,6 +19,7 @@ from vllm_omni.diffusion.models.soulx_singer.modules import (
 )
 from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import (
     FlowMatchingAudioPipeline,
+    convert_soulx_audio_output_to_numpy,
 )
 from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
     SOULX_PREPROCESSED_KEY,
@@ -96,7 +97,7 @@ def get_soulxsinger_svc_pre_process_func(od_config: OmniDiffusionConfig):
 
 def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig):
     def post_process_func(audio: torch.Tensor):
-        return audio.detach().cpu().float().numpy()
+        return convert_soulx_audio_output_to_numpy(audio)
 
     return post_process_func
 
@@ -400,7 +401,7 @@ class PipelineSoulXSingerSVC(FlowMatchingAudioPipeline):
         return self._forward_batch_from_request(
             req,
             kind="svc",
-            custom_output_key="pitch_shift",
+            metadata_key="pitch_shift",
             infer_batch_fn=self.infer_svc_batch,
         )
 
diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py
index 9be70e651..fb2c086f2 100644
--- a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py
+++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py
@@ -15,6 +15,7 @@ from vllm_omni.diffusion.models.soulx_singer.modules import (
 )
 from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import (
     FlowMatchingAudioPipeline,
+    convert_soulx_audio_output_to_numpy,
 )
 from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
     SOULX_PREPROCESSED_KEY,
@@ -97,7 +98,7 @@ def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig):
     """Convert pipeline audio tensor output for offline consumers."""
 
     def post_process_func(audio: torch.Tensor):
-        return audio.detach().cpu().float().numpy()
+        return convert_soulx_audio_output_to_numpy(audio)
 
     return post_process_func
 
@@ -343,7 +344,7 @@ class PipelineSoulXSingerSVS(FlowMatchingAudioPipeline):
         return self._forward_batch_from_request(
             req,
             kind="svs",
-            custom_output_key="f0_shift",
+            metadata_key="f0_shift",
             infer_batch_fn=self.infer_svs_batch,
             prepare_extra_args=lambda extra_args, _sampling: normalize_svs_control_extra_args(extra_args),
         )
diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py
index 83f81e481..ff4e4b7ca 100644
--- a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py
+++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py
@@ -13,7 +13,7 @@ SOULX_PRECOMPUTED_KEYS_BY_KIND: dict[str, tuple[str, ...]] = {
 
 
 # Key under which the SoulX preprocess payload is stored in
-# prompt["additional_information"] or DiffusionOutput.custom_output.
+# prompt["additional_information"] or diffusion output metadata.
 SOULX_PREPROCESSED_KEY = "soulx_preprocessed"
 
 
diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py
index 23b61209a..41f12f654 100644
--- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py
+++ b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py
@@ -385,37 +385,21 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        audio_end_in_s: float | None = None,
-        audio_start_in_s: float = 0.0,
-        guidance_scale: float = 7.0,
-        num_waveforms_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str = "np",
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """
         Generate audio from text prompt.
 
         Args:
-            req: OmniDiffusionRequest containing generation parameters
-            prompt: Text prompt for audio generation
-            negative_prompt: Negative prompt for CFG
-            audio_end_in_s: Audio end time in seconds (max ~47s for stable-audio-open-1.0)
-            audio_start_in_s: Audio start time in seconds
-            guidance_scale: CFG scale
-            num_waveforms_per_prompt: Number of audio outputs per prompt
-            generator: Random generator for reproducibility
-            latents: Pre-generated latents
-            prompt_embeds: Pre-computed prompt embeddings
-            negative_prompt_embeds: Pre-computed negative prompt embeddings
-            output_type: Output format ("np", "pt", or "latent")
+            req: OmniDiffusionRequest containing generation parameters.
+                The `req.sampling_params.extra_args` can include the following keys:
+                - audio_start_in_s (`float`, *optional*, defaults to 0.0):
+                    Start time of the audio in seconds.
+                - audio_end_in_s (`float`, *optional*):
+                    End time of the audio in seconds.
+                - num_waveforms_per_prompt (`int`, *optional*, defaults to 1):
+                    Number of audio outputs per prompt.
+                - output_type (`str`, *optional*, defaults to "np"):
+                    Output format ("np", "pt", or "latent").
 
         Returns:
             DiffusionOutput containing generated audio
@@ -423,28 +407,33 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         # Extract from request
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
-        num_inference_steps = req.sampling_params.num_inference_steps
-        if num_inference_steps is None:
-            num_inference_steps = 100  # Default steps
-        num_inference_steps = req.sampling_params.num_inference_steps
-        if num_inference_steps is None:
-            num_inference_steps = 50  # Default steps
+
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 7.0
 
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
+        latents = req.sampling_params.latents
 
         # Get audio duration from request extra params or defaults
-        audio_start_in_s = req.sampling_params.extra_args.get("audio_start_in_s", audio_start_in_s)
-        audio_end_in_s = req.sampling_params.extra_args.get("audio_end_in_s", audio_end_in_s)
+        audio_start_in_s: float = req.sampling_params.extra_args.get("audio_start_in_s", 0.0)
+        audio_end_in_s: float | None = req.sampling_params.extra_args.get("audio_end_in_s", None)
+        num_waveforms_per_prompt: int = req.sampling_params.extra_args.get(
+            "num_waveforms_per_prompt", req.sampling_params.num_outputs_per_prompt or 1
+        )
+        output_type = req.sampling_params.output_type or "np"
 
         # Calculate audio length
         downsample_ratio = self.vae.hop_length
@@ -474,12 +463,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         )
 
         # Determine batch size
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = len(prompt)
 
         device = self.device
         do_classifier_free_guidance = guidance_scale > 1.0
@@ -500,7 +484,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
             audio_start_in_s,
             audio_end_in_s,
             device,
-            do_classifier_free_guidance and (negative_prompt is not None or negative_prompt_embeds is not None),
+            do_classifier_free_guidance and negative_prompt is not None,
             batch_size,
         )
 
@@ -515,7 +499,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         )
 
         # Handle CFG without negative prompt
-        if do_classifier_free_guidance and negative_prompt_embeds is None and negative_prompt is None:
+        if do_classifier_free_guidance and negative_prompt is None:
             negative_text_audio_duration_embeds = torch.zeros_like(text_audio_duration_embeds)
             text_audio_duration_embeds = torch.cat(
                 [negative_text_audio_duration_embeds, text_audio_duration_embeds],
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
index b2870549c..d3726e087 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
@@ -187,7 +187,7 @@ def get_wan22_post_process_func(
     ):
         if output_type == "latent":
             return video
-        custom_output = {}
+        video_metadata = {}
         if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
             video, multiplier = interpolate_video_tensor(
                 video,
@@ -195,10 +195,10 @@ def get_wan22_post_process_func(
                 scale=sampling_params.frame_interpolation_scale,
                 model_path=sampling_params.frame_interpolation_model_path,
             )
-            custom_output["video_fps_multiplier"] = multiplier
+            video_metadata["video_fps_multiplier"] = multiplier
         return {
-            "video": video_processor.postprocess_video(video, output_type=output_type),
-            "custom_output": custom_output,
+            "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)},
+            "metadata": {"video": video_metadata} if video_metadata else {},
         }
 
     return post_process_func
@@ -540,38 +540,27 @@ class Wan22Pipeline(
 
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | None = None,
-        negative_prompt: str | None = None,
-        height: int = 480,
-        width: int = 832,
-        num_inference_steps: int = 40,
-        guidance_scale: float | tuple[float, float] = 4.0,
-        frame_num: int = 81,
-        output_type: str | None = "np",
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        attention_kwargs: dict | None = None,
-        **kwargs,
-    ) -> DiffusionOutput:
-        # Get parameters from request or arguments
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        prompt: str | None = None
+        negative_prompt: str | None = None
+        prompt_embeds: torch.Tensor | None = None
+        negative_prompt_embeds: torch.Tensor | None = None
         if len(req.prompts) > 1:
             raise ValueError(
                 """This model only supports a single prompt, not a batched request.""",
                 """Please pass in a single prompt object or string, or a single-item list.""",
             )
-        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
-            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
-            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
+        if len(req.prompts) == 1:
+            first_prompt = req.prompts[0]
+            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num
+        if not prompt:
+            raise ValueError("Prompt is required for Wan2.2 generation.")
+
+        height = req.sampling_params.height or 480
+        width = req.sampling_params.width or 832
+        num_frames = req.sampling_params.num_frames or 81
 
         # Ensure dimensions are compatible with VAE and patch size
         # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size)
@@ -579,11 +568,16 @@ class Wan22Pipeline(
         mod_value = self.vae_scale_factor_spatial * patch_size[1]  # 16*2=32 for TI2V, 8*2=16 for I2V
         height = (height // mod_value) * mod_value
         width = (width // mod_value) * mod_value
-        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 4.0
+
+        output_type = req.sampling_params.output_type or "np"
+        attention_kwargs: dict | None = None
 
         guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
         guidance_high = (
@@ -634,8 +628,7 @@ class Wan22Pipeline(
             dtype = self.text_encoder.dtype
 
         # Seed / generator
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
@@ -644,24 +637,16 @@ class Wan22Pipeline(
             current_omni_platform.synchronize()
             _t_pipeline_start = time.perf_counter()
             _t_text_enc_start = _t_pipeline_start
-        if prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
-                device=device,
-                dtype=dtype,
-            )
-        else:
-            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
-            if negative_prompt_embeds is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
-            elif guidance_low > 1.0 or guidance_high > 1.0:
-                raise ValueError(
-                    "negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1."
-                )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+            max_sequence_length=req.sampling_params.max_sequence_length or 512,
+            device=device,
+            dtype=dtype,
+        )
+
         if DEBUG_PERF:
             current_omni_platform.synchronize()
             _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000
diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
index fc2d79cec..8b042b595 100644
--- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
+++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
@@ -65,7 +65,7 @@ def get_wan22_i2v_post_process_func(
     ):
         if output_type == "latent":
             return video
-        custom_output = {}
+        video_metadata = {}
         if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
             video, multiplier = interpolate_video_tensor(
                 video,
@@ -73,10 +73,10 @@ def get_wan22_i2v_post_process_func(
                 scale=sampling_params.frame_interpolation_scale,
                 model_path=sampling_params.frame_interpolation_model_path,
             )
-            custom_output["video_fps_multiplier"] = multiplier
+            video_metadata["video_fps_multiplier"] = multiplier
         return {
-            "video": video_processor.postprocess_video(video, output_type=output_type),
-            "custom_output": custom_output,
+            "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)},
+            "metadata": {"video": video_metadata} if video_metadata else {},
         }
 
     return post_process_func
@@ -429,66 +429,59 @@ class Wan22I2VPipeline(
         quant_config = getattr(self.od_config, "quantization_config", None)
         return create_transformer_from_config(config, quant_config=quant_config)
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | None = None,
-        negative_prompt: str | None = None,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        height: int = 480,
-        width: int = 832,
-        num_inference_steps: int = 40,
-        guidance_scale: float | tuple[float, float] = 5.0,
-        frame_num: int = 81,
-        output_type: str | None = "np",
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        image_embeds: torch.Tensor | None = None,
-        last_image: PIL.Image.Image | torch.Tensor | None = None,
-        attention_kwargs: dict | None = None,
-        **kwargs,
-    ) -> DiffusionOutput:
-        # Get parameters from request or arguments
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        prompt: str | None = None
+        negative_prompt: str | None = None
+        prompt_embeds: torch.Tensor | None = None
+        negative_prompt_embeds: torch.Tensor | None = None
+        image_embeds: torch.Tensor | None = None
         if len(req.prompts) > 1:
             raise ValueError(
                 """This model only supports a single prompt, not a batched request.""",
                 """Please pass in a single prompt object or string, or a single-item list.""",
             )
-        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
-            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
-            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
+        if len(req.prompts) == 1:
+            first_prompt = req.prompts[0]
+            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+
+        if not prompt:
+            raise ValueError("Prompt is required for Wan2.2 generation.")
 
         # Get image from request
-        if image is None:
-            multi_modal_data = (
-                req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
-            )
-            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
-            if raw_image is None:
-                raise ValueError("Image is required for I2V generation.")
-            if isinstance(raw_image, list):
-                if len(raw_image) > 1:
-                    logger.warning(
-                        """Received a list of image. Only a single image is supported by this model."""
-                        """Taking only the first image for now."""
-                    )
-                raw_image = raw_image[0]
-            if isinstance(raw_image, str):
-                image = PIL.Image.open(raw_image)
-            else:
-                image = cast(PIL.Image.Image | torch.Tensor, raw_image)
+        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
+        raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+        if raw_image is None:
+            raise ValueError("Image is required for I2V generation.")
+        if isinstance(raw_image, list):
+            if len(raw_image) > 1:
+                logger.warning(
+                    """Received a list of image. Only a single image is supported by this model."""
+                    """Taking only the first image for now."""
+                )
+            raw_image = raw_image[0]
+        if isinstance(raw_image, str):
+            image = PIL.Image.open(raw_image)
+        else:
+            image = cast(PIL.Image.Image | torch.Tensor, raw_image)
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_frames = req.sampling_params.num_frames or frame_num
-        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        last_image: PIL.Image.Image | torch.Tensor | None = None
+        if multi_modal_data is not None:
+            last_image = multi_modal_data.get("last_image", None)
+
+        height = req.sampling_params.height or 480
+        width = req.sampling_params.width or 832
+        num_frames = req.sampling_params.num_frames or 81
+        num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 5.0
+
+        output_type = req.sampling_params.output_type or "np"
+        attention_kwargs: dict | None = None
 
         # Handle guidance scales
         guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
@@ -533,8 +526,7 @@ class Wan22I2VPipeline(
         dtype = self.transformer.dtype
 
         # Generator setup
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
@@ -544,20 +536,15 @@ class Wan22I2VPipeline(
             _t_pipeline_start = time.perf_counter()
             _t_text_enc_start = _t_pipeline_start
 
-        if prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
-                device=device,
-                dtype=dtype,
-            )
-        else:
-            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
-            if negative_prompt_embeds is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+            max_sequence_length=req.sampling_params.max_sequence_length or 512,
+            device=device,
+            dtype=dtype,
+        )
 
         if DEBUG_PERF:
             current_omni_platform.synchronize()
diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
index 104576f41..b77d2265c 100644
--- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
+++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py
@@ -24,7 +24,7 @@ from typing import Any, ClassVar
 import PIL.Image
 import torch
 import torch.nn as nn
-from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.image_processor import VaeImageProcessor
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import logging
 from diffusers.utils.torch_utils import randn_tensor
@@ -410,124 +410,21 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        image: PipelineImageInput = None,
-        strength: float = 0.6,
-        height: int = 1024,
-        width: int = 1024,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 5.0,
-        cfg_normalization: bool = False,
-        cfg_truncation: float = 1.0,
-        negative_prompt: str | list[str] | None = None,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: list[torch.FloatTensor] | None = None,
-        negative_prompt_embeds: list[torch.FloatTensor] | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            image (`PipelineImageInput`, *optional*):
-                The image to use for img2img generation. If provided, the pipeline
-                will perform img2img instead of text-to-image.
-            strength (`float`, *optional*, defaults to 0.6):
-                Indicates extent to transform the reference `image`. Must be between 0 and 1.
-            height (`int`, *optional*, defaults to 1024):
-                The height in pixels of the generated image.
-            width (`int`, *optional*, defaults to 1024):
-                The width in pixels of the generated image.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`list[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            guidance_scale (`float`, *optional*, defaults to 5.0):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                0`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            cfg_normalization (`bool`, *optional*, defaults to False):
-                Whether to apply configuration normalization.
-            cfg_truncation (`float`, *optional*, defaults to 1.0):
-                The truncation value for configuration.
-            negative_prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than or equal to `0`).
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.FloatTensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`list[torch.FloatTensor]`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain
-                tuple.
-            joint_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`list`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int`, *optional*, defaults to 512):
-                Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if
-            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
-            generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
-        # Handle img2img: extract image from request
-        if image is None and req.prompts:
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        image = None
+        if req.prompts:
             if len(req.prompts) > 1:
                 logger.warning(
                     "This model only supports a single prompt for img2img, not a batched request. "
@@ -542,10 +439,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
                     else:
                         image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else raw_image
 
-        # strength is currently only applicable for Z-Image I2I; other pipelines ignore this parameter
         explicit_strength = req.sampling_params.strength is not None
-        if explicit_strength:
-            strength = req.sampling_params.strength
+        strength = req.sampling_params.strength if explicit_strength else 0.6
         if explicit_strength and image is None:
             logger.warning(
                 "strength parameter (%.2f) is only applicable for image-to-image (I2I) generation. "
@@ -556,20 +451,24 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
         if image is not None and strength is not None and (strength < 0 or strength > 1):
             raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}")
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        height = req.sampling_params.height or 1024
+        width = req.sampling_params.width or 1024
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         generator = req.sampling_params.generator
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_rescale is not None else guidance_scale
-        )
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 512
+        guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+
+        cfg_normalization = req.sampling_params.cfg_normalize
+        cfg_truncation = req.sampling_params.extra_args.get("cfg_truncation", 1.0)
+        joint_attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
+        output_type = req.sampling_params.output_type or "pil"
 
         vae_scale = self.vae_scale_factor * 2
         if height % vae_scale != 0:
@@ -591,33 +490,20 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
         self._cfg_normalization = cfg_normalization
         self._cfg_truncation = cfg_truncation
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = len(prompt_embeds)
+        batch_size = len(prompt)
 
-        # If prompt_embeds is provided and prompt is None, skip encoding
-        if prompt_embeds is not None and prompt is None:
-            if self.do_classifier_free_guidance and negative_prompt_embeds is None:
-                raise ValueError(
-                    "When `prompt_embeds` is provided without `prompt`, "
-                    "`negative_prompt_embeds` must also be provided for classifier-free guidance."
-                )
-        else:
-            (
-                prompt_embeds,
-                negative_prompt_embeds,
-            ) = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=self.do_classifier_free_guidance,
-                prompt_embeds=prompt_embeds,
-                negative_prompt_embeds=negative_prompt_embeds,
-                device=device,
-                max_sequence_length=max_sequence_length,
-            )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            device=device,
+            max_sequence_length=max_sequence_length,
+        )
 
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.in_channels
diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py
index 320f9bf38..17d4ca2fd 100644
--- a/vllm_omni/diffusion/offloader/module_collector.py
+++ b/vllm_omni/diffusion/offloader/module_collector.py
@@ -22,6 +22,38 @@ class PipelineModules:
     resident_modules: list[nn.Module] = field(default_factory=list)
     resident_names: list[str] = field(default_factory=list)
 
+    def outermost_dits(self) -> tuple[list[str], list[nn.Module]]:
+        """Discovered DiTs that are not nested inside another discovered DiT.
+
+        Some pipelines declare a DiT and one of its own submodules as separate
+        DiTs (e.g. Cosmos3 lists both ``transformer`` and the nested
+        ``transformer.language_model``) so that offloading can treat them as
+        independent rings. HSDP instead shards each DiT through that DiT's own
+        ``_hsdp_shard_conditions``; an inner DiT's blocks are already covered by
+        its ancestor's conditions, so sharding it again would double-wrap them
+        (and the inner stack may not declare conditions at all). Keeping only the
+        outermost DiTs leaves the ancestor's conditions as the single source of
+        truth for every nested block. Order follows ``dit_names``.
+        """
+
+        def is_nested_in_another_dit(dit: nn.Module) -> bool:
+            for other in self.dits:
+                if other is dit:
+                    continue
+                for submodule in other.modules():
+                    if submodule is dit:
+                        return True
+            return False
+
+        outer_names: list[str] = []
+        outer_modules: list[nn.Module] = []
+        for name, dit in zip(self.dit_names, self.dits):
+            if is_nested_in_another_dit(dit):
+                continue
+            outer_names.append(name)
+            outer_modules.append(dit)
+        return outer_names, outer_modules
+
 
 class ModuleDiscovery:
     """Discovers pipeline components.
diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py
index 0edf77651..0b031be3a 100644
--- a/vllm_omni/diffusion/offloader/sequential_backend.py
+++ b/vllm_omni/diffusion/offloader/sequential_backend.py
@@ -197,12 +197,35 @@ class ModelLevelOffloadBackend(OffloadBackend):
     def __init__(self, config: OffloadConfig, device: torch.device):
         super().__init__(config, device)
         self._offload_modules: list[nn.Module] = []  # Track modules with hooks
+        self._custom_pipeline: nn.Module | None = None
 
     def enable(self, pipeline: nn.Module) -> None:
         if self.enabled:
             logger.warning("ModelLevelOffloadBackend already enabled")
             return
 
+        # Pipelines with a nested transformer (e.g. Cosmos3's reasoner/generator
+        # pathways) own their own mutual-exclusion swaps and cannot be offloaded
+        # with the generic encoder<->DiT hooks. Delegate to the custom hook.
+        # TODO: Cosmos3 is the only implementer today, so we duck-type the
+        # enable/disable_omni_model_cpu_offload pair via getattr. Once a second
+        # pipeline needs it, formalize this as a Protocol (like
+        # SupportsComponentDiscovery) instead of getattr detection.
+        custom_enable = getattr(pipeline, "enable_omni_model_cpu_offload", None)
+        if callable(custom_enable):
+            custom_enable(
+                device=self.device,
+                pin_memory=self.config.pin_cpu_memory,
+                use_hsdp=self.config.use_hsdp,
+            )
+            self._custom_pipeline = pipeline
+            self.enabled = True
+            logger.info(
+                "Model-level offloading enabled through %s.enable_omni_model_cpu_offload",
+                pipeline.__class__.__name__,
+            )
+            return
+
         modules = ModuleDiscovery.discover(pipeline)
 
         # Move encoders to GPU
@@ -259,6 +282,15 @@ class ModelLevelOffloadBackend(OffloadBackend):
         if not self.enabled:
             return
 
+        if self._custom_pipeline is not None:
+            custom_disable = getattr(self._custom_pipeline, "disable_omni_model_cpu_offload", None)
+            if callable(custom_disable):
+                custom_disable()
+            self._custom_pipeline = None
+            self.enabled = False
+            logger.info("Model-level offloading disabled")
+            return
+
         remove_sequential_offload(self._offload_modules)
 
         self._offload_modules.clear()
diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py
index 0611c6b04..a36305fce 100644
--- a/vllm_omni/diffusion/output_formatter.py
+++ b/vllm_omni/diffusion/output_formatter.py
@@ -3,8 +3,9 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import Any
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import TypeGuard
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.io_support import supports_audio_output
@@ -12,6 +13,18 @@ from vllm_omni.diffusion.registry import DiffusionModelRegistry
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniPromptType
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.outputs.output_metadata import (
+    DiffusionMetadata,
+    DiffusionMultimodalOutput,
+    DiffusionOutputEnvelope,
+    DiffusionPayload,
+    DiffusionPayloadValue,
+    DiffusionPostprocessRawOutput,
+    DiffusionTrajectoryPayload,
+    strip_internal_metadata,
+    validate_diffusion_metadata,
+    validate_public_diffusion_metadata,
+)
 
 
 @dataclass(frozen=True)
@@ -24,53 +37,79 @@ class DiffusionStepTimings:
 
 @dataclass(frozen=True)
 class DiffusionPostprocessOutput:
-    outputs: Any
-    custom_output: dict[str, Any]
-    audio_payload: Any | None = None
-    action_payload: Any | None = None
-    audio_sample_rate: int | None = None
-    fps: float | None = None
-    has_video_payload: bool = False
+    outputs: DiffusionPayload
+    metadata: DiffusionMetadata = field(default_factory=dict)
+    primary_key: str | None = None
+
+
+def _is_output_envelope(outputs: DiffusionPostprocessRawOutput) -> TypeGuard[DiffusionOutputEnvelope]:
+    return isinstance(outputs, dict) and isinstance(outputs.get("payload"), dict)
 
 
 def normalize_diffusion_postprocess_output(
-    outputs: Any,
-    custom_output: dict[str, Any],
+    outputs: DiffusionPostprocessRawOutput,
 ) -> DiffusionPostprocessOutput:
-    """Normalize the legacy postprocess dict shape used by diffusion models.
+    """Normalize diffusion postprocess output into payload and metadata."""
 
-    The returned envelope owns a shallow merged copy of ``custom_output`` so
-    postprocess values keep legacy override precedence without mutating the
-    caller-owned dict.
-    """
+    if _is_output_envelope(outputs):
+        payload: DiffusionPayload = outputs.get("payload") or {}
+        metadata = outputs.get("metadata") or {}
+        if not isinstance(metadata, dict):
+            metadata = {}
+        validate_diffusion_metadata(metadata)
 
-    audio_payload = None
-    action_payload = None
-    audio_sample_rate = None
-    fps = None
-    has_video_payload = False
-    merged_custom_output = dict(custom_output)
+        public_metadata = strip_internal_metadata(metadata)
+        if "text" in payload and "text" not in public_metadata:
+            public_metadata = {**public_metadata, "text": {}}
+        validate_public_diffusion_metadata(public_metadata)
+
+        return DiffusionPostprocessOutput(
+            outputs=payload,
+            metadata=public_metadata,
+            primary_key=_infer_primary_payload_key(payload),
+        )
 
     if isinstance(outputs, dict):
-        has_video_payload = "video" in outputs
-        audio_payload = outputs.get("audio")
-        action_payload = outputs.get("actions")
-        merged_custom_output.update(outputs.get("custom_output") or {})
-        audio_sample_rate = outputs.get("audio_sample_rate")
-        fps = outputs.get("fps")
-        outputs = outputs.get("video", outputs)
+        payload = {key: value for key, value in outputs.items() if key not in {"audio_sample_rate", "fps"}}
+        metadata = _metadata_from_legacy_payload(outputs)
+        if "text" in payload and "text" not in metadata:
+            metadata["text"] = {}
+        if metadata:
+            validate_public_diffusion_metadata(metadata)
+        return DiffusionPostprocessOutput(
+            outputs=payload,
+            metadata=metadata,
+            primary_key=_infer_primary_payload_key(payload),
+        )
 
     return DiffusionPostprocessOutput(
-        outputs=outputs,
-        custom_output=merged_custom_output,
-        audio_payload=audio_payload,
-        action_payload=action_payload,
-        audio_sample_rate=audio_sample_rate,
-        fps=fps,
-        has_video_payload=has_video_payload,
+        outputs={"output": outputs},
+        primary_key="output",
     )
 
 
+def _metadata_from_legacy_payload(payload: dict[str, DiffusionPayloadValue]) -> DiffusionMetadata:
+    metadata: DiffusionMetadata = {}
+    audio_sample_rate = payload.get("audio_sample_rate")
+    if audio_sample_rate is not None:
+        metadata["audio"] = {"sample_rate": audio_sample_rate}
+    fps = payload.get("fps")
+    if fps is not None:
+        metadata["video"] = {"fps": fps}
+    return metadata
+
+
+def _infer_primary_payload_key(payload: DiffusionPayload) -> str | None:
+    for key in ("video", "image", "text", "audio", "output"):
+        if key in payload:
+            return key
+    if set(payload).issubset({"actions", "trajectory"}):
+        return None
+    if "actions" in payload:
+        return None
+    return next(iter(payload), None)
+
+
 def format_empty_diffusion_outputs(
     request: OmniDiffusionRequest,
     *,
@@ -93,13 +132,14 @@ def format_diffusion_outputs(
     request: OmniDiffusionRequest,
     od_config: OmniDiffusionConfig,
     diffusion_output: DiffusionOutput,
-    output_data: Any,
+    output_data: DiffusionPostprocessRawOutput,
     postprocess_output: DiffusionPostprocessOutput,
     timings: DiffusionStepTimings,
 ) -> list[OmniRequestOutput]:
     """Convert a finished diffusion model output into API-facing outputs."""
 
-    outputs = _ensure_list(postprocess_output.outputs)
+    primary_payload = _primary_payload(postprocess_output)
+    outputs = _ensure_list(primary_payload)
     metrics = {
         "preprocess_time_ms": timings.preprocess_time_s * 1000,
         "diffusion_engine_exec_time_ms": timings.exec_time_s * 1000,
@@ -112,10 +152,10 @@ def format_diffusion_outputs(
     # Detect text output: when the pipeline returns a string (e.g.,
     # SenseNova-U1 / BAGEL single-stage img2text / text2text), wrap it
     # as a text-type response instead of an image.
-    is_text_output = isinstance(output_data, str) and postprocess_output.custom_output.get("text_output") is not None
+    is_text_output = postprocess_output.primary_key == "text" or "text" in postprocess_output.metadata
 
     is_audio_output = supports_audio_output(od_config.model_class_name)
-    audio_sample_rate = postprocess_output.audio_sample_rate
+    audio_sample_rate = _metadata_audio_sample_rate(postprocess_output.metadata)
     if is_audio_output and audio_sample_rate is None:
         model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
         audio_sample_rate = getattr(model_cls, "audio_sample_rate", None)
@@ -134,14 +174,42 @@ def format_diffusion_outputs(
     )
 
 
-def _ensure_list(outputs: Any) -> list[Any]:
+def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]:
     if isinstance(outputs, list):
         return outputs
     return [outputs] if outputs is not None else []
 
 
-def _format_audio_multimodal_output(payload: Any, audio_sample_rate: int | None) -> dict[str, Any]:
-    mm_output: dict[str, Any] = {"audio": payload}
+def _primary_payload(postprocess_output: DiffusionPostprocessOutput) -> DiffusionPayloadValue | None:
+    if postprocess_output.primary_key is None:
+        return []
+    return postprocess_output.outputs.get(postprocess_output.primary_key)
+
+
+def _metadata_audio_sample_rate(metadata: DiffusionMetadata) -> int | None:
+    audio_metadata = metadata.get("audio")
+    if not isinstance(audio_metadata, dict):
+        return None
+    sample_rate = audio_metadata.get("sample_rate")
+    return sample_rate if isinstance(sample_rate, int) else None
+
+
+def _metadata_video_fps(metadata: DiffusionMetadata) -> float | None:
+    video_metadata = metadata.get("video")
+    if not isinstance(video_metadata, dict):
+        return None
+    fps = video_metadata.get("fps")
+    return fps if isinstance(fps, (int, float)) and not isinstance(fps, bool) else None
+
+
+def _format_audio_multimodal_output(
+    payload: DiffusionPayloadValue,
+    audio_sample_rate: int | None,
+    metadata: DiffusionMetadata | None = None,
+) -> DiffusionMultimodalOutput:
+    mm_output: DiffusionMultimodalOutput = {"audio": payload}
+    if metadata:
+        mm_output["metadata"] = metadata
     if audio_sample_rate is not None:
         mm_output["audio_sample_rate"] = audio_sample_rate
     return mm_output
@@ -149,25 +217,29 @@ def _format_audio_multimodal_output(payload: Any, audio_sample_rate: int | None)
 
 def _has_non_audio_postprocess_payload(postprocess_output: DiffusionPostprocessOutput) -> bool:
     return (
-        postprocess_output.has_video_payload
-        or postprocess_output.action_payload is not None
-        or postprocess_output.fps is not None
+        "video" in postprocess_output.outputs
+        or "image" in postprocess_output.outputs
+        or "text" in postprocess_output.metadata
+        or "actions" in postprocess_output.outputs
+        or "trajectory" in postprocess_output.outputs
+        or _metadata_video_fps(postprocess_output.metadata) is not None
     )
 
 
 def _build_multimodal_output(
     postprocess_output: DiffusionPostprocessOutput,
     audio_sample_rate: int | None,
-) -> dict[str, Any]:
-    mm_output: dict[str, Any] = {}
-    if postprocess_output.audio_payload is not None:
-        mm_output["audio"] = postprocess_output.audio_payload
+) -> DiffusionMultimodalOutput:
+    mm_output: DiffusionMultimodalOutput = {}
+    if postprocess_output.metadata:
+        mm_output["metadata"] = postprocess_output.metadata
+    for key, value in postprocess_output.outputs.items():
+        if key in {"audio", "actions", "trajectory"}:
+            mm_output[key] = value
     if audio_sample_rate is not None:
         mm_output["audio_sample_rate"] = audio_sample_rate
-    if postprocess_output.fps is not None:
-        mm_output["fps"] = postprocess_output.fps
-    if postprocess_output.action_payload is not None:
-        mm_output["actions"] = postprocess_output.action_payload
+    if (fps := _metadata_video_fps(postprocess_output.metadata)) is not None:
+        mm_output["fps"] = fps
     return mm_output
 
 
@@ -176,8 +248,8 @@ def _format_single_prompt_output(
     request: OmniDiffusionRequest,
     prompt: OmniPromptType,
     diffusion_output: DiffusionOutput,
-    outputs: list[Any],
-    metrics: dict[str, Any],
+    outputs: list[DiffusionPayloadValue],
+    metrics: dict[str, object],
     postprocess_output: DiffusionPostprocessOutput,
     is_text_output: bool,
     is_audio_output: bool,
@@ -186,6 +258,13 @@ def _format_single_prompt_output(
 ) -> list[OmniRequestOutput]:
     request_id = request.request_id
     mm_output = _build_multimodal_output(postprocess_output, audio_sample_rate)
+    if is_text_output:
+        mm_output["text"] = outputs[0] if len(outputs) == 1 else outputs
+    trajectory_payload = _trajectory_payload(postprocess_output, diffusion_output)
+    trajectory_latents = trajectory_payload.get("latents")
+    trajectory_timesteps = trajectory_payload.get("timesteps")
+    trajectory_log_probs = trajectory_payload.get("log_probs")
+    trajectory_decoded = trajectory_payload.get("decoded")
 
     if is_text_output:
         return [
@@ -194,7 +273,6 @@ def _format_single_prompt_output(
                 images=[],
                 prompt=prompt,
                 metrics=metrics,
-                custom_output=postprocess_output.custom_output,
                 multimodal_output=mm_output,
                 final_output_type="text",
                 stage_durations=diffusion_output.stage_durations,
@@ -204,7 +282,7 @@ def _format_single_prompt_output(
         ]
 
     if is_audio_output and not _has_non_audio_postprocess_payload(postprocess_output):
-        request_audio_payload = postprocess_output.audio_payload
+        request_audio_payload = postprocess_output.outputs.get("audio")
         if request_audio_payload is None:
             request_audio_payload = outputs[0] if len(outputs) == 1 else outputs
         return [
@@ -213,14 +291,15 @@ def _format_single_prompt_output(
                 images=[],
                 prompt=prompt,
                 metrics=metrics,
-                latents=diffusion_output.trajectory_latents,
-                trajectory_latents=diffusion_output.trajectory_latents,
-                trajectory_timesteps=diffusion_output.trajectory_timesteps,
-                trajectory_log_probs=diffusion_output.trajectory_log_probs,
-                trajectory_decoded=diffusion_output.trajectory_decoded,
+                latents=trajectory_latents,
+                trajectory_latents=trajectory_latents,
+                trajectory_timesteps=trajectory_timesteps,
+                trajectory_log_probs=trajectory_log_probs,
+                trajectory_decoded=trajectory_decoded,
                 multimodal_output=_format_audio_multimodal_output(
                     request_audio_payload,
                     audio_sample_rate,
+                    postprocess_output.metadata,
                 ),
                 final_output_type="audio",
                 stage_durations=diffusion_output.stage_durations,
@@ -235,15 +314,43 @@ def _format_single_prompt_output(
             images=outputs,
             prompt=prompt,
             metrics=metrics,
-            latents=diffusion_output.trajectory_latents,
-            trajectory_latents=diffusion_output.trajectory_latents,
-            trajectory_timesteps=diffusion_output.trajectory_timesteps,
-            trajectory_log_probs=diffusion_output.trajectory_log_probs,
-            trajectory_decoded=diffusion_output.trajectory_decoded,
-            custom_output=postprocess_output.custom_output,
+            latents=trajectory_latents,
+            trajectory_latents=trajectory_latents,
+            trajectory_timesteps=trajectory_timesteps,
+            trajectory_log_probs=trajectory_log_probs,
+            trajectory_decoded=trajectory_decoded,
             multimodal_output=mm_output,
             stage_durations=diffusion_output.stage_durations,
             peak_memory_mb=diffusion_output.peak_memory_mb,
             finished=finished,
         ),
     ]
+
+
+def _trajectory_payload(
+    postprocess_output: DiffusionPostprocessOutput,
+    diffusion_output: DiffusionOutput,
+) -> DiffusionTrajectoryPayload:
+    trajectory: DiffusionTrajectoryPayload = {}
+    payload = postprocess_output.outputs.get("trajectory")
+    if isinstance(payload, Mapping):
+        for source_key, target_key in (
+            ("latents", "latents"),
+            ("timesteps", "timesteps"),
+            ("log_probs", "log_probs"),
+            ("decoded", "decoded"),
+        ):
+            value = payload.get(source_key)
+            if value is not None:
+                trajectory[target_key] = value
+
+    fallback_fields = (
+        ("latents", diffusion_output.trajectory_latents),
+        ("timesteps", diffusion_output.trajectory_timesteps),
+        ("log_probs", diffusion_output.trajectory_log_probs),
+        ("decoded", diffusion_output.trajectory_decoded),
+    )
+    for key, value in fallback_fields:
+        if key not in trajectory and value is not None:
+            trajectory[key] = value
+    return trajectory
diff --git a/vllm_omni/diffusion/output_metadata.py b/vllm_omni/diffusion/output_metadata.py
new file mode 100644
index 000000000..abea8759e
--- /dev/null
+++ b/vllm_omni/diffusion/output_metadata.py
@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Compatibility imports for diffusion output metadata.
+
+The shared output payload and metadata schema lives in
+``vllm_omni.outputs.output_metadata``.  Keep this module as a stable import
+path for existing diffusion callers.
+"""
+
+from vllm_omni.outputs.output_metadata import *  # noqa: F403
diff --git a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py
index 8d4bf41de..5a293b905 100644
--- a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py
+++ b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py
@@ -14,6 +14,7 @@ logger = init_logger(__name__)
 
 def profiler(name: str, func: Callable, instance: Any) -> Callable:
     """Timing a function execution."""
+    metric_name = f"{name.rsplit('.', 1)[0]}.diffuse" if name.endswith(".denoise_step") else name
 
     @functools.wraps(func)
     def wrapper(*args, **kwargs) -> Any:
@@ -28,10 +29,10 @@ def profiler(name: str, func: Callable, instance: Any) -> Callable:
             if current_omni_platform.is_available():
                 current_omni_platform.synchronize()
             duration = time.perf_counter() - start_time
-            logger.info(f"[DiffusionPipelineProfiler] {name} took {duration:.6f}s")
+            logger.info(f"[DiffusionPipelineProfiler] {metric_name} took {duration:.6f}s")
             # record the profiling data: duration of stages
             with instance._profiler_lock:
-                instance._stage_durations[name] = instance._stage_durations.get(name, 0.0) + duration
+                instance._stage_durations[metric_name] = instance._stage_durations.get(metric_name, 0.0) + duration
 
     return wrapper
 
@@ -93,7 +94,11 @@ class DiffusionPipelineProfilerMixin:
         self._stage_durations: dict[str, float] = {}
         targets = profiler_targets if profiler_targets is not None else self._PROFILER_TARGETS
         if not targets:
-            return
+            targets = []
+        else:
+            targets = list(targets)
+        if not profiler_targets and hasattr(self, "denoise_step"):
+            targets.append("denoise_step")
 
         targets = ["forward"] + [
             t for t in targets if t != "forward"
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index 341e62ede..5b8b2fba1 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -276,6 +276,11 @@ _DIFFUSION_MODELS = {
         "pipeline_cosmos3",
         "Cosmos3OmniDiffusersPipeline",
     ),
+    "Cosmos3OmniPipeline": (
+        "cosmos3",
+        "pipeline_cosmos3",
+        "Cosmos3OmniDiffusersPipeline",
+    ),
     "SoulXSingerPipeline": (
         "soulx_singer",
         "pipeline_soulx_singer_svs",
@@ -306,6 +311,11 @@ _DIFFUSION_MODELS = {
         "pipeline_sdxl",
         "StableDiffusionXLPipeline",
     ),
+    "Krea2Pipeline": (
+        "krea2",
+        "pipeline_krea2",
+        "Krea2Pipeline",
+    ),
 }
 
 
@@ -531,15 +541,10 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",
     "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func",
+    "Cosmos3OmniPipeline": "get_cosmos3_post_process_func",
     "HiDreamImagePipeline": "get_hidream_image_post_process_func",
     "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func",
-}
-
-_DIFFUSION_ACTION_POST_PROCESS_FUNCS = {
-    # arch: action_post_process_func
-    # `action_post_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
-    # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key.
-    "Cosmos3OmniDiffusersPipeline": "get_cosmos3_action_post_process_func",
+    "Krea2Pipeline": "get_krea2_post_process_func",
 }
 
 _DIFFUSION_IR_OP_PRIORITY_FUNCS = {
@@ -547,6 +552,7 @@ _DIFFUSION_IR_OP_PRIORITY_FUNCS = {
     # `ir_op_priority_func` function must be placed in {mod_folder}/{mod_relname}.py,
     # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key.
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_ir_op_priority_func",
+    "Cosmos3OmniPipeline": "get_cosmos3_ir_op_priority_func",
 }
 
 _DIFFUSION_PRE_PROCESS_FUNCS = {
@@ -571,6 +577,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func",
     "MagiHumanPipeline": "get_magi_human_pre_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
+    "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",
     "SoulXSingerPipeline": "get_soulxsinger_pre_process_func",
     "SoulXSingerSVCPipeline": "get_soulxsinger_svc_pre_process_func",
 }
@@ -582,8 +589,8 @@ def register_diffusion_model(
     class_name: str,
     pre_process_func_name: str | None = None,
     post_process_func_name: str | None = None,
-    action_post_process_func_name: str | None = None,
     ir_op_priority_func_name: str | None = None,
+    action_post_process_func_name: str | None = None,
 ) -> None:
     """Register a diffusion model pipeline from an out-of-tree plugin.
 
@@ -602,13 +609,23 @@ def register_diffusion_model(
         post_process_func_name: Optional name of the post-process function
             located in *module_name*.  Pass ``None`` to keep the existing
             entry when replacing a built-in model.
-        action_post_process_func_name: Optional name of the action post-process
-            function located in *module_name*.  Pass ``None`` to keep the
-            existing entry when replacing a built-in model.
         ir_op_priority_func_name: Optional name of the IR op priority merge
             function located in *module_name*. Pass ``None`` to keep the
             existing entry when replacing a built-in model.
+        action_post_process_func_name: Deprecated compatibility-only keyword
+            for out-of-tree plugins. Action postprocess hooks are no longer
+            registered separately; move action handling into
+            ``post_process_func_name`` and return a payload/metadata envelope.
     """
+    if action_post_process_func_name is not None:
+        logger.warning(
+            "Ignoring deprecated action_post_process_func_name=%r for diffusion "
+            "model %s. Move action postprocess logic into post_process_func_name "
+            "and return payload/metadata output.",
+            action_post_process_func_name,
+            model_arch,
+        )
+
     # Register model class in DiffusionModelRegistry
     DiffusionModelRegistry.register_model(
         model_arch,
@@ -625,8 +642,6 @@ def register_diffusion_model(
         _DIFFUSION_PRE_PROCESS_FUNCS[model_arch] = pre_process_func_name
     if post_process_func_name is not None:
         _DIFFUSION_POST_PROCESS_FUNCS[model_arch] = post_process_func_name
-    if action_post_process_func_name is not None:
-        _DIFFUSION_ACTION_POST_PROCESS_FUNCS[model_arch] = action_post_process_func_name
     if ir_op_priority_func_name is not None:
         _DIFFUSION_IR_OP_PRIORITY_FUNCS[model_arch] = ir_op_priority_func_name
 
@@ -659,13 +674,6 @@ def get_diffusion_post_process_func(od_config: OmniDiffusionConfig):
     return _load_process_func(od_config, func_name)
 
 
-def get_diffusion_action_post_process_func(od_config: OmniDiffusionConfig):
-    if od_config.model_class_name not in _DIFFUSION_ACTION_POST_PROCESS_FUNCS:
-        return None
-    func_name = _DIFFUSION_ACTION_POST_PROCESS_FUNCS[od_config.model_class_name]
-    return _load_process_func(od_config, func_name)
-
-
 def get_diffusion_ir_op_priority_func(od_config: OmniDiffusionConfig):
     if od_config.model_class_name not in _DIFFUSION_IR_OP_PRIORITY_FUNCS:
         return None
diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py
index 84eb32044..cb16bdeac 100644
--- a/vllm_omni/diffusion/worker/diffusion_model_runner.py
+++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py
@@ -39,7 +39,15 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, DiffusionRequestState, RunnerOutput
+from vllm_omni.diffusion.worker.utils import (
+    BatchRunnerOutput,
+    DiffusionRequestState,
+    RunnerOutput,
+    attach_stage_durations,
+    clear_pipeline_stage_durations,
+    consume_pipeline_stage_durations,
+    merge_stage_durations,
+)
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
@@ -487,6 +495,19 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         return self._runner_output_from_outputs(reqs, outputs)
 
+    def _attach_stepwise_metrics(
+        self,
+        state: DiffusionRequestState,
+        output: DiffusionOutput,
+        *,
+        is_primary: bool,
+    ) -> None:
+        merge_stage_durations(
+            state,
+            consume_pipeline_stage_durations(self.pipeline),
+        )
+        attach_stage_durations(state, output)
+
     def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
         """
         Execute a forward pass for the given requests.
@@ -601,8 +622,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     else:
                         gen_device = self.device
                     state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
+                clear_pipeline_stage_durations(self.pipeline)
                 # encode
                 self.pipeline.prepare_encode(state)
+                merge_stage_durations(
+                    state,
+                    consume_pipeline_stage_durations(self.pipeline),
+                )
 
         input_batch = InputBatch.make_batch(
             states,
@@ -658,19 +684,27 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
+            had_active_states = bool(self.state_cache)
             states, new_request_ids = self._update_states(scheduler_output)
+            is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+            if new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available():
+                current_omni_platform.reset_peak_memory_stats()
             input_batch = self._prepare_batch_inputs(states, new_request_ids)
             attn_metadata = self._prepare_attn_metadata(input_batch)
-            is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
-            if is_primary:
-                current_omni_platform.reset_peak_memory_stats()
 
             with set_forward_context(
                 vllm_config=self.vllm_config,
                 omni_diffusion_config=self.od_config,
                 attn_metadata=attn_metadata,
             ):
+                clear_pipeline_stage_durations(self.pipeline)
                 noise_pred = self.pipeline.denoise_step(input_batch, states=states)
+                denoise_stage_durations = consume_pipeline_stage_durations(self.pipeline)
+                for state in states:
+                    merge_stage_durations(
+                        state,
+                        denoise_stage_durations,
+                    )
 
                 runner_output_list = []
                 pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
@@ -693,14 +727,20 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             self.pipeline.step_scheduler(
                                 req, noise_pred[offset : offset + row_num] if noise_pred is not None else None
                             )
-                            offset = offset + row_num
                             if self.od_config.streaming_output:
                                 should_decode = req.chunk_denoise_completed
                             else:
                                 should_decode = req.denoise_completed
 
                             if should_decode:
+                                clear_pipeline_stage_durations(self.pipeline)
                                 result = self.pipeline.post_decode(req)
+                                if result is not None:
+                                    self._attach_stepwise_metrics(
+                                        req,
+                                        result,
+                                        is_primary=is_primary,
+                                    )
                             else:
                                 result = None
                             # finished should be computed after post_decode() advanced chunk_index
@@ -717,6 +757,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                     result=result,
                                 )
                             )
+                            offset = offset + row_num
                         except Exception as per_req_exc:
                             offset = offset + row_num
                             logger.error(
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index efb280a65..2ce234b69 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/vllm_omni/diffusion/worker/diffusion_worker.py
@@ -219,8 +219,32 @@ class DiffusionWorker:
         self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
         self.stage_id = getattr(od_config, "stage_id", 0)
         self.init_device()
-        # Create model runner using the platform-specified class
-        model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls()
+        # Create model runner — one decision chain, in precedence order:
+        #   1. explicit od_config.diffusion_model_runner_cls (user override),
+        #   2. the runner declared by the engine class that engine_backend
+        #      selects (e.g. ARDiffusionEngine -> ARDiffusionModelRunner),
+        #   3. the platform default.
+        # Routing policy therefore lives on the engine class / config surface;
+        # engines never mutate od_config. Overrides must be import-path
+        # strings — guard with isinstance so a non-string (e.g. a Mock
+        # od_config in tests) doesn't shadow the platform hook.
+        runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None)
+        engine_runner = None
+        if not (isinstance(runner_override, str) and runner_override):
+            try:
+                from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+                engine_cls = DiffusionEngine.resolve_engine_class(self.od_config)
+                engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
+            except Exception:
+                logger.warning("Worker %s: engine_backend resolution failed; using platform runner", self.rank)
+                engine_runner = None
+        if isinstance(runner_override, str) and runner_override:
+            model_runner_cls_path = runner_override
+        elif isinstance(engine_runner, str) and engine_runner:
+            model_runner_cls_path = engine_runner
+        else:
+            model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls()
         model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path)
         self.model_runner = model_runner_cls(
             vllm_config=self.vllm_config,
@@ -270,6 +294,12 @@ class DiffusionWorker:
         # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping.
         # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion.
         vllm_config.kernel_config.ir_op_priority = _resolve_ir_op_priority(self.od_config, vllm_config)
+        if self.od_config.moe_backend != "auto":
+            logger.warning(
+                "Overriding MoE backend from default 'auto' to '%s' per deploy config.",
+                self.od_config.moe_backend,
+            )
+            vllm_config.kernel_config.moe_backend = self.od_config.moe_backend
         logger.info(
             "Final IR op priority after setting vLLM-Omni overrides: %s", vllm_config.kernel_config.ir_op_priority
         )
diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py
index 4fd374d34..ba6165a9d 100644
--- a/vllm_omni/diffusion/worker/utils.py
+++ b/vllm_omni/diffusion/worker/utils.py
@@ -15,6 +15,41 @@ if TYPE_CHECKING:
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
 
+def clear_pipeline_stage_durations(pipeline: Any) -> None:
+    clear_records = getattr(pipeline, "clear_profiler_records", None)
+    if getattr(pipeline, "enable_diffusion_pipeline_profiler", False) and callable(clear_records):
+        clear_records()
+
+
+def consume_pipeline_stage_durations(pipeline: Any) -> dict[str, float]:
+    if not getattr(pipeline, "enable_diffusion_pipeline_profiler", False):
+        return {}
+    stage_durations = getattr(pipeline, "stage_durations", None)
+    if not isinstance(stage_durations, dict):
+        return {}
+    result = {stage: float(duration) for stage, duration in stage_durations.items()}
+    clear_pipeline_stage_durations(pipeline)
+    return result
+
+
+def merge_stage_durations(
+    state: DiffusionRequestState,
+    stage_durations: dict[str, float],
+) -> None:
+    if not stage_durations:
+        return
+    for stage, duration in stage_durations.items():
+        state.stage_durations[stage] = float(state.stage_durations.get(stage, 0.0)) + float(duration)
+
+
+def attach_stage_durations(
+    state: DiffusionRequestState,
+    output: DiffusionOutput,
+) -> None:
+    if state.stage_durations:
+        output.stage_durations = dict(state.stage_durations)
+
+
 @dataclass
 class DiffusionRequestState:
     """Per-request mutable state across all pipeline stages.
@@ -79,6 +114,9 @@ class DiffusionRequestState:
     # For example: Wan condition tensors / masks, or Bagel KV contexts.
     extra: dict[str, Any] = field(default_factory=dict)
 
+    # ── Runner-owned profiling metadata ──
+    stage_durations: dict[str, float] = field(default_factory=dict)
+
     # Peak device memory observed while this request is active in step mode.
     peak_memory_mb: float = 0.0
 
diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
index 304489bbb..51e604336 100644
--- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
+++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py
@@ -856,8 +856,20 @@ class OmniKVTransferManager:
 
             payloads.append(local_payload)
 
+        padding = None
+        if not branch_roles and cfg_size == 2 and main_payload:
+            # AR didn't split KV by branch; cfg follower needs the same
+            # positive KV for shared-prefix reuse. Only safe for cfg_size == 2:
+            # for cfg_size > 2 every rank must get its own KV, which we can't
+            # assign here, so padding stays None and the runner raises.
+            # Shallow dict copy: shares PKV tensors with main_payload. Safe
+            # only because the send path (cfg_scatter -> pack -> to_gpu_tensor)
+            # reads tensors to serialize, never mutates in place. Any async/
+            # KV-pool path that mutates req.past_key_values mid-scatter must
+            # clone tensors here instead.
+            padding = dict(main_payload)
         while len(payloads) < cfg_size:
-            payloads.append(None)
+            payloads.append(padding)
 
         return payloads[:cfg_size]
 
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index 9097237fd..bdf43711c 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/vllm_omni/engine/arg_utils.py
@@ -9,7 +9,7 @@ from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.logger import init_logger
 
 from vllm_omni.config import OmniModelConfig
-from vllm_omni.engine.output_modality import OutputModality
+from vllm_omni.outputs.output_modality import OutputModality
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.plugins import load_omni_general_plugins
 
@@ -434,6 +434,9 @@ class OrchestratorArgs:
     stage_configs_path: str | None = None
     deploy_config: str | None = None
     stage_overrides: str | None = None  # raw JSON string; parsed downstream
+    # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays
+    # derived sizing onto merged stages, then drops it before per-stage engine args.
+    strategy_config: str | None = None
 
     # === Mode Switches (orchestrator reads, DeployConfig redistributes) ===
     async_chunk: bool | None = None
diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py
index 77e0a8405..ce9662884 100644
--- a/vllm_omni/engine/async_omni_engine.py
+++ b/vllm_omni/engine/async_omni_engine.py
@@ -30,6 +30,7 @@ from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.input_processor import InputProcessor
 
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.stage_config import strip_parent_engine_args
 from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config
 from vllm_omni.diffusion.diffusion_engine import supports_audio_output
@@ -57,7 +58,10 @@ from vllm_omni.engine.stage_runtime import (
     create_stage_runtime,
 )
 from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
-from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
+from vllm_omni.entrypoints.utils import (
+    load_and_resolve_stage_configs,
+    parse_stage_overrides,
+)
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.metrics.prometheus import OmniRequestCounter
 
@@ -216,6 +220,7 @@ class AsyncOmniEngine:
         transfer_emitter: Any = None,
         log_stats: bool = False,
         tokenizer: str | None = None,
+        trust_remote_code: bool = False,
         **kwargs: Any,
     ) -> None:
         self.model = model
@@ -273,6 +278,17 @@ class AsyncOmniEngine:
                 self._omni_master_port,
             )
 
+        # Stage resolution pops deploy_config, so get the pipeline endpoint
+        # restriction beforehand. TODO (Alex) make this cleaner and refactor
+        # stage config resolution to remove kwargs hacks.
+        deploy_config_path = kwargs.get("deploy_config")
+        self.endpoint_restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions(
+            model=model,
+            trust_remote_code=trust_remote_code,
+            deploy_config_path=deploy_config_path,
+        )
+
+        kwargs["trust_remote_code"] = trust_remote_code
         self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs)
 
         self.num_stages = len(self.stage_configs)
@@ -1027,6 +1043,7 @@ class AsyncOmniEngine:
             "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True),
             "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4),
             "quantization": kwargs.get("quantization", None),
+            "quantization_config": kwargs.get("quantization_config", None),
             "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None),
             "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None),
             "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None),
@@ -1110,11 +1127,41 @@ class AsyncOmniEngine:
 
         return result
 
+    def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None:
+        """Apply a strategy-derived ``omni_lb_policy`` to the engine.
+
+        Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is
+        the engine default and is treated as "unset" (indistinguishable from no
+        flag), so a strategy value overrides it. If the user explicitly passed a
+        non-default policy that conflicts with the strategy-derived one, raise so
+        the mismatch is not silently ignored.
+        """
+        if not derived:
+            return
+        explicit = kwargs.get("omni_lb_policy")
+        user_set = explicit is not None and str(explicit) != "random"
+        if user_set:
+            if str(explicit) != str(derived):
+                raise ValueError(
+                    f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given "
+                    f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. "
+                    "Drop --omni-lb-policy to use the strategy value, or make them match."
+                )
+            return
+        if self._omni_lb_policy != str(derived):
+            logger.info(
+                "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).",
+                derived,
+                self._omni_lb_policy,
+            )
+            self._omni_lb_policy = str(derived)
+
     def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]:
         """Resolve stage configs and inject defaults shared by orchestrator/headless."""
 
         stage_configs_path = kwargs.get("stage_configs_path", None)
         deploy_config_path = kwargs.pop("deploy_config", None)
+        strategy_config_path = kwargs.pop("strategy_config", None)
         stage_overrides_json = kwargs.pop("stage_overrides", None)
         explicit_stage_configs = kwargs.pop("stage_configs", None)
         if explicit_stage_configs is not None:
@@ -1129,27 +1176,23 @@ class AsyncOmniEngine:
             base_kwargs = kwargs
 
         # Parse --stage-overrides JSON string if provided
-        stage_overrides = None
-        if stage_overrides_json:
-            if isinstance(stage_overrides_json, str):
-                try:
-                    stage_overrides = json.loads(stage_overrides_json)
-                except json.JSONDecodeError as exc:
-                    raise ValueError(
-                        f"--stage-overrides is not valid JSON: {exc}. Got: {stage_overrides_json!r}"
-                    ) from exc
-            else:
-                stage_overrides = stage_overrides_json
+        stage_overrides = parse_stage_overrides(stage_overrides_json)
 
-        config_path, stage_configs = load_and_resolve_stage_configs(
+        config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
             model,
             stage_configs_path,
             base_kwargs,
             default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
             deploy_config_path=deploy_config_path,
             stage_overrides=stage_overrides,
+            strategy_config_path=strategy_config_path,
         )
 
+        # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is
+        # an orchestrator-level knob (read once at construction), so apply it here
+        # rather than as a per-stage config field.
+        self._apply_strategy_lb_policy(strategy_lb_policy, kwargs)
+
         # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
         for cfg in stage_configs:
             try:
@@ -1191,7 +1234,7 @@ class AsyncOmniEngine:
                             kwargs.get("diffusion_attention_config"),
                             attention_backend=kwargs.get("diffusion_attention_backend"),
                         )
-                quantization_config = kwargs.get("diffusion_quantization_config")
+                quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config")
                 if quantization_config is not None:
                     if (
                         not hasattr(cfg.engine_args, "quantization_config")
@@ -1525,16 +1568,18 @@ class AsyncOmniEngine:
         # and can cause CUDA OOM for subsequent engine instances (especially
         # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB).
         #
-        # Discard mode (level=2) is correct at shutdown: there is no benefit to
-        # keeping a CPU backup when the engine is being torn down.
+        # CuMemAllocator.sleep() is NOT idempotent — calling it on already-
+        # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator
+        # cuMemRelease (double-free of the memory handle).  Use release_pools()
+        # instead, which is the designed cleanup path: it drops MemPool refs
+        # and lets the destructor/free path handle asleep entries correctly
+        # (returns a null handle so the C extension skips unmap/release).
         try:
             from vllm.device_allocator.cumem import CuMemAllocator, cumem_available
 
             if cumem_available:
                 allocator = CuMemAllocator.get_instance()
-                # Sleep at level 2 discards all pool memory from the GPU
-                # without creating CPU backups — cheapest and fastest.
-                allocator.sleep()
+                allocator.release_pools()
                 logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown")
         except Exception:
             pass
diff --git a/vllm_omni/engine/mm_outputs.py b/vllm_omni/engine/mm_outputs.py
index 3c9da2496..336bf7d9a 100644
--- a/vllm_omni/engine/mm_outputs.py
+++ b/vllm_omni/engine/mm_outputs.py
@@ -1,127 +1,9 @@
-"""Multimodal output data structures for vLLM-Omni.
+import warnings
 
-This module defines structured types for multimodal outputs.
+from vllm_omni.outputs.mm_outputs import *  # noqa: F401,F403
 
-"""
-
-from __future__ import annotations
-
-from collections.abc import Iterator, Mapping
-from dataclasses import dataclass, field
-from typing import Any
-
-import torch
-from vllm.outputs import CompletionOutput
-
-
-@dataclass(eq=False)
-class MultimodalPayload(Mapping):
-    """Structured multimodal output payload.
-
-    Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)``
-    style checks in downstream code can be replaced with duck-typing, and
-    ``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)``
-    all work seamlessly for both tensors and metadata.
-
-    Attributes:
-        tensors: Dictionary mapping modality/key names to their tensors.
-        metadata: Optional dictionary for non-tensor metadata
-            (e.g., sample rate for audio, image dimensions).
-    """
-
-    tensors: dict[str, torch.Tensor] = field(default_factory=dict)
-    metadata: dict[str, Any] = field(default_factory=dict)
-
-    @property
-    def primary_tensor(self) -> torch.Tensor | None:
-        """Return the first tensor in the payload, or None if empty."""
-        if self.tensors:
-            return next(iter(self.tensors.values()))
-        return None
-
-    @property
-    def is_empty(self) -> bool:
-        """Return True if the payload has no tensors and no metadata."""
-        return len(self.tensors) == 0 and len(self.metadata) == 0
-
-    def get(self, key: str, default: Any = None) -> Any:
-        """Get a value by key, searching tensors first then metadata."""
-        if key in self.tensors:
-            return self.tensors[key]
-        return self.metadata.get(key, default)
-
-    def __contains__(self, key: str) -> bool:
-        return key in self.tensors or key in self.metadata
-
-    def __getitem__(self, key: str) -> Any:
-        """Dict-like indexing: search tensors first, then metadata."""
-        if key in self.tensors:
-            return self.tensors[key]
-        if key in self.metadata:
-            return self.metadata[key]
-        raise KeyError(key)
-
-    def __len__(self) -> int:
-        return len(self.tensors) + len(self.metadata)
-
-    def __iter__(self) -> Iterator[str]:
-        yield from self.tensors
-        yield from self.metadata
-
-    def __bool__(self) -> bool:
-        return bool(self.tensors) or bool(self.metadata)
-
-    def __eq__(self, other: object) -> bool:
-        """Support equality with plain dicts and other Mappings."""
-        if isinstance(other, MultimodalPayload):
-            return self.tensors == other.tensors and self.metadata == other.metadata
-        if isinstance(other, Mapping):
-            return dict(self) == dict(other)
-        return NotImplemented
-
-    def to_dict(self) -> dict[str, Any]:
-        """Convert back to a plain dict (tensors + metadata merged)."""
-        result: dict[str, Any] = dict(self.tensors)
-        result.update(self.metadata)
-        return result
-
-    @classmethod
-    def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None:
-        """Create a MultimodalPayload from a raw dictionary.
-
-        Separates torch.Tensor values into tensors and everything
-        else into metadata.
-        """
-        if not data:
-            return None
-        tensors: dict[str, torch.Tensor] = {}
-        metadata: dict[str, Any] = {}
-        for k, v in data.items():
-            if isinstance(v, torch.Tensor):
-                tensors[k] = v
-            else:
-                metadata[k] = v
-        if not tensors and not metadata:
-            return None
-        return cls(tensors=tensors, metadata=metadata)
-
-
-@dataclass
-class MultimodalCompletionOutput(CompletionOutput):
-    """CompletionOutput with multimodal support.
-
-    Inherits all CompletionOutput fields and adds multimodal_output.
-    As a CompletionOutput subclass, compatible with all existing vLLM consumers.
-    """
-
-    def __init__(
-        self,
-        multimodal_output: MultimodalPayload | None = None,
-        **kwargs: Any,
-    ):
-        super().__init__(**kwargs)
-        self.multimodal_output = multimodal_output
-
-    def __repr__(self) -> str:
-        base = super().__repr__()
-        return f"{base[:-1]}, multimodal_output={self.multimodal_output!r})"
+warnings.warn(
+    "Importing from 'vllm_omni.engine.mm_outputs' is deprecated. Use 'vllm_omni.outputs.mm_outputs' instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py
index df05cb669..bdd678f27 100644
--- a/vllm_omni/engine/orchestrator.py
+++ b/vllm_omni/engine/orchestrator.py
@@ -1287,7 +1287,7 @@ class Orchestrator:
                             [req_id, *self._cfg_tracker.cleanup_parent(req_id)],
                         )
                         return
-                    if already_submitted and len(diffusion_prompt) == 1:
+                    if len(diffusion_prompt) == 1:
                         diffusion_prompt = diffusion_prompt[0]
             else:
                 diffusion_prompt = req_state.prompt
diff --git a/vllm_omni/engine/output_modality.py b/vllm_omni/engine/output_modality.py
index ac72910dd..364519bed 100644
--- a/vllm_omni/engine/output_modality.py
+++ b/vllm_omni/engine/output_modality.py
@@ -1,123 +1,9 @@
-"""Output modality types for vLLM-Omni.
+import warnings
 
-This module defines the OutputModality enum and TensorAccumulationStrategy
-for type-safe multimodal output routing and tensor merging.
+from vllm_omni.outputs.output_modality import *  # noqa: F401,F403
 
-"""
-
-from __future__ import annotations
-
-import re
-from enum import Enum, Flag, StrEnum, auto
-from typing import Literal, TypeAlias
-
-FinalOutputModalityType: TypeAlias = Literal["text", "image", "audio", "video"]
-
-_MODALITY_ALIASES: dict[str, str] = {
-    "speech": "audio",
-    "images": "image",
-    "latents": "latent",
-    "wav": "audio",
-    "waveform": "audio",
-    "pixel_values": "image",
-    "pixels": "image",
-    "token_ids": "text",
-    "tokens": "text",
-}
-
-
-class OutputModalityNames(StrEnum):
-    """Keys for output modalities.
-
-    TODO: (Alex) Integrate this with the big-flag enum below + throughout the code
-    for better type safety (currently only used for output processor).
-    """
-
-    TEXT = "text"
-    IMAGE = "image"
-    AUDIO = "audio"
-    LATENT = "latent"
-
-
-# Specify which output modalities may be drained when handling delta messages.
-# For some types, e.g., latents, we need to be careful to ensure the full context
-# is passed as the stream yields due to assumptions in the I/O processing and model
-# when async chunk isn't enabled.
-NON_DRAINABLE_MODALITIES = {OutputModalityNames.TEXT, OutputModalityNames.LATENT}
-DRAINABLE_MODALITIES = {mod for mod in OutputModalityNames if mod not in NON_DRAINABLE_MODALITIES}
-
-
-class OutputModality(Flag):
-    """Bit-flag enum for output modalities.
-
-    Compose freely with ``|`` — no need to enumerate every combination.
-
-    Single:   ``OutputModality.TEXT``, ``OutputModality.IMAGE``, ...
-    Compound: ``OutputModality.TEXT | OutputModality.IMAGE``  (text+image)
-
-    Note: POOLING is intentionally excluded. Pooling/embedding is vLLM's
-    native path (pooling_output → PoolingRequestOutput), handled entirely
-    by the base OutputProcessor. vLLM-Omni's layer does not participate.
-    """
-
-    TEXT = auto()
-    IMAGE = auto()
-    AUDIO = auto()
-    LATENT = auto()
-
-    @classmethod
-    def from_string(cls, s: str | None) -> OutputModality:
-        """Parse a free-text modality string into an OutputModality flag.
-
-        Handles common aliases and compound strings separated by + or ,.
-
-        Examples::
-
-            OutputModality.from_string("text+image")
-            # → OutputModality.TEXT | OutputModality.IMAGE
-        """
-        if not s or not s.strip():
-            return cls.TEXT
-
-        parts = [p.strip().lower() for p in re.split(r"[+,]", s.strip())]
-        result = cls(0)
-        for p in parts:
-            p = _MODALITY_ALIASES.get(p, p)
-            try:
-                result |= cls[p.upper()]
-            except KeyError:
-                raise ValueError(f"Unknown modality: {p!r}. Supported: {[m.name.lower() for m in cls]}")
-        return result
-
-    @property
-    def has_text(self) -> bool:
-        return OutputModality.TEXT in self
-
-    @property
-    def has_multimodal(self) -> bool:
-        return bool(self & ~OutputModality.TEXT)
-
-
-class TensorAccumulationStrategy(Enum):
-    """Strategy for merging incremental multimodal tensors."""
-
-    CONCAT_DIM0 = "concat_dim0"
-    """Concatenate along dimension 0. Used for image/latent tensors."""
-
-    CONCAT_LAST = "concat_last"
-    """Concatenate along the last dimension. Used for audio waveforms."""
-
-    APPEND_LIST = "append_list"
-    """Append to a list (no tensor concatenation)."""
-
-    REPLACE = "replace"
-    """Replace previous tensor entirely with the latest one."""
-
-
-def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy:
-    """Determine tensor merge strategy from the multimodal flags."""
-    if OutputModality.AUDIO in modality:
-        return TensorAccumulationStrategy.CONCAT_LAST
-    if OutputModality.IMAGE in modality or OutputModality.LATENT in modality:
-        return TensorAccumulationStrategy.CONCAT_DIM0
-    return TensorAccumulationStrategy.CONCAT_DIM0  # default
+warnings.warn(
+    "Importing from 'vllm_omni.engine.output_modality' is deprecated. Use 'vllm_omni.outputs.output_modality' instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py
index c9044e29e..0d87583a8 100644
--- a/vllm_omni/engine/output_processor.py
+++ b/vllm_omni/engine/output_processor.py
@@ -1,744 +1,10 @@
-from dataclasses import fields as dataclass_fields
-from typing import Any
+import warnings
 
-import torch
-from vllm.logger import init_logger
-from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput
-from vllm.sampling_params import RequestOutputKind
-from vllm.tokenizers import TokenizerLike
-from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
-from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
-from vllm.v1.engine.output_processor import (
-    OutputProcessorOutput,
-    RequestOutputCollector,
-    RequestState,
+from vllm_omni.outputs.output_processor import *  # noqa: F401,F403
+
+warnings.warn(
+    "Importing from 'vllm_omni.engine.output_processor' is deprecated. "
+    "Use 'vllm_omni.outputs.output_processor' instead.",
+    DeprecationWarning,
+    stacklevel=2,
 )
-from vllm.v1.engine.parallel_sampling import ParentRequest
-from vllm.v1.metrics.stats import IterationStats, RequestStateStats
-
-from vllm_omni.data_entry_keys import unflatten_payload
-from vllm_omni.engine.mm_outputs import MultimodalCompletionOutput, MultimodalPayload
-from vllm_omni.engine.output_modality import (
-    DRAINABLE_MODALITIES,
-    OutputModality,
-    TensorAccumulationStrategy,
-    get_accumulation_strategy,
-)
-from vllm_omni.outputs import OmniRequestOutput
-
-logger = init_logger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# Module-level helpers
-# ---------------------------------------------------------------------------
-
-
-def _to_cpu(x: Any) -> Any:
-    """Move a tensor to CPU; pass through non-tensors unchanged."""
-    if isinstance(x, torch.Tensor):
-        try:
-            return x.detach().to("cpu", non_blocking=True).contiguous()
-        except (RuntimeError, AttributeError):
-            return x
-    if isinstance(x, dict):
-        return {str(sk): _to_cpu(sv) for sk, sv in x.items()}
-    return x
-
-
-def _merge_payload(target: MultimodalPayload, incoming: MultimodalPayload) -> None:
-    """Merge *incoming* into *target* using list-based deferred concatenation."""
-    for k, v in incoming.tensors.items():
-        if k not in target.tensors:
-            target.tensors[k] = v
-        else:
-            existing = target.tensors[k]
-            if isinstance(existing, list):
-                existing.append(v)
-            elif isinstance(existing, torch.Tensor) and isinstance(v, torch.Tensor):
-                target.tensors[k] = [existing, v]
-            else:
-                target.tensors[k] = v
-
-    for k, v in incoming.metadata.items():
-        if k not in target.metadata:
-            target.metadata[k] = v
-        else:
-            existing = target.metadata[k]
-            if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor):
-                target.metadata[k] = [existing, v]
-            elif isinstance(v, torch.Tensor) and isinstance(existing, list):
-                existing.append(v)
-            else:
-                # Metadata: replace with latest value
-                target.metadata[k] = v
-
-
-def _cat_tensors(
-    tensors: list[torch.Tensor],
-    strategy: TensorAccumulationStrategy,
-) -> torch.Tensor:
-    """Concatenate a list of tensors according to *strategy*."""
-    if strategy == TensorAccumulationStrategy.CONCAT_LAST:
-        return torch.cat(tensors, dim=-1)
-    if strategy == TensorAccumulationStrategy.REPLACE:
-        return tensors[-1]
-    # CONCAT_DIM0 / APPEND_LIST / default
-    return torch.cat(tensors, dim=0)
-
-
-def _modality_to_type_string(modality: OutputModality) -> str:
-    """Convert an OutputModality flag to a lowercase type string."""
-    try:
-        if OutputModality.AUDIO in modality:
-            return "audio"
-        if OutputModality.IMAGE in modality:
-            return "image"
-        if OutputModality.LATENT in modality:
-            return "latent"
-    except TypeError:
-        # Flag identity mismatch (e.g. after module reload in tests).
-        name = getattr(modality, "name", "") or ""
-        lowered = name.lower()
-        if "audio" in lowered:
-            return "audio"
-        if "image" in lowered:
-            return "image"
-        if "latent" in lowered:
-            return "latent"
-    return "text"
-
-
-class OmniRequestState(RequestState):
-    """Request state for omni models, tracking multimodal outputs.
-
-    Extends the base RequestState with support for accumulating
-    multimodal tensor outputs (e.g., images, audio, latents) that
-    are produced incrementally during generation.
-    """
-
-    def __init__(
-        self,
-        *args,
-        **kwargs,
-    ):
-        # arrival_time is always passed as a keyword argument by the sole caller
-        # (from_new_request).  The *args fallback via positional index was removed
-        # because upstream RequestState.__init__ reorders parameters between
-        # releases, making args[12] fragile.
-        arrival_time = kwargs.get("arrival_time")
-        super().__init__(*args, **kwargs)
-        self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0))
-        # Omni-specific: multimodal output accumulation
-        # TODO: mm_type is per-request, not per-key. If a model ever produces
-        # both audio and latent outputs, the modality type would flip on each
-        # add_multimodal_tensor() call. Consider tracking per-key modality
-        # types (e.g. dict[str, str]) for future multi-output models.
-        self.mm_type: str | None = None
-        self.mm_accumulated: MultimodalPayload = MultimodalPayload()
-
-    def apply_streaming_update(self, update) -> None:
-        super().apply_streaming_update(update)
-        self.native_text_stats.arrival_time = update.arrival_time
-
-    @staticmethod
-    def _to_cpu(x):
-        """Try to convert to CPU tensor if needed."""
-        if isinstance(x, torch.Tensor):
-            try:
-                return x.detach().to("cpu", non_blocking=True).contiguous()
-            except Exception:
-                return x
-        return x
-
-    def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
-        """Accumulate a multimodal tensor payload into the request state.
-
-        Normalizes incoming payloads (dict or raw tensor) into a
-        MultimodalPayload and merges with any previously accumulated data.
-        Uses list-based deferred concatenation to avoid O(n²) repeated
-        torch.cat calls.
-        """
-        if payload is None:
-            return
-        try:
-            if mm_type:
-                self.mm_type = (mm_type or "").lower()
-
-            target_key = self.mm_type or "hidden"
-
-            # Build incoming MultimodalPayload
-            if isinstance(payload, dict):
-                remapped: dict[str, Any] = {}
-                for k, v in payload.items():
-                    # Normalize producer keys to the modality name.
-                    # AR runners produce {"hidden": ...} and generation
-                    # runners produce {"model_outputs": ...}; remap both
-                    # to the semantic modality key (e.g. "audio", "latent").
-                    if k == "model_outputs":
-                        k = target_key
-                    elif k == "hidden" and target_key != "hidden":
-                        k = target_key
-                    remapped[k] = _to_cpu(v)
-                incoming = MultimodalPayload.from_dict(remapped)
-            elif isinstance(payload, MultimodalPayload):
-                incoming = payload
-            else:
-                incoming = MultimodalPayload.from_dict({target_key: _to_cpu(payload)})
-
-            if incoming is None:
-                return
-
-            if self.mm_accumulated.is_empty:
-                self.mm_accumulated = incoming
-            else:
-                _merge_payload(self.mm_accumulated, incoming)
-        except (ValueError, TypeError, RuntimeError):
-            logger.exception("Error accumulating multimodal tensor")
-
-    def _consolidate_multimodal_tensors(self) -> None:
-        """Consolidate accumulated tensor lists into single tensors.
-
-        Uses TensorAccumulationStrategy derived from the output modality
-        to determine concatenation behavior. Metadata values always use
-        REPLACE (keep latest).
-        """
-        if self.mm_accumulated.is_empty:
-            return
-
-        try:
-            modality = OutputModality.from_string(self.mm_type)
-        except (ValueError, KeyError):
-            modality = OutputModality.TEXT
-        strategy = get_accumulation_strategy(modality)
-
-        try:
-            for k, v in list(self.mm_accumulated.tensors.items()):
-                if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
-                    try:
-                        self.mm_accumulated.tensors[k] = _cat_tensors(v, strategy)
-                    except RuntimeError:
-                        if k == "audio":
-                            try:
-                                self.mm_accumulated.tensors[k] = torch.cat(v, dim=-1)
-                            except RuntimeError:
-                                self.mm_accumulated.tensors[k] = torch.cat([t.reshape(-1) for t in v], dim=0)
-                        else:
-                            logger.warning(
-                                "Error concatenating tensor for key %s; keeping last tensor",
-                                k,
-                            )
-                            self.mm_accumulated.tensors[k] = v[-1]
-
-            # Metadata: consolidate any deferred tensor lists (REPLACE semantics)
-            for k, v in list(self.mm_accumulated.metadata.items()):
-                if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
-                    self.mm_accumulated.metadata[k] = v[-1]
-                elif isinstance(v, dict):
-                    for sk, sv in list(v.items()):
-                        if isinstance(sv, list) and sv and isinstance(sv[0], torch.Tensor):
-                            v[sk] = sv[-1]
-        except (RuntimeError, TypeError, KeyError):
-            logger.exception("Error consolidating multimodal tensors")
-
-    # Override: do not route to pooling-only path; always create completion
-    # outputs, and attach pooling_result into the CompletionOutput.
-    def make_request_output(
-        self,
-        new_token_ids: list[int],
-        pooling_output: torch.Tensor | None,
-        finish_reason: FinishReason | None,
-        stop_reason: int | str | None,
-        kv_transfer_params: dict[str, Any] | None = None,
-        routed_experts: Any = None,
-    ) -> OmniRequestOutput | PoolingRequestOutput | None:
-        """Create a request output from generation results.
-
-        Creates a RequestOutput or PoolingRequestOutput from the generated
-        tokens and accumulated multimodal outputs. Attaches multimodal
-        tensors to the completion output if available.
-
-        Args:
-            new_token_ids: List of newly generated token IDs
-            pooling_output: Optional pooling output tensor
-            finish_reason: Optional finish reason indicating why generation stopped
-            stop_reason: Optional stop reason (token ID or stop string)
-            kv_transfer_params: Optional KV cache transfer parameters
-            routed_experts: Optional MoE routed-expert ids for this step,
-                attached to the completion output for generation stages
-
-        Returns:
-            OmniRequestOutput or PoolingRequestOutput if output should be
-            emitted (based on finish status and output kind), None otherwise
-        """
-        # Pooling-only requests should follow base behavior.
-        if self.detokenizer is None and pooling_output is not None:
-            return super().make_request_output(
-                new_token_ids,
-                pooling_output,
-                finish_reason,
-                stop_reason,
-                kv_transfer_params,
-            )
-
-        finished = finish_reason is not None
-        final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
-        is_delta = self.output_kind == RequestOutputKind.DELTA
-
-        if not finished and final_only:
-            return None
-
-        # Consolidate accumulated tensors when finishing, or for every
-        # CUMULATIVE step so the consumer always sees a single tensor.
-        if finished or not is_delta:
-            self._consolidate_multimodal_tensors()
-
-        if self.stream_interval > 1 and self.detokenizer is not None:
-            # Send output request only when
-            # 1. It has finished, or
-            # 2. It is the first token, or
-            # 3. It has reached the stream interval number of tokens
-            if not (
-                finished
-                or self.sent_tokens_offset == 0
-                or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval
-            ):
-                return None
-
-            if self.output_kind == RequestOutputKind.DELTA:
-                # Send tokens from the offset in DELTA mode, otherwise all
-                # tokens are sent.
-                new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :]
-                self.sent_tokens_offset = self.detokenizer.num_output_tokens()
-
-        external_req_id = self.external_req_id
-
-        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts)
-
-        if self.parent_req is None:
-            outputs = [output]
-        else:
-            outputs, finished = self.parent_req.get_outputs(self.request_id, output)
-            if not outputs:
-                return None
-            external_req_id = self.parent_req.external_req_id
-
-        return self._new_request_output(
-            external_req_id,
-            outputs,
-            finished,
-            kv_transfer_params,
-        )
-
-    def _new_completion_output(
-        self,
-        token_ids: list[int],
-        finish_reason: FinishReason | None,
-        stop_reason: int | str | None,
-        routed_experts: Any = None,
-    ) -> MultimodalCompletionOutput | CompletionOutput:
-        """Create a completion output with multimodal data attached.
-
-        Returns a MultimodalCompletionOutput when multimodal data has been
-        accumulated, otherwise returns the base CompletionOutput.  Snapshots
-        the accumulated payload before draining DELTA-mode modality keys so
-        that callers receive an immutable view.
-        """
-        # When there is no detokenizer (generation stages), build a minimal
-        # CompletionOutput directly since upstream asserts detokenizer != None.
-        if self.detokenizer is None:
-            finished = finish_reason is not None
-            base_output = CompletionOutput(
-                index=self.request_index,
-                text="",
-                token_ids=token_ids,
-                logprobs=None,
-                cumulative_logprob=None,
-                finish_reason=str(finish_reason) if finished else None,
-                stop_reason=stop_reason if finished else None,
-                routed_experts=routed_experts,
-            )
-        else:
-            base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
-
-        # Always provide cumulative token IDs for inter-stage processors.
-        if self.detokenizer is not None:
-            base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids)
-        else:
-            base_output.cumulative_token_ids = list(token_ids)
-
-        # Attach cumulative_text only at the final step for inter-stage use.
-        if finish_reason is not None and hasattr(self.detokenizer, "output_text"):
-            base_output.cumulative_text = self.detokenizer.output_text
-
-        try:
-            if self.mm_accumulated and not self.mm_accumulated.is_empty:
-                # Snapshot: copy current tensors/metadata so drain doesn't
-                # mutate the payload already handed to the caller.
-                # Unflatten dotted keys (e.g. "hidden_states.layer_0") back
-                # to nested dicts so downstream consumers (thinker2talker etc.)
-                # can access with .get("hidden_states", {}).get("layers", {}).
-                merged = {**self.mm_accumulated.tensors, **self.mm_accumulated.metadata}
-                nested = unflatten_payload(merged)
-                snapshot = MultimodalPayload.from_dict(nested) or MultimodalPayload()
-                kwargs = {f.name: getattr(base_output, f.name) for f in dataclass_fields(CompletionOutput)}
-                output = MultimodalCompletionOutput(
-                    multimodal_output=snapshot,
-                    **kwargs,
-                )
-                output.cumulative_token_ids = base_output.cumulative_token_ids
-                if hasattr(base_output, "cumulative_text"):
-                    output.cumulative_text = base_output.cumulative_text
-
-                # DELTA mode: drain modality keys (e.g. audio) so the next
-                # step only sees freshly accumulated data for those keys.
-                if self.output_kind == RequestOutputKind.DELTA:
-                    for modality_key in DRAINABLE_MODALITIES:
-                        self.mm_accumulated.tensors.pop(modality_key, None)
-
-                return output
-        except (RuntimeError, TypeError, AttributeError):
-            logger.exception("Error creating MultimodalCompletionOutput")
-        return base_output
-
-    def _new_request_output(
-        self,
-        external_req_id: str,
-        outputs: list,
-        finished: bool,
-        kv_transfer_params: dict[str, Any] | None = None,
-    ) -> RequestOutput | PoolingRequestOutput:
-        """Create request output, handling no-detokenizer generation stages.
-
-        Upstream asserts ``self.logprobs_processor is not None`` which fails
-        for generation stages (talker, code2wav) that have no tokenizer.
-        When the logprobs processor is absent we build the RequestOutput
-        directly with ``prompt_logprobs=None``.
-        """
-        if self.logprobs_processor is not None:
-            return super()._new_request_output(
-                external_req_id,
-                outputs,
-                finished,
-                kv_transfer_params,
-            )
-
-        # No-detokenizer path: build RequestOutput directly.
-        prompt_token_ids = self.prompt_token_ids
-        if prompt_token_ids is None and self.prompt_embeds is not None:
-            prompt_token_ids = [0] * len(self.prompt_embeds)
-        if prompt_token_ids is None:
-            prompt_token_ids = []
-
-        return RequestOutput(
-            request_id=external_req_id,
-            lora_request=self.lora_request,
-            prompt=self.prompt,
-            prompt_token_ids=prompt_token_ids,
-            prompt_logprobs=None,
-            outputs=outputs,
-            finished=finished,
-            kv_transfer_params=kv_transfer_params,
-            num_cached_tokens=self.num_cached_tokens,
-            metrics=self.stats,
-        )
-
-
-class MultimodalOutputProcessor(VLLMOutputProcessor):
-    """Handles multimodal output processing.
-
-    Captures multimodal outputs from OmniEngineCoreOutput and accumulates
-    them as MultimodalPayload in OmniRequestState, before delegating to
-    the base vLLM OutputProcessor for text handling.
-
-    The data flow is:
-    1. For each EngineCoreOutput with multimodal_output:
-       - Capture into OmniRequestState.add_multimodal_tensor()
-    2. Base vLLM OutputProcessor handles text detokenization
-    3. On finish, _consolidate_multimodal_tensors() concatenates accumulated
-       tensors using strategy-based dispatch
-    4. _new_completion_output() returns MultimodalCompletionOutput
-    """
-
-    def __init__(
-        self,
-        tokenizer: TokenizerLike | None,
-        *,
-        log_stats: bool,
-        stream_interval: int = 1,
-        tracing_enabled: bool = False,
-        engine_core_output_type: str | None = None,
-        output_modality: OutputModality = OutputModality.TEXT,
-    ):
-        """Initialize the multimodal output processor.
-
-        Args:
-            tokenizer: Tokenizer for detokenizing text outputs
-            log_stats: Whether to log statistics
-            stream_interval: Stream interval for output generation
-            engine_core_output_type: Optional output type string (e.g.,
-                "image", "audio", "latent"). Converted to OutputModality
-                internally. Kept for backward compatibility with
-                stage_init_utils.
-            output_modality: Type-safe output modality flag. Used to tag
-                multimodal outputs with the correct modality key when
-                per-output type info is unavailable.
-        """
-        super().__init__(
-            tokenizer=tokenizer,
-            log_stats=log_stats,
-            stream_interval=stream_interval,
-            tracing_enabled=tracing_enabled,
-        )
-        # Convert string-based engine_core_output_type to OutputModality
-        if engine_core_output_type is not None:
-            self.output_modality = OutputModality.from_string(engine_core_output_type)
-        else:
-            self.output_modality = output_modality
-        self.engine_core_output_type = engine_core_output_type
-        self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {}
-
-    def _native_text_metric_record(self, request_id: str) -> dict[str, Any]:
-        return self._native_text_metrics_by_request.setdefault(
-            request_id,
-            {
-                "vllm_ttft_ms": 0.0,
-                "vllm_tpot_ms": 0.0,
-                "vllm_itl_ms": 0.0,
-                "vllm_itls_ms": [],
-            },
-        )
-
-    def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]:
-        return self._native_text_metrics_by_request.pop(request_id, {})
-
-    def abort_requests(self, request_ids, internal: bool) -> list[str]:
-        request_ids = list(request_ids)
-        for request_id in request_ids:
-            if internal:
-                req_state = self.request_states.get(request_id)
-                if req_state is not None:
-                    self._native_text_metrics_by_request.pop(req_state.external_req_id, None)
-            else:
-                self._native_text_metrics_by_request.pop(request_id, None)
-        return super().abort_requests(request_ids, internal)
-
-    def add_request(
-        self,
-        request: EngineCoreRequest,
-        prompt: str | None,
-        parent_req: ParentRequest | None = None,
-        request_index: int = 0,
-        queue: RequestOutputCollector | None = None,
-    ) -> None:
-        """Add a new request to be processed.
-
-        Creates an OmniRequestState for the request and registers it
-        for output processing.
-
-        Args:
-            request: Engine core request to add
-            prompt: Optional prompt string for the request
-            parent_req: Optional parent request for parallel sampling
-            request_index: Index of the request in the batch
-            queue: Optional queue for collecting outputs
-
-        Raises:
-            ValueError: If the request ID is already registered
-        """
-        request_id = request.request_id
-        req_state = self.request_states.get(request_id)
-        if req_state is not None:
-            self._update_streaming_request_state(req_state, request, prompt)
-            return
-
-        req_state = OmniRequestState.from_new_request(
-            tokenizer=self.tokenizer,
-            request=request,
-            prompt=prompt,
-            parent_req=parent_req,
-            request_index=request_index,
-            queue=queue,
-            log_stats=self.log_stats,
-            stream_interval=self.stream_interval,
-        )
-        self.request_states[request_id] = req_state
-        if parent_req:
-            self.parent_requests[parent_req.request_id] = parent_req
-        self.external_req_ids[req_state.external_req_id].append(request_id)
-
-    def remove_request(self, request_id: str) -> None:
-        """Rollback one previously registered request if it was never submitted."""
-        req_state = self.request_states.pop(request_id, None)
-        if req_state is None:
-            return
-
-        external_req_id = getattr(req_state, "external_req_id", None)
-        if external_req_id is not None:
-            self._native_text_metrics_by_request.pop(external_req_id, None)
-            request_ids = self.external_req_ids.get(external_req_id)
-            if request_ids is not None:
-                self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id]
-                if not self.external_req_ids[external_req_id]:
-                    self.external_req_ids.pop(external_req_id, None)
-
-        parent_req = getattr(req_state, "parent_req", None)
-        if parent_req is not None:
-            self.parent_requests.pop(parent_req.request_id, None)
-
-    def process_outputs(
-        self,
-        engine_core_outputs: list[EngineCoreOutput],
-        engine_core_timestamp: float | None = None,
-        iteration_stats: IterationStats | None = None,
-    ) -> OutputProcessorOutput:
-        default_mm_type = _modality_to_type_string(self.output_modality)
-
-        # Separate outputs that upstream can handle (has detokenizer or
-        # pooling) from multimodal-only outputs (no detokenizer, no pooling)
-        # that would trigger upstream's `assert detokenizer is not None`.
-        upstream_outputs: list[EngineCoreOutput] = []
-        mm_only_outputs: list[EngineCoreOutput] = []
-
-        for eco in engine_core_outputs:
-            req_state = self.request_states.get(eco.request_id)
-            if req_state is None:
-                continue
-
-            # Accumulate multimodal tensors regardless of path.
-            if isinstance(req_state, OmniRequestState):
-                mm_output = getattr(eco, "multimodal_output", None)
-                if mm_output is not None:
-                    mm_type = getattr(eco, "output_type", None) or default_mm_type
-                    req_state.add_multimodal_tensor(mm_output, mm_type)
-
-            # Route: if no detokenizer and no pooling output, handle locally
-            # to avoid upstream's assert on detokenizer.
-            if req_state.detokenizer is None and eco.pooling_output is None:
-                mm_only_outputs.append(eco)
-            else:
-                upstream_outputs.append(eco)
-
-        # Handle multimodal-only outputs (generation stages) locally.
-        self._process_mm_only_outputs(mm_only_outputs)
-
-        # Delegate text/pooling outputs to upstream.
-        return super().process_outputs(
-            upstream_outputs,
-            engine_core_timestamp=engine_core_timestamp,
-            iteration_stats=iteration_stats,
-        )
-
-    def _process_mm_only_outputs(
-        self,
-        engine_core_outputs: list[EngineCoreOutput],
-    ) -> None:
-        """Handle outputs from generation stages that have no detokenizer.
-
-        These cannot go through upstream process_outputs because it asserts
-        detokenizer is not None when pooling_output is None.
-        """
-        for eco in engine_core_outputs:
-            req_state = self.request_states.get(eco.request_id)
-            if req_state is None or not isinstance(req_state, OmniRequestState):
-                continue
-
-            new_token_ids = eco.new_token_ids
-            finish_reason = eco.finish_reason
-            stop_reason = eco.stop_reason
-            kv_transfer_params = eco.kv_transfer_params
-            routed_experts = eco.routed_experts
-            req_state.num_cached_tokens = eco.num_cached_tokens
-            req_state.is_prefilling = False
-
-            if request_output := req_state.make_request_output(
-                new_token_ids,
-                None,  # pooling_output
-                finish_reason,
-                stop_reason,
-                kv_transfer_params,
-                routed_experts,
-            ):
-                if req_state.queue is not None:
-                    req_state.queue.put(request_output)
-
-            if finish_reason is not None:
-                self._finish_request(req_state)
-
-    def _update_stats_from_output(
-        self,
-        req_state: RequestState,
-        engine_core_output: EngineCoreOutput,
-        engine_core_timestamp: float | None,
-        iteration_stats: IterationStats | None,
-    ):
-        was_prefilling = req_state.is_prefilling
-        native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
-        previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0
-
-        # NOTE: We pass ``None`` for  *iteration_stats* to the parent so that
-        # the upstream's ``_update_stats_from_output`` logs stats via its own
-        # ``RequestStateStats`` (req_state.stats) without touching the omni
-        # ``IterationStats`` object.  Below we re-apply the same update using
-        # *native_stats* (the omni-specific ``RequestStateStats`` attached to
-        # ``OmniRequestState``), which correctly feeds the omni metrics path.
-        # If the upstream adds new logic in ``_update_stats_from_output`` (e.g.
-        # additional stat fields), this override must be kept in sync.
-        super()._update_stats_from_output(
-            req_state,
-            engine_core_output,
-            engine_core_timestamp,
-            None,
-        )
-
-        if iteration_stats is None or engine_core_timestamp is None or native_stats is None:
-            return
-
-        iteration_stats.update_from_output(
-            engine_core_output,
-            engine_core_timestamp,
-            was_prefilling,
-            native_stats,
-            self.lora_states,
-            req_state.lora_name,
-        )
-        record = self._native_text_metric_record(req_state.external_req_id)
-        if was_prefilling:
-            record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0)
-            return
-
-        if previous_last_token_ts > 0:
-            itl_ms = max((engine_core_timestamp - previous_last_token_ts) * 1000.0, 0.0)
-            itls_ms = record.setdefault("vllm_itls_ms", [])
-            itls_ms.append(itl_ms)
-            record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms))
-
-    def _update_stats_from_finished(
-        self,
-        req_state: RequestState,
-        finish_reason: FinishReason | None,
-        iteration_stats: IterationStats | None,
-    ):
-        super()._update_stats_from_finished(req_state, finish_reason, None)
-
-        native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
-        if iteration_stats is None or native_stats is None or finish_reason is None:
-            return
-
-        iteration_stats.update_from_finished_request(
-            finish_reason=finish_reason,
-            request_id=req_state.external_req_id,
-            num_prompt_tokens=req_state.prompt_len,
-            max_tokens_param=req_state.max_tokens_param,
-            req_stats=native_stats,
-            num_cached_tokens=req_state.num_cached_tokens,
-        )
-        if not iteration_stats.finished_requests:
-            return
-
-        finished_request = iteration_stats.finished_requests[-1]
-        if finished_request.request_id != req_state.external_req_id:
-            return
-        self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = (
-            float(finished_request.mean_time_per_output_token) * 1000.0
-        )
diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py
index 276775f66..6db090955 100644
--- a/vllm_omni/engine/stage_client.py
+++ b/vllm_omni/engine/stage_client.py
@@ -11,8 +11,8 @@ if TYPE_CHECKING:
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniTokensPrompt
     from vllm_omni.outputs import OmniRequestOutput
 
-from vllm_omni.engine.output_modality import FinalOutputModalityType
 from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.outputs.output_metadata import FinalOutputModalityType
 
 
 class StageClient(Protocol):
diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py
index 78c209b0b..5fbbbc46d 100644
--- a/vllm_omni/engine/stage_engine_core_client.py
+++ b/vllm_omni/engine/stage_engine_core_client.py
@@ -230,7 +230,7 @@ class StageEngineCoreClientBase(StageClientBase):
 
     async def add_request_async(self, request: EngineCoreRequest) -> None:
         """Add request to the stage engine core."""
-        logger.info(
+        logger.debug(
             "[%s] stage-%s [rep-%s] add request: %s",
             self.__class__.__name__,
             self.stage_id,
diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py
index 8f6c199ec..7a55bb93f 100644
--- a/vllm_omni/engine/stage_engine_startup.py
+++ b/vllm_omni/engine/stage_engine_startup.py
@@ -138,6 +138,21 @@ class StageReplicaResources:
 # ---------------------------------------------------------------------------
 
 
+def _port_from_zmq_address(address: str | None) -> int | None:
+    """Extract the TCP port from a ``tcp://host:port`` ZMQ address, else ``None``.
+
+    Non-TCP transports (``ipc://``, ``inproc://``) and unparsable values return
+    ``None`` so callers can skip seeding them into the port-dedup set.
+    """
+    if not address:
+        return None
+    tail = address.rsplit(":", 1)[-1]
+    try:
+        return int(tail)
+    except ValueError:
+        return None
+
+
 class OmniMasterServer:
     """Registration server for single-stage engine startup."""
 
@@ -158,10 +173,25 @@ class OmniMasterServer:
         self._stage_configs: dict[StageRoute, Any] = {}
         self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {}
         self._stage_config_events: dict[StageRoute, threading.Event] = {}
+        # Ports already handed out by *this* server. ``get_open_ports_list`` only
+        # guarantees uniqueness *within* a single call; a per-route call cannot
+        # see ports drawn for other routes, so two stages/replicas can draw the
+        # same ephemeral port and the second engine to ``bind()`` it dies with
+        # ``zmq.error.ZMQError: Address already in use``. Multi-stage models
+        # (e.g. Qwen3-Omni: thinker/talker/code2wav) allocate many routes and hit
+        # this intermittently. We dedup every allocated port against this set,
+        # seeded with the registration port (and coordinator ROUTER port) that
+        # are already bound on the same host.
+        self._allocated_ports: set[int] = set()
+        if master_port:
+            self._allocated_ports.add(int(master_port))
         # Coordinator ROUTER address echoed back in every registration reply
         # so OmniCoordClientForStage knows where to connect from inside the
         # engine subprocess.
         self._coordinator_router_address = coordinator_router_address
+        coord_port = _port_from_zmq_address(coordinator_router_address)
+        if coord_port is not None:
+            self._allocated_ports.add(coord_port)
         # Fires only for *newly assigned* (auto-assigned) replicas, not for
         # head-side pre-allocated slots that already have head-side clients.
         self._on_register = on_register
@@ -235,7 +265,7 @@ class OmniMasterServer:
 
         self._stage_config_events[route] = threading.Event()
         self._stage_coordinator_addresses[route] = StageCoordinatorAddresses()
-        hs_port, inp_port, out_port = get_open_ports_list(count=3)
+        hs_port, inp_port, out_port = self._alloc_unique_ports(3)
         alloc = StageAllocation(
             handshake_bind_address=f"tcp://{self._address}:{hs_port}",
             handshake_connect_address=f"tcp://{self._address}:{hs_port}",
@@ -247,6 +277,32 @@ class OmniMasterServer:
         self._stage_routes[route] = alloc
         return alloc
 
+    def _alloc_unique_ports(self, count: int) -> list[int]:
+        """Return ``count`` open ports unique across every route this server owns.
+
+        ``get_open_ports_list`` dedups only within its own call, so we redraw any
+        port that collides with one already recorded in ``self._allocated_ports``
+        (registration/coordinator ports plus every previously-allocated route
+        port) and register the winners before returning. Callers already hold the
+        relevant allocation context (``__init__`` is single-threaded;
+        registration holds ``self._alloc_lock``).
+        """
+        picked: list[int] = []
+        # Bounded retry budget so a pathological host (near port exhaustion)
+        # fails loudly instead of spinning forever.
+        for _ in range(64):
+            for port in get_open_ports_list(count=count - len(picked)):
+                if port in self._allocated_ports:
+                    continue
+                self._allocated_ports.add(port)
+                picked.append(port)
+                if len(picked) == count:
+                    return picked
+        raise RuntimeError(
+            f"[OmniMasterServer] Could not allocate {count} unique open ports "
+            f"(host={self._address}); {len(self._allocated_ports)} already in use."
+        )
+
     def _next_free_replica_id(self, stage_id: int) -> int:
         """Return the next replica id to assign for an auto-assign registration.
 
diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py
index c3cf11b72..d43160895 100644
--- a/vllm_omni/engine/stage_init_utils.py
+++ b/vllm_omni/engine/stage_init_utils.py
@@ -27,11 +27,11 @@ from vllm.v1.executor import Executor
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
-from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
 from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
 from vllm_omni.inputs.preprocess import OmniInputPreprocessor
+from vllm_omni.outputs.output_processor import MultimodalOutputProcessor
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization.inc_config import OmniINCConfig
 
diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py
index 48afad0e5..8897043ca 100644
--- a/vllm_omni/engine/stage_runtime.py
+++ b/vllm_omni/engine/stage_runtime.py
@@ -32,7 +32,6 @@ from vllm_omni.engine.messages import (
     EngineQueueMessage,
     RegisterRemoteReplicaMessage,
 )
-from vllm_omni.engine.output_modality import FinalOutputModalityType
 from vllm_omni.engine.stage_client import StageClient, StagePoolClient
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase
 from vllm_omni.engine.stage_engine_startup import (
@@ -63,6 +62,7 @@ from vllm_omni.engine.stage_init_utils import (
 from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.entrypoints.stage_utils import resolve_stage_physical_devices
 from vllm_omni.entrypoints.utils import inject_omni_kv_config
+from vllm_omni.outputs.output_metadata import FinalOutputModalityType
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py
index ab2d8836b..74911af9e 100644
--- a/vllm_omni/entrypoints/async_omni.py
+++ b/vllm_omni/entrypoints/async_omni.py
@@ -26,7 +26,7 @@ from vllm.tasks import SupportedTask
 from vllm.utils import random_uuid
 from vllm.v1.engine.exceptions import EngineDeadError
 
-from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
+from vllm_omni.diffusion.data import CuMemTag, OmniACK, OmniSleepTask, OmniWakeTask
 from vllm_omni.engine.messages import ErrorMessage, OutputMessage
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni_base import (
@@ -141,12 +141,14 @@ class AsyncOmni(EngineClient, OmniBase):
         OmniBase.__init__(self, model=model, **kwargs)
         self._pause_cond: asyncio.Condition = asyncio.Condition()
         self._paused: bool = False
-        self._is_sleeping: bool = False
+        self._sleeping_tags: set[str] = set()
+        self._level2_sleeping: bool = False
         self.final_output_task: asyncio.Task | None = None
         self.event_resolver = AsyncEventResolver(orchestrator=self)
         self.config_path = self.engine.config_path
         self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None)
         self.input_processor = self.engine.input_processor
+        self.endpoint_restrictions = self.engine.endpoint_restrictions
 
         stage_index = self._get_comprehension_stage_index()
         if stage_index is None:
@@ -208,10 +210,13 @@ class AsyncOmni(EngineClient, OmniBase):
 
     def get_diffusion_od_config(self) -> Any | None:
         """Return the diffusion-stage config when the pipeline has one."""
+        saw_diffusion_stage = False
         for stage_client in self.engine.stage_clients:
             if getattr(stage_client, "stage_type", None) != "diffusion":
                 continue
 
+            saw_diffusion_stage = True
+
             od_config = getattr(stage_client, "od_config", None)
             if od_config is not None:
                 return od_config
@@ -221,6 +226,11 @@ class AsyncOmni(EngineClient, OmniBase):
             if od_config is not None:
                 return od_config
 
+        # Out-of-process diffusion clients don't carry od_config (it lives in the
+        # worker); fall back to the engine's model_class_name resolution.
+        if saw_diffusion_stage:
+            return self.engine.get_diffusion_od_config()
+
         return None
 
     @property
@@ -306,6 +316,14 @@ class AsyncOmni(EngineClient, OmniBase):
 
         logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}")
 
+        _sleeping_tags = getattr(self, "_sleeping_tags", None)
+        if _sleeping_tags:
+            raise RuntimeError(
+                f"Generation rejected: Engine is partially or fully asleep. "
+                f"Currently sleeping tags: {list(_sleeping_tags)}. "
+                f"Please perform a full wake_up before generating."
+            )
+
         # Reject diffusion list-prompt early with a clear API error.
         if isinstance(prompt, list) and any(
             getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", [])
@@ -922,11 +940,32 @@ class AsyncOmni(EngineClient, OmniBase):
                 if ack is not None:
                     await self.event_resolver.resolve(ack)
                     final_acks.append(ack)
-        self._is_sleeping = True
+        if not hasattr(self, "_sleeping_tags"):
+            self._sleeping_tags = set()
+        self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value])
+        if level == 2:
+            self._level2_sleeping = True
         return final_acks
 
     async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]:
         self._final_output_handler()
+
+        if getattr(self, "_level2_sleeping", False):
+            raise NotImplementedError(
+                "wake_up() after sleep(level=2) is not yet implemented: weights were "
+                "discarded from GPU and reloading from disk is not yet supported. "
+                "Use sleep(level=1) instead, which offloads weights to CPU RAM "
+                "and supports fast DMA restore."
+            )
+        _current_tags = getattr(self, "_sleeping_tags", set())
+        if tags is None:
+            requested_tags = list(_current_tags)
+        else:
+            requested_tags = [t for t in tags if t in _current_tags]
+        if not requested_tags:
+            logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.")
+            return []
+
         if stage_ids is None:
             stage_ids = list(range(len(self.engine.stage_clients)))
         total_workers = 0
@@ -940,7 +979,7 @@ class AsyncOmni(EngineClient, OmniBase):
         task_id = str(uuid.uuid4())
         self.event_resolver.watch_task(task_id, expected_count=total_workers)
         logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...")
-        task = OmniWakeTask(tags=tags, task_id=task_id)
+        task = OmniWakeTask(tags=requested_tags, task_id=task_id)
         rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids)
         final_acks = []
         for stage_res in rpc_results:
@@ -951,7 +990,14 @@ class AsyncOmni(EngineClient, OmniBase):
                     final_acks.append(ack)
         current_omni_platform.synchronize()
         await asyncio.sleep(0.1)
-        self._is_sleeping = False
+
+        for t in requested_tags:
+            if hasattr(self, "_sleeping_tags"):
+                self._sleeping_tags.discard(t)
+        # Only clear the level-2 flag once all tags are warm, in case partial
+        # wake support (e.g. tags=["kv_cache"] only) is added in the future.
+        if not getattr(self, "_sleeping_tags", None):
+            self._level2_sleeping = False
         logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.")
         return final_acks
 
@@ -961,7 +1007,7 @@ class AsyncOmni(EngineClient, OmniBase):
         TODO(AsyncOmni): query the orchestrator once all stage backends expose
         a real sleeping-state RPC. For now we track the requested state locally.
         """
-        return self._is_sleeping
+        return bool(getattr(self, "_sleeping_tags", None))
 
     async def add_lora(self, lora_request: LoRARequest) -> bool:
         """Load a new LoRA adapter into all stages.
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index b1ab30dee..c1ca73e44 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -148,7 +148,8 @@ class OmniServeCommand(CLISubcommand):
                     "The following CLI args are not supported under --omni: "
                     f"{', '.join(offenders)}. Configure parallelism through the "
                     "per-stage YAML (`--deploy-config` / `--stage-configs-path`) "
-                    "and replica count via `--omni-dp-size-local`."
+                    "and replica count via the per-stage `num_replicas` config field "
+                    "(single-runtime) or `--omni-dp-size-local` (headless / multi-runtime)."
                 )
 
         # --omni-lb-policy is validated against the LoadBalancingPolicy enum.
@@ -254,6 +255,15 @@ class OmniServeCommand(CLISubcommand):
             help="Path to a deploy config YAML (new format with stages/engine_args). "
             "Mutually exclusive with --stage-configs-path.",
         )
+        omni_config_group.add_argument(
+            "--strategy-config",
+            type=str,
+            default=None,
+            help="Path to a composable-parallel strategy.yaml. Only applies to "
+            "registry-based models (e.g. qwen2_5_omni): its derived parallel sizing "
+            "is overlaid onto the registry-merged stages before per-stage engine "
+            "args are built.",
+        )
         omni_config_group.add_argument(
             "--stage-overrides",
             type=str,
@@ -469,7 +479,7 @@ class OmniServeCommand(CLISubcommand):
             default=None,
             help=(
                 "JSON string for diffusion quantization_config. "
-                'Example: \'{"method":"gguf","gguf_model":"/path/to/model.gguf"}\'.'
+                'Example: \'{"method":"fp8","activation_scheme":"dynamic"}\'.'
             ),
         )
         omni_config_group.add_argument(
@@ -766,7 +776,10 @@ def run_headless(args: TrackingNamespace) -> None:
         load_omni_transfer_config_for_model,
         prepare_engine_environment,
     )
-    from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs
+    from vllm_omni.entrypoints.utils import (
+        load_and_resolve_stage_configs,
+        parse_stage_overrides,
+    )
 
     model = args.model
     stage_id: int | None = args.stage_id
@@ -802,11 +815,18 @@ def run_headless(args: TrackingNamespace) -> None:
             args.replica_id,
         )
 
-    config_path, stage_configs = load_and_resolve_stage_configs(
+    # Parse --stage-overrides (raw JSON string) exactly like the standard
+    # engine path (AsyncOmniEngine._resolve_stage_configs) so headless and
+    # standard launches resolve to the same per-stage device layout.
+    stage_overrides = parse_stage_overrides(args_dict.get("stage_overrides"))
+
+    config_path, stage_configs, _ = load_and_resolve_stage_configs(
         model,
         stage_configs_path,
         args_dict,
         deploy_config_path=args_dict.get("deploy_config"),
+        stage_overrides=stage_overrides,
+        strategy_config_path=args_dict.get("strategy_config"),
     )
 
     # Locate the stage config that matches stage_id.
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 40f3ed5a2..a9251a8d8 100755
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/vllm_omni/entrypoints/openai/api_server.py
@@ -32,7 +32,7 @@ from starlette.routing import Route
 from starlette.types import ASGIApp, Receive, Scope, Send
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
-from vllm.entrypoints.chat_utils import load_chat_template
+from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template
 from vllm.entrypoints.launcher import serve_http, terminate_if_errored
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
@@ -60,12 +60,11 @@ from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
 from vllm.entrypoints.pooling.scoring.serving import ServingScores
-from vllm.entrypoints.serve.disagg.serving import ServingTokens
+from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
 
 # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
 # Keep a fallback for older/newer upstream layouts during rebase windows.
 from vllm.entrypoints.serve.instrumentator.basic import base
-from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import ServingTokenization
 from vllm.entrypoints.serve.utils.api_utils import (
     load_aware_call,
@@ -85,12 +84,14 @@ from vllm.entrypoints.speech_to_text.translation.serving import (
     OpenAIServingTranslation,
 )
 from vllm.logger import init_logger
+from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.tasks import POOLING_TASKS
 from vllm.tool_parsers import ToolParserManager
 from vllm.utils import random_uuid
 from vllm.utils.system_utils import decorate_logs
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
+from vllm_omni.config.endpoint_policy import shutdown_unsupported_routes
 from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
@@ -259,7 +260,7 @@ async def _get_vllm_config(engine_client: EngineClient) -> Any:
     return getattr(engine_client, "vllm_config", None)
 
 
-def _remove_route_from_app(app, path: str, methods: set[str] | None = None):
+def _remove_route_from_app(app, path: str, methods: frozenset[str] | None = None):
     """Remove a route from the app by path and optionally by methods.
 
     OMNI: used to override upstream /v1/chat/completions with omni behavior.
@@ -450,7 +451,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
     # Add process-specific prefix to stdout and stderr.
     decorate_logs("APIServer", skip_if_decorated=True)
 
-    listen_address, sock = setup_openai_server(args)
+    listen_address, sock = setup_openai_server(args, reuse_port=False)
 
     # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models
     await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
@@ -509,6 +510,12 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         await omni_init_app_state(engine_client, app.state, args)
 
+        # After initializing the app state, shut down any endpoints that are model specific
+        if hasattr(engine_client, "endpoint_restrictions"):
+            shutdown_unsupported_routes(app, engine_client.endpoint_restrictions)
+        else:
+            logger.warning("engine client has no endpoint restrictions attribute")
+
         # Start background processes
         await STORAGE_MANAGER.start()
 
@@ -751,7 +758,7 @@ async def omni_init_app_state(
         state.diffusion_engine = engine_client
         state.openai_serving_models = _DiffusionServingModels(base_model_paths)
         # OMNI: tokenization endpoints are not supported in pure diffusion mode.
-        state.openai_serving_tokenization = None
+        state.serving_tokenization = None
 
         # Use for_diffusion method to create chat handler
         state.openai_serving_chat = OmniOpenAIServingChat.for_diffusion(
@@ -823,6 +830,12 @@ async def omni_init_app_state(
         if tokenizer is None or getattr(tokenizer, "chat_template", None) is None:
             resolved_chat_template = _load_model_chat_template_json(args.model)
 
+    chat_template_config = ChatTemplateConfig(
+        chat_template=resolved_chat_template,
+        chat_template_content_format=args.chat_template_content_format,
+        trust_request_chat_template=args.trust_request_chat_template,
+    )
+
     if args.tool_server == "demo":
         tool_server: ToolServer | None = DemoToolServer()
         assert isinstance(tool_server, DemoToolServer)
@@ -888,11 +901,10 @@ async def omni_init_app_state(
     await state.openai_serving_models.init_static_loras()
 
     # NOTE: kept aligned with upstream `init_app_state`:
-    # Use OpenAIServingRender (replaces the old OnlineRenderer + OnlineDerenderer + ServingRender split).
-    state.openai_serving_render = OpenAIServingRender(
+    # Use OnlineRenderer (replaced OpenAIServingRender which was removed upstream).
+    state.online_renderer = OnlineRenderer(
         model_config=engine_client.model_config,
         renderer=engine_client.renderer,
-        model_registry=state.openai_serving_models.registry,
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
@@ -909,7 +921,7 @@ async def omni_init_app_state(
         OpenAIServingResponses(
             engine_client,
             state.openai_serving_models,
-            openai_serving_render=state.openai_serving_render,
+            state.online_renderer,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -930,7 +942,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             args.response_role,
-            openai_serving_render=state.openai_serving_render,
+            online_renderer=state.online_renderer,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -957,7 +969,7 @@ async def omni_init_app_state(
         OpenAIServingCompletion(
             engine_client,
             state.openai_serving_models,
-            openai_serving_render=state.openai_serving_render,
+            online_renderer=state.online_renderer,
             request_logger=request_logger,
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
@@ -972,9 +984,7 @@ async def omni_init_app_state(
             state.openai_serving_models,
             supported_tasks=tuple(supported_tasks),
             request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
         )
         if any(task in POOLING_TASKS for task in supported_tasks)
         else None
@@ -984,9 +994,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
         )
         if "embed" in supported_tasks
         else None
@@ -996,9 +1004,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             request_logger=request_logger,
-            chat_template=resolved_chat_template,
-            chat_template_content_format=args.chat_template_content_format,
-            trust_request_chat_template=args.trust_request_chat_template,
+            chat_template_config=chat_template_config,
         )
         if "classify" in supported_tasks
         else None
@@ -1007,16 +1013,17 @@ async def omni_init_app_state(
         ServingScores(
             engine_client,
             state.openai_serving_models,
+            supported_tasks=tuple(supported_tasks),
             request_logger=request_logger,
-            score_template=resolved_chat_template,
+            chat_template_config=chat_template_config,
             log_error_stack=args.log_error_stack,
         )
         if any(t in supported_tasks for t in ("embed", "score", "token_embed"))
         else None
     )
-    state.openai_serving_tokenization = ServingTokenization(
+    state.serving_tokenization = ServingTokenization(
         state.openai_serving_models,
-        state.openai_serving_render,
+        state.online_renderer,
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
@@ -1048,7 +1055,7 @@ async def omni_init_app_state(
             engine_client,
             state.openai_serving_models,
             args.response_role,
-            openai_serving_render=state.openai_serving_render,
+            online_renderer=state.online_renderer,
             request_logger=request_logger,
             chat_template=resolved_chat_template,
             chat_template_content_format=args.chat_template_content_format,
@@ -1067,7 +1074,7 @@ async def omni_init_app_state(
         ServingTokens(
             engine_client,
             state.openai_serving_models,
-            state.openai_serving_render,
+            state.online_renderer,
             request_logger=request_logger,
             return_tokens_as_token_ids=args.return_tokens_as_token_ids,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
@@ -1154,7 +1161,7 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
     metrics_header_format = raw_request.headers.get(ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, "")
     handler = Omnichat(raw_request)
     if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
         if base_server is None:
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND.value,
@@ -1245,7 +1252,7 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
     """
     handler = Omnispeech(raw_request)
     if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
         if base_server is None:
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND.value,
@@ -1283,7 +1290,7 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request
 async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request):
     handler = Omnispeech(raw_request)
     if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
         if base_server is None:
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND.value,
@@ -1326,7 +1333,7 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request)
 async def create_audio_generate(request: OpenAICreateAudioGenerateRequest, raw_request: Request):
     handler = OmniAudioGenerate(raw_request)
     if handler is None:
-        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        base_server = getattr(raw_request.app.state, "serving_tokenization", None)
         if base_server is None:
             raise HTTPException(
                 status_code=HTTPStatus.NOT_FOUND.value,
diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py
index 12024013d..6c7043201 100644
--- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py
+++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py
@@ -5,7 +5,7 @@ import torch
 import torchaudio
 from vllm.logger import init_logger
 
-from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
+from vllm_omni.entrypoints.openai.protocol.audio import DEFAULT_AUDIO_FORMAT, AudioResponse, CreateAudio
 
 try:
     import soundfile
@@ -49,13 +49,12 @@ class AudioMixin:
             "pcm": ("RAW", "audio/pcm", {"subtype": "PCM_16"}),
             "flac": ("FLAC", "audio/flac", {}),
             "mp3": ("MP3", "audio/mpeg", {}),
-            "aac": ("AAC", "audio/aac", {}),
             "opus": ("OGG", "audio/ogg", {"subtype": "OPUS"}),
         }
 
         if response_format not in supported_formats:
-            logger.warning(f"Unsupported response format '{response_format}', defaulting to 'wav'.")
-            response_format = "wav"
+            logger.warning(f"Unsupported response format '{response_format}', defaulting to '{DEFAULT_AUDIO_FORMAT}'.")
+            response_format = DEFAULT_AUDIO_FORMAT
 
         soundfile_format, media_type, kwargs = supported_formats[response_format]
 
diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py
index b958d786a..5bd79e936 100644
--- a/vllm_omni/entrypoints/openai/image_api_utils.py
+++ b/vllm_omni/entrypoints/openai/image_api_utils.py
@@ -14,7 +14,7 @@ import io
 import PIL.Image
 
 SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024)
-SUPPORTED_LAYERED_LAYERS_RANGE = range(3, 11)
+SUPPORTED_LAYERED_LAYERS_RANGE = range(2, 11)
 
 
 def parse_size(size_str: str) -> tuple[int, int]:
diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py
index ad7bf8fae..92b2efd14 100644
--- a/vllm_omni/entrypoints/openai/protocol/audio.py
+++ b/vllm_omni/entrypoints/openai/protocol/audio.py
@@ -6,6 +6,10 @@ from pydantic import AliasChoices, BaseModel, Field, field_validator, model_vali
 
 _MAX_EMBEDDING_DIM = 8192
 
+SUPPORTED_AUDIO_FORMATS: frozenset[str] = frozenset({"wav", "pcm", "flac", "mp3", "opus"})
+SUPPORTED_CHAT_AUDIO_FORMATS: frozenset[str] = SUPPORTED_AUDIO_FORMATS | {"pcm16"}
+DEFAULT_AUDIO_FORMAT: str = "wav"
+
 
 def _normalize_ref_audio_value(value):
     if value is None:
@@ -59,7 +63,7 @@ class OpenAICreateSpeechRequest(BaseModel):
         default=None,
         description="Instructions for voice style/emotion (maps to 'instruct' for Qwen3-TTS)",
     )
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(
         default=1.0,
         ge=0.25,
@@ -325,7 +329,7 @@ class OpenAICreateAudioGenerateRequest(BaseModel):
         description="Text prompt describing the audio to generate",
     )
     model: str | None = None
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(
         default=1.0,
         ge=0.25,
@@ -391,7 +395,7 @@ class SpeechBatchItem(BaseModel):
     input: str
     voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
     instructions: str | None = None
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] | None = None
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] | None = None
     speed: float | None = Field(default=None, ge=0.25, le=4.0)
     task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
     language: str | None = None
@@ -411,7 +415,7 @@ class BatchSpeechRequest(BaseModel):
     items: list[SpeechBatchItem] = Field(..., min_length=1)
     voice: str | None = Field(default=None, validation_alias=AliasChoices("voice", "speaker"))
     instructions: str | None = None
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
     task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
     language: str | None = None
@@ -500,7 +504,7 @@ class StreamingSpeechSessionConfig(BaseModel):
     task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None
     language: str | None = None
     instructions: str | None = None
-    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT
     speed: float | None = Field(default=1.0, ge=0.25, le=4.0)
     max_new_tokens: int | None = Field(default=None, ge=1)
     initial_codec_chunk_frames: int | None = Field(
diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py
index 78f63d387..b7247962c 100644
--- a/vllm_omni/entrypoints/openai/protocol/images.py
+++ b/vllm_omni/entrypoints/openai/protocol/images.py
@@ -60,7 +60,7 @@ class ImageGenerationRequest(BaseModel):
     user: str | None = Field(default=None, description="User identifier for tracking")
     layers: int | None = Field(
         default=None,
-        description="Number of output layers for layered image models. Supported range: 3-10.",
+        description="Number of output layers for layered image models. Supported range: 2-10.",
     )
 
     @field_validator("size")
diff --git a/vllm_omni/entrypoints/openai/serving_audio_generate.py b/vllm_omni/entrypoints/openai/serving_audio_generate.py
index 81226439c..f8fedd33f 100644
--- a/vllm_omni/entrypoints/openai/serving_audio_generate.py
+++ b/vllm_omni/entrypoints/openai/serving_audio_generate.py
@@ -3,7 +3,7 @@ import asyncio
 import torch
 from fastapi import Request
 from fastapi.responses import Response
-from vllm.entrypoints.openai.engine.serving import OpenAIServing
+from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
 from vllm.logger import init_logger
 from vllm.utils import random_uuid
 
diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py
index 2878487fe..eec6a03b9 100755
--- a/vllm_omni/entrypoints/openai/serving_chat.py
+++ b/vllm_omni/entrypoints/openai/serving_chat.py
@@ -41,6 +41,7 @@ except ImportError:
     soundfile = None
 
 
+from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionNamedToolChoiceParam,
@@ -64,7 +65,6 @@ from vllm.entrypoints.openai.engine.protocol import (
     ToolCall,
     UsageInfo,
 )
-from vllm.entrypoints.openai.engine.serving import clamp_prompt_logprobs
 from vllm.entrypoints.openai.parser.harmony_utils import (
     get_streamable_parser_for_assistant,
 )
@@ -97,7 +97,12 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
 from vllm_omni.entrypoints.openai.image_api_utils import encode_image_base64_with_compression, validate_layered_layers
 from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse
-from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    DEFAULT_AUDIO_FORMAT,
+    SUPPORTED_CHAT_AUDIO_FORMATS,
+    AudioResponse,
+    CreateAudio,
+)
 from vllm_omni.entrypoints.openai.protocol.images import (
     ImageData,
     ImageEditARDeltaChunk,
@@ -120,6 +125,7 @@ from vllm_omni.entrypoints.openai.utils import (
 from vllm_omni.errors import OmniClientError
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.outputs.output_metadata import DiffusionMetadataMapping, DiffusionMetadataValue
 from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate
 
 logger = init_logger(__name__)
@@ -150,6 +156,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     # Harmony flag (always False for vllm-omni models)
     use_harmony: bool = False
 
+    @property
+    def tool_call_id_type(self) -> str:
+        """Return the tool call ID type, delegating to model config.
+
+        Upstream vLLM removed the stored ``tool_call_id_type`` attribute
+        from ``OpenAIServingChat`` after the ParserManager refactor; the
+        field is now resolved on demand via ``get_tool_call_id_type``.
+        """
+        try:
+            from vllm.entrypoints.chat_utils import get_tool_call_id_type
+
+            return get_tool_call_id_type(self.model_config)
+        except Exception:
+            return "random"
+
     def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest) -> bool:
         """Check if streamed tokens should go through the tool-call parser.
 
@@ -215,11 +236,15 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
     def _get_diffusion_extra_output_params(
         self,
-        custom_output: dict[str, Any] | None,
-    ) -> dict[str, Any] | None:
-        """Pick model-specific extra output keys from *custom_output*."""
-        if not custom_output:
-            return None
+        output: object,
+    ) -> dict[str, DiffusionMetadataValue] | None:
+        """Pick model-specific extra output keys from diffusion metadata."""
+        metadata: DiffusionMetadataMapping = {}
+        mm_output = getattr(output, "multimodal_output", None)
+        if isinstance(mm_output, dict):
+            raw_metadata = mm_output.get("metadata")
+            if isinstance(raw_metadata, dict):
+                metadata = raw_metadata
 
         if self._diffusion_extra_output_params is None:
             params: frozenset[str] = frozenset()
@@ -233,9 +258,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         if not self._diffusion_extra_output_params:
             return None
-        out = {k: custom_output[k] for k in self._diffusion_extra_output_params if k in custom_output}
+        flat_metadata: dict[str, DiffusionMetadataValue] = {}
+        for section in metadata.values():
+            if isinstance(section, dict):
+                flat_metadata.update(section)
+        out = {k: flat_metadata[k] for k in self._diffusion_extra_output_params if k in flat_metadata}
         return out or None
 
+    @staticmethod
+    def _get_diffusion_text_output(output: object) -> str:
+        mm_output = getattr(output, "multimodal_output", None)
+        if isinstance(mm_output, dict) and mm_output.get("text") is not None:
+            return str(mm_output["text"])
+        return ""
+
     def _get_supported_speakers(self) -> set[str]:
         """Load supported speakers from model config (cached)."""
         if self._supported_speakers is not None:
@@ -368,7 +404,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 tool_dicts = [tool.model_dump() for tool in request.tools]
 
             if not self.use_harmony:
-                error_check_ret = self.openai_serving_render.validate_chat_template(
+                error_check_ret = self.online_renderer.validate_chat_template(
                     request_chat_template=request.chat_template,
                     chat_template_kwargs=request.chat_template_kwargs,
                     trust_request_chat_template=self.trust_request_chat_template,
@@ -396,7 +432,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 )
             else:
                 should_include_tools = tool_dicts is not None
-                conversation, engine_prompts = self.openai_serving_render._make_request_with_harmony(
+                conversation, engine_prompts = self.online_renderer._make_request_with_harmony(
                     request, should_include_tools
                 )
 
@@ -419,6 +455,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         output_modalities = getattr(request, "modalities", engine_output_modalities)
         request.modalities = output_modalities if output_modalities is not None else engine_output_modalities
 
+        if request.modalities and "audio" in request.modalities:
+            audio_format_check = self._resolve_audio_format(request)
+            if isinstance(audio_format_check, ErrorResponse):
+                return audio_format_check
+
         num_inference_steps = None
         extra_body: dict[str, Any] = {}
         # Omni multistage image generation: Stage-0 (AR) should receive a clean
@@ -2152,8 +2193,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 else:
                     # Diffusion pipeline text output (e.g. single-stage
                     # img2text / text2text) — no AR request_output, so build
-                    # a simple text choice from custom_output.
-                    text_body = (omni_outputs.custom_output or {}).get("text_output", "")
+                    # a simple text choice from diffusion multimodal output.
+                    text_body = self._get_diffusion_text_output(omni_outputs)
                     message = ChatMessage(role=role, content=text_body)
                     choices_data = [
                         ChatCompletionResponseChoice(
@@ -2179,7 +2220,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     response_metrics = {}
                 response_metrics.setdefault("stage_durations", omni_outputs.stage_durations or {})
                 response_metrics.setdefault("peak_memory_mb", float(omni_outputs.peak_memory_mb or 0.0))
-                extra = self._get_diffusion_extra_output_params(omni_outputs.custom_output)
+                extra = self._get_diffusion_extra_output_params(omni_outputs)
                 if extra:
                     response_metrics.update(extra)
             choices.extend(choices_data)
@@ -2530,10 +2571,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         else:
             sample_rate = int(sr_raw)
 
+        audio_format = self._resolve_audio_format(request)
+        if isinstance(audio_format, ErrorResponse):
+            return audio_format
+
         audio_obj = CreateAudio(
             audio_tensor=audio_tensor,
             sample_rate=sample_rate,
-            response_format="wav",
+            response_format=audio_format,
             speed=1.0,
             base64_encode=True,
         )
@@ -3413,7 +3458,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             # Text output path (img2text / text2text)
             if is_text_request and result.final_output_type == "text":
-                text_body = (result.custom_output or {}).get("text_output", "")
+                text_body = self._get_diffusion_text_output(result)
                 message = ChatMessage(role="assistant", content=text_body)
                 choice = ChatCompletionResponseChoice(
                     index=0,
@@ -3432,7 +3477,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         completion_tokens=len(text_body.split()),
                         total_tokens=len(prompt.split()) + len(text_body.split()),
                     ),
-                    metrics=self._get_diffusion_extra_output_params(result.custom_output),
+                    metrics=self._get_diffusion_extra_output_params(result),
                 )
                 logger.info(
                     "Diffusion chat completed for request %s: text output (%d chars)",
@@ -3484,10 +3529,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     raise ValueError(f"Unexpected audio tensor rank {audio_tensor.ndim}; expected 1-3 dims.")
                 audio_array = audio_tensor.numpy()
 
+                audio_format = self._resolve_audio_format(request)
+                if isinstance(audio_format, ErrorResponse):
+                    return audio_format
+
                 audio_obj = CreateAudio(
                     audio_tensor=audio_array,
                     sample_rate=sample_rate,
-                    response_format="wav",
+                    response_format=audio_format,
                     speed=1.0,
                     base64_encode=True,
                 )
@@ -3565,7 +3614,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     completion_tokens=1,
                     total_tokens=len(prompt.split()) + 1,
                 ),
-                metrics=self._get_diffusion_extra_output_params(result.custom_output),
+                metrics=self._get_diffusion_extra_output_params(result),
             )
 
             logger.info(
@@ -3703,6 +3752,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         return height, width
 
+    def _resolve_audio_format(self, request: ChatCompletionRequest) -> str | ErrorResponse:
+        """Extract and validate the audio output format from a chat completion request."""
+        audio_params = getattr(request, "audio", None)
+        if isinstance(audio_params, dict):
+            audio_format = audio_params.get("format", DEFAULT_AUDIO_FORMAT)
+        else:
+            audio_format = DEFAULT_AUDIO_FORMAT
+        if audio_format not in SUPPORTED_CHAT_AUDIO_FORMATS:
+            return self._create_error_response(
+                f"Invalid audio format '{audio_format}'. Supported formats: {sorted(SUPPORTED_CHAT_AUDIO_FORMATS)}",
+            )
+        if audio_format == "pcm16":
+            audio_format = "pcm"
+        return audio_format
+
     def _create_error_response(
         self,
         message: str,
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index ac37a63a2..3db6b63a6 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/vllm_omni/entrypoints/openai/serving_speech.py
@@ -21,12 +21,12 @@ import torch
 from fastapi import HTTPException, Request, UploadFile
 from fastapi.responses import Response, StreamingResponse
 from transformers.utils.hub import cached_file
+from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing
 from vllm.entrypoints.launcher import terminate_if_errored
 from vllm.entrypoints.openai.engine.protocol import (
     ErrorResponse,
     RequestResponseMetadata,
 )
-from vllm.entrypoints.openai.engine.serving import OpenAIServing
 from vllm.inputs import tokens_input
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
@@ -1395,7 +1395,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         emb_dim = len(embedding)
         if self._tts_model_type == "ming_tts":
             if emb_dim != 192:
-                logger.warning("speaker_embedding has %d dimensions; Ming dense expects 192", emb_dim)
+                raise ValueError(f"Ming speaker embedding must have 192 dims, got {emb_dim}")
         else:
             dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim)
             if dim_err is not None:
@@ -2620,16 +2620,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         self,
         request,
         *,
-        include_language=False,
         include_voice=False,
         plain_text_passthrough=False,
     ):
         instruction_text = request.instructions.strip() if isinstance(request.instructions, str) else None
         instruction_dict: dict[str, Any] = {}
 
-        if include_language and request.language not in (None, "", "Auto"):
-            instruction_dict["方言"] = request.language
-
         voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
         if include_voice and request.voice and not (voice_lower and voice_lower in self.uploaded_speakers):
             instruction_dict["IP"] = request.voice
@@ -2652,7 +2648,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         """Build a Ming instruction payload from OpenAI speech fields."""
         return self._parse_ming_instruction_fields(
             request,
-            include_language=True,
             include_voice=True,
             plain_text_passthrough=True,
         )
@@ -3500,12 +3495,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_id = request_id or f"speech-{random_uuid()}"
         qwen3_ref_audio_warmup_artifact_key: str | None = None
 
-        # If this is a streaming request, we need to coerce
-        # cumulative outputs to delta outputs; this ensures
-        # we don't emit redundant MM data & drain after emitting.
+        # If this is a streaming request with real async chunks, we need to
+        # coerce cumulative outputs to delta outputs; this ensures we don't
+        # emit redundant MM data & drain after emitting. Qwen3-TTS full-payload
+        # (async_chunk=False) has no incremental audio chunks, so keep
+        # FINAL_ONLY semantics and let the streaming response send the final
+        # waveform once. Scoped to qwen3_tts: other async_chunk=False models
+        # keep the DELTA coercion they stream with today.
         # list() makes a copy to avoid mutating the params.
         sampling_params_list = list(self.engine_client.default_sampling_params_list)
-        is_streaming_request = request.is_streaming()
+        async_chunk = getattr(self.model_config, "async_chunk", True)
+        qwen3_full_payload = self._tts_model_type == "qwen3_tts" and not bool(async_chunk)
+        is_streaming_request = request.is_streaming() and not qwen3_full_payload
         sampling_params_list = coerce_param_message_types(sampling_params_list, is_streaming_request)
 
         # Build prompt + tts_params via the per-model adapter (RFC #4327). Every
@@ -3518,11 +3519,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         # apply uploaded speakers inside validate(), which sets request.ref_audio
         # in place. The builders need to know whether the caller supplied audio
         # inline vs. via an uploaded voice.
+        model_type: str | None = None
         has_inline_ref_audio = request.ref_audio is not None
         if self._tts_model_type == "ming_flash_omni_tts":
             # ming_flash_omni is intentionally NOT migrated onto the adapter
             # framework in this PR (it has no registered adapter); keep it on the
             # legacy inline dispatch so serving still works.
+            model_type = "ming_flash_omni_tts"
             validation_error = self._validate_ming_flash_omni_tts_request(request)
             if validation_error:
                 raise ValueError(validation_error)
@@ -3536,6 +3539,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             prepared = await adapter.build(request, sampling_params_list, has_inline_ref_audio)
             prompt = prepared.prompt
             tts_params = prepared.tts_params
+            model_type = prepared.model_type
             qwen3_ref_audio_warmup_artifact_key = prepared.warmup_artifact_key
         else:
             # Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker"
@@ -3559,36 +3563,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             tts_params = {}
             prompt = {"prompt": request.input}
 
-        if self._tts_model_type == "step_audio2":
-            model_type = "step_audio2"
-        elif self._is_fish_speech:
-            model_type = "fish_speech"
-        elif self._tts_model_type == "covo_audio":
-            model_type = "covo_audio"
-        elif self._tts_model_type == "voxtral_tts":
-            model_type = "voxtral_tts"
-        elif self._tts_model_type == "cosyvoice3":
-            model_type = "cosyvoice3"
-        elif self._tts_model_type == "voxcpm2":
-            model_type = "voxcpm2"
-        elif self._tts_model_type == "ming_flash_omni_tts":
-            model_type = "ming_flash_omni_tts"
-        elif self._tts_model_type == "ming_tts":
-            model_type = "ming_tts"
-        elif self._tts_model_type == "moss_tts_nano":
-            model_type = "moss_tts_nano"
-        elif self._tts_model_type == "moss_tts":
-            model_type = "moss_tts"
-        elif self._tts_model_type == "higgs_audio_v2":
-            model_type = "higgs_audio_v2"
-        elif self._tts_model_type == "glm_tts":
-            model_type = "glm_tts"
-        elif self._tts_model_type == "indextts2":
-            model_type = "indextts2"
-        elif self._is_tts:
-            model_type = tts_params.get("task_type", ["unknown"])[0]
-        else:
-            model_type = "generic"
+        if model_type is None:
+            if self._is_tts:
+                model_type = tts_params.get("task_type", ["unknown"])[0]
+            else:
+                model_type = "generic"
         logger.info(
             "TTS speech request %s: text=%r, model=%s",
             request_id,
@@ -3705,15 +3684,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 sampling_params_list[0].max_tokens = int(request.max_new_tokens) + 1
 
         if request.seed is not None and sampling_params_list:
-            if sampling_params_list is self.engine_client.default_sampling_params_list:
+            import copy
+
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            stage0_params = sampling_params_list[0]
+            stage0_params.seed = request.seed
+            if stage0_params.extra_args is None:
+                stage0_params.extra_args = {}
+            stage0_params.extra_args["tts_local_seed"] = request.seed
+
+        if self._tts_model_type == "qwen3_tts" and sampling_params_list:
+            stage0_params = sampling_params_list[0]
+            default_seed = getattr(stage0_params, "seed", None)
+            if default_seed is not None:
                 import copy
 
                 sampling_params_list = copy.deepcopy(sampling_params_list)
-            sampling_params_list[0].seed = request.seed
-            if self._tts_model_type == "qwen3_tts":
-                if sampling_params_list[0].extra_args is None:
-                    sampling_params_list[0].extra_args = {}
-                sampling_params_list[0].extra_args["qwen3_tts_request_seed"] = request.seed
+                stage0_params = sampling_params_list[0]
+                if stage0_params.extra_args is None:
+                    stage0_params.extra_args = {}
+                stage0_params.extra_args.setdefault("tts_local_seed", int(default_seed))
 
         generator = self.engine_client.generate(
             prompt=prompt,
diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py
index 3e2927a27..ed846e7e8 100644
--- a/vllm_omni/entrypoints/openai/serving_video.py
+++ b/vllm_omni/entrypoints/openai/serving_video.py
@@ -29,6 +29,11 @@ from vllm_omni.entrypoints.openai.stage_params import (
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+from vllm_omni.outputs.output_metadata import (
+    DiffusionMetadataMapping,
+    DiffusionMultimodalOutput,
+    DiffusionPayloadValue,
+)
 
 logger = init_logger(__name__)
 
@@ -58,8 +63,8 @@ class ReferenceAudio:
 class VideoGenerationArtifacts:
     """Normalized outputs and profiler metadata extracted from one request."""
 
-    videos: list[Any]
-    audios: list[Any | None]
+    videos: list[DiffusionPayloadValue]
+    audios: list[DiffusionPayloadValue | None]
     actions: list[VideoAction | None]
     audio_sample_rate: int
     output_fps: int
@@ -225,8 +230,10 @@ class OmniOpenAIServingVideo:
         )
 
         result = await self._run_generation(prompt, gen_params, reference_id)
-        custom_output = self._extract_custom_output(result)
-        action_only = isinstance(custom_output, dict) and bool(custom_output.get("action_only_output"))
+        multimodal_output = self._extract_multimodal_output(result)
+        metadata = multimodal_output.get("metadata") if isinstance(multimodal_output, dict) else {}
+        common_metadata = metadata.get("common") if isinstance(metadata, dict) else {}
+        action_only = bool(isinstance(common_metadata, dict) and common_metadata.get("action_only_output"))
         videos = [{"action_only_output": True}] if action_only else self._extract_video_outputs(result)
         audios = self._extract_audio_outputs(result, expected_count=len(videos))
         actions = self._extract_action_outputs(result, expected_count=len(videos))
@@ -340,10 +347,12 @@ class OmniOpenAIServingVideo:
         return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb, artifacts.actions[0]
 
     @staticmethod
-    def _resolve_video_fps_multiplier(result: Any) -> int:
-        custom_output = OmniOpenAIServingVideo._extract_custom_output(result)
-        if isinstance(custom_output, dict):
-            multiplier = custom_output.get("video_fps_multiplier")
+    def _resolve_video_fps_multiplier(result: object) -> int:
+        multimodal_output = OmniOpenAIServingVideo._extract_multimodal_output(result)
+        metadata = multimodal_output.get("metadata")
+        video_metadata = metadata.get("video") if isinstance(metadata, Mapping) else None
+        if isinstance(video_metadata, Mapping):
+            multiplier = video_metadata.get("video_fps_multiplier")
             if multiplier is not None:
                 return int(multiplier)
         return 1
@@ -390,7 +399,7 @@ class OmniOpenAIServingVideo:
         prompt: OmniTextPrompt,
         gen_params: OmniDiffusionSamplingParams,
         request_id: str,
-    ) -> Any:
+    ) -> object:
         stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None)
 
         if not stage_configs:
@@ -433,7 +442,7 @@ class OmniOpenAIServingVideo:
         return result
 
     @staticmethod
-    def _normalize_video_outputs(videos: Any) -> list[Any]:
+    def _normalize_video_outputs(videos: DiffusionPayloadValue | None) -> list[DiffusionPayloadValue]:
         if videos is None:
             return []
         if hasattr(videos, "ndim") and videos.ndim == 5:
@@ -443,7 +452,7 @@ class OmniOpenAIServingVideo:
                 return []
             first = videos[0]
             if hasattr(first, "ndim") and first.ndim == 5:
-                flattened: list[Any] = []
+                flattened: list[DiffusionPayloadValue] = []
                 for item in videos:
                     if hasattr(item, "ndim") and item.ndim == 5:
                         flattened.extend([item[i] for i in range(item.shape[0])])
@@ -459,7 +468,7 @@ class OmniOpenAIServingVideo:
             return videos
         return [videos]
 
-    def _extract_video_outputs(self, result: Any) -> list[Any]:
+    def _extract_video_outputs(self, result: object) -> list[DiffusionPayloadValue]:
         videos = None
         if hasattr(result, "images") and result.images:
             videos = result.images
@@ -469,10 +478,14 @@ class OmniOpenAIServingVideo:
                 videos = request_output["images"]
             elif hasattr(request_output, "images") and request_output.images:
                 videos = request_output.images
-            elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
-                videos = request_output.multimodal_output.get("video")
-        if videos is None and hasattr(result, "multimodal_output") and result.multimodal_output:
-            videos = result.multimodal_output.get("video")
+            else:
+                request_multimodal_output = getattr(request_output, "multimodal_output", None)
+                if isinstance(request_multimodal_output, Mapping):
+                    videos = request_multimodal_output.get("video")
+        if videos is None:
+            multimodal_output = getattr(result, "multimodal_output", None)
+            if isinstance(multimodal_output, Mapping):
+                videos = multimodal_output.get("video")
 
         normalized = self._normalize_video_outputs(videos)
         if not normalized:
@@ -483,17 +496,21 @@ class OmniOpenAIServingVideo:
         return normalized
 
     @staticmethod
-    def _extract_audio_outputs(result: Any, expected_count: int) -> list[Any | None]:
+    def _extract_audio_outputs(result: object, expected_count: int) -> list[DiffusionPayloadValue | None]:
         audio = None
-        if hasattr(result, "multimodal_output") and result.multimodal_output:
-            audio = result.multimodal_output.get("audio")
+        multimodal_output = getattr(result, "multimodal_output", None)
+        if isinstance(multimodal_output, Mapping):
+            audio = multimodal_output.get("audio")
         elif hasattr(result, "request_output"):
             request_output = result.request_output
             if isinstance(request_output, dict) and request_output.get("multimodal_output"):
                 mm_output = request_output.get("multimodal_output") or {}
-                audio = mm_output.get("audio")
-            elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output:
-                audio = request_output.multimodal_output.get("audio")
+                if isinstance(mm_output, Mapping):
+                    audio = mm_output.get("audio")
+            else:
+                request_multimodal_output = getattr(request_output, "multimodal_output", None)
+                if isinstance(request_multimodal_output, Mapping):
+                    audio = request_multimodal_output.get("audio")
 
         if audio is None:
             return [None] * expected_count
@@ -514,7 +531,7 @@ class OmniOpenAIServingVideo:
 
         return [audio] + [None] * max(expected_count - 1, 0)
 
-    def _resolve_audio_sample_rate(self, result: Any) -> int:
+    def _resolve_audio_sample_rate(self, result: object) -> int:
         result_sample_rate = self._extract_audio_sample_rate_from_result(result)
         if result_sample_rate is not None:
             return result_sample_rate
@@ -528,38 +545,45 @@ class OmniOpenAIServingVideo:
         return 24000
 
     @classmethod
-    def _extract_action_outputs(cls, result: Any, expected_count: int) -> list[VideoAction | None]:
-        custom_output = cls._extract_custom_output(result)
-        if not custom_output or "action" not in custom_output:
-            return [None] * expected_count
+    def _extract_action_outputs(cls, result: object, expected_count: int) -> list[VideoAction | None]:
+        multimodal_output = cls._extract_multimodal_output(result)
+        if "actions" in multimodal_output:
+            action_payload = multimodal_output["actions"]
+            metadata = multimodal_output.get("metadata")
+            action_metadata = metadata.get("actions") if isinstance(metadata, Mapping) else {}
+            action_metadata = action_metadata if isinstance(action_metadata, Mapping) else {}
+            action_items = cls._split_action_payload(action_payload, expected_count)
+            return [
+                cls._make_video_action(action_item, action_metadata) if action_item is not None else None
+                for action_item in action_items
+            ]
 
-        action_payload = custom_output.get("actions", custom_output["action"])
-        action_items = cls._split_action_payload(action_payload, expected_count)
-        return [
-            cls._make_video_action(action_item, custom_output) if action_item is not None else None
-            for action_item in action_items
-        ]
+        return [None] * expected_count
 
     @staticmethod
-    def _extract_custom_output(result: Any) -> dict[str, Any]:
-        custom_output = getattr(result, "custom_output", None)
-        if isinstance(custom_output, dict):
-            return custom_output
+    def _extract_multimodal_output(result: object) -> DiffusionMultimodalOutput:
+        multimodal_output = getattr(result, "multimodal_output", None)
+        if isinstance(multimodal_output, Mapping):
+            return dict(multimodal_output)
 
         request_output = getattr(result, "request_output", None)
         if isinstance(request_output, dict):
-            custom_output = request_output.get("custom_output")
-            if custom_output is None:
-                custom_output = request_output.get("_custom_output")
+            multimodal_output = request_output.get("multimodal_output")
+            if multimodal_output is None:
+                multimodal_output = request_output.get("_multimodal_output")
         elif request_output is not None:
-            custom_output = getattr(request_output, "custom_output", None)
-            if custom_output is None:
-                custom_output = getattr(request_output, "_custom_output", None)
+            multimodal_output = getattr(request_output, "multimodal_output", None)
+            if multimodal_output is None:
+                multimodal_output = getattr(request_output, "_multimodal_output", None)
 
-        return custom_output if isinstance(custom_output, dict) else {}
+        return dict(multimodal_output) if isinstance(multimodal_output, Mapping) else {}
 
     @classmethod
-    def _split_action_payload(cls, action: Any, expected_count: int) -> list[Any | None]:
+    def _split_action_payload(
+        cls,
+        action: DiffusionPayloadValue,
+        expected_count: int,
+    ) -> list[DiffusionPayloadValue | None]:
         if expected_count <= 0:
             return []
 
@@ -573,30 +597,34 @@ class OmniOpenAIServingVideo:
         return [action] + [None] * (expected_count - 1)
 
     @classmethod
-    def _make_video_action(cls, action: Any, custom_output: dict[str, Any]) -> VideoAction:
+    def _make_video_action(
+        cls,
+        action: DiffusionPayloadValue,
+        action_metadata: DiffusionMetadataMapping,
+    ) -> VideoAction:
         data = cls._to_jsonable(action)
         if not isinstance(data, list):
             data = [data]
 
-        action_mode = custom_output.get("action_mode")
+        action_mode = action_metadata.get("action_mode")
         return VideoAction(
             data=data,
             shape=cls._shape_of(action),
             dtype=cls._dtype_of(action),
-            raw_action_dim=cls._coerce_optional_int(custom_output.get("raw_action_dim")),
+            raw_action_dim=cls._coerce_optional_int(action_metadata.get("raw_action_dim")),
             action_mode=str(action_mode) if action_mode is not None else None,
-            domain_id=cls._coerce_optional_int(custom_output.get("domain_id")),
+            domain_id=cls._coerce_optional_int(action_metadata.get("domain_id")),
         )
 
     @staticmethod
-    def _index_action(action: Any, index: int) -> Any:
+    def _index_action(action: DiffusionPayloadValue, index: int) -> DiffusionPayloadValue | None:
         try:
             return action[index]
         except (IndexError, KeyError, TypeError):
             return None
 
     @classmethod
-    def _to_jsonable(cls, value: Any) -> Any:
+    def _to_jsonable(cls, value: DiffusionPayloadValue) -> DiffusionPayloadValue:
         if hasattr(value, "detach"):
             value = value.detach()
         if hasattr(value, "cpu"):
@@ -613,7 +641,7 @@ class OmniOpenAIServingVideo:
         return value
 
     @classmethod
-    def _shape_of(cls, value: Any) -> list[int]:
+    def _shape_of(cls, value: DiffusionPayloadValue) -> list[int]:
         shape = getattr(value, "shape", None)
         if shape is not None:
             try:
@@ -627,12 +655,12 @@ class OmniOpenAIServingVideo:
         return []
 
     @staticmethod
-    def _dtype_of(value: Any) -> str | None:
+    def _dtype_of(value: DiffusionPayloadValue) -> str | None:
         dtype = getattr(value, "dtype", None)
         return str(dtype) if dtype is not None else None
 
     @staticmethod
-    def _coerce_optional_int(value: Any) -> int | None:
+    def _coerce_optional_int(value: object) -> int | None:
         if value is None:
             return None
         try:
@@ -642,7 +670,7 @@ class OmniOpenAIServingVideo:
             return None
 
     @staticmethod
-    def _resolve_fps(result: Any) -> int | None:
+    def _resolve_fps(result: object) -> int | None:
         """Extract fps from multimodal_output if the model reported it."""
         multimodal_output = getattr(result, "multimodal_output", None)
         if isinstance(multimodal_output, Mapping):
@@ -682,7 +710,7 @@ class OmniOpenAIServingVideo:
         return None
 
     @classmethod
-    def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None:
+    def _extract_audio_sample_rate_from_result(cls, result: object) -> int | None:
         multimodal_output = getattr(result, "multimodal_output", None)
         if isinstance(multimodal_output, Mapping):
             sample_rate = cls._coerce_audio_sample_rate(
@@ -717,7 +745,7 @@ class OmniOpenAIServingVideo:
         return None
 
     @classmethod
-    def _extract_audio_sample_rate_from_config(cls, config: Any) -> int | None:
+    def _extract_audio_sample_rate_from_config(cls, config: object) -> int | None:
         if config is None:
             return None
 
@@ -748,7 +776,7 @@ class OmniOpenAIServingVideo:
         return None
 
     @staticmethod
-    def _coerce_audio_sample_rate(value: Any) -> int | None:
+    def _coerce_audio_sample_rate(value: object) -> int | None:
         if value is None:
             return None
 
@@ -761,12 +789,12 @@ class OmniOpenAIServingVideo:
         return sample_rate if sample_rate > 0 else None
 
     @staticmethod
-    def _extract_stage_durations(result: Any) -> dict[str, float]:
+    def _extract_stage_durations(result: object) -> dict[str, float]:
         stage_durations = getattr(result, "stage_durations", None)
         return stage_durations if isinstance(stage_durations, dict) else {}
 
     @staticmethod
-    def _extract_peak_memory_mb(result: Any) -> float:
+    def _extract_peak_memory_mb(result: object) -> float:
         peak_memory_mb = getattr(result, "peak_memory_mb", 0.0)
         try:
             return float(peak_memory_mb or 0.0)
diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py
index 37318a521..7e1a8e66b 100644
--- a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py
+++ b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py
@@ -25,9 +25,18 @@ class MingTTSAdapter(ARTTSAdapter):
         ref_audio_source = request.ref_audio
         voice_lower = request.voice.lower() if isinstance(request.voice, str) else None
         if ref_audio_source is None and voice_lower in server.uploaded_speakers:
-            ref_audio_source = server._get_uploaded_audio_data(request.voice)
-            if request.ref_text is None:
-                request.ref_text = server.uploaded_speakers[voice_lower].get("ref_text")
+            speaker_info = server.uploaded_speakers[voice_lower]
+            if speaker_info.get("embedding_source") == "direct":
+                if request.speaker_embedding is None:
+                    request.speaker_embedding = server._get_uploaded_speaker_embedding(request.voice)
+                if request.speaker_embedding is None:
+                    raise ValueError(f"Speaker embedding for uploaded voice '{request.voice}' is missing")
+            else:
+                ref_audio_source = server._get_uploaded_audio_data(request.voice)
+                if not ref_audio_source:
+                    raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing")
+                if request.ref_text is None:
+                    request.ref_text = speaker_info.get("ref_text")
         ref_audio_data = None
         if isinstance(ref_audio_source, list):
             ref_audio_data = await server._resolve_ref_audio_many(ref_audio_source)
diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py
index 961638e14..a2c598b31 100644
--- a/vllm_omni/entrypoints/utils.py
+++ b/vllm_omni/entrypoints/utils.py
@@ -1,3 +1,4 @@
+import json
 import os
 import types
 from collections import Counter
@@ -316,7 +317,8 @@ def load_stage_configs_from_model(
     base_engine_args: dict | None = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
-) -> list:
+    strategy_config_path: str | None = None,
+) -> tuple[list, str | None]:
     """Load stage configurations from model's default config file.
 
     For models registered in the pipeline registry (new path), uses
@@ -330,9 +332,15 @@ def load_stage_configs_from_model(
         base_engine_args: Base engine args to merge as CLI overrides.
         deploy_config_path: Optional explicit deploy config path.
         stage_overrides: Per-stage overrides from --stage-overrides.
+        strategy_config_path: Optional path to a composable-parallel
+            ``strategy.yaml`` whose derived sizing is overlaid onto the
+            registry-merged stages (opt-in; ignored on the legacy YAML path).
 
     Returns:
-        List of stage configuration dictionaries
+        ``(stage_configs, omni_lb_policy)``: the list of stage configuration
+        dictionaries plus the strategy-derived pipeline-wide ``omni_lb_policy``
+        (``None`` when no strategy set one). The policy is returned rather than
+        written into a caller-provided mutable dict.
     """
     if base_engine_args is None:
         base_engine_args = {}
@@ -343,25 +351,44 @@ def load_stage_configs_from_model(
             for key, val in overrides.items():
                 cli_overrides[f"stage_{stage_id_str}_{key}"] = val
 
-    stages = StageConfigFactory.create_from_model(
+    strategy_specs = None
+    if strategy_config_path is not None:
+        from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs
+
+        strategy_specs = load_strategy_specs(strategy_config_path)
+
+    stages, omni_lb_policy = StageConfigFactory.create_from_model(
         model,
+        trust_remote_code=cli_overrides.get("trust_remote_code", False),
         cli_overrides=cli_overrides,
         deploy_config_path=deploy_config_path,
+        strategy_specs=strategy_specs,
     )
     if stages is not None:
         # Convert StageConfig objects to OmegaConf for backward compat
-        return [stage.to_omegaconf() for stage in stages]
+        return [stage.to_omegaconf() for stage in stages], omni_lb_policy
 
-    # Legacy fallback: load from YAML
+    # Legacy fallback: load from YAML. A composable-parallel strategy cannot be
+    # applied here (it overlays onto registry-merged stages), so warn rather than
+    # silently dropping the operator's --strategy-config.
+    if strategy_config_path is not None:
+        logger.warning(
+            "--strategy-config (%s) was provided but model %r 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.",
+            strategy_config_path,
+            model,
+        )
     stage_config_path = resolve_model_config_path(model)
     if stage_config_path is None:
-        return []
+        return [], None
     stage_configs = load_stage_configs_from_yaml(
         config_path=stage_config_path,
         base_engine_args=base_engine_args,
         prefer_stage_engine_args=True,
     )
-    return stage_configs
+    return stage_configs, None
 
 
 def load_stage_configs_from_yaml(
@@ -500,6 +527,25 @@ def filter_stages(
         return stage_configs
 
 
+def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None:
+    """Parse the ``--stage-overrides`` value into a per-stage override dict.
+
+    ``value`` may be a raw JSON string (as supplied on the CLI) or an
+    already-parsed mapping. Returns ``None`` when no overrides are given.
+
+    Raises:
+        ValueError: when ``value`` is a string that is not valid JSON.
+    """
+    if not value:
+        return None
+    if isinstance(value, str):
+        try:
+            return json.loads(value)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"--stage-overrides is not valid JSON: {exc}. Got: {value!r}") from exc
+    return value
+
+
 def load_and_resolve_stage_configs(
     model: str,
     stage_configs_path: str | None,
@@ -507,7 +553,8 @@ def load_and_resolve_stage_configs(
     default_stage_cfg_factory: Any = None,
     deploy_config_path: str | None = None,
     stage_overrides: dict[str, dict[str, Any]] | None = None,
-) -> tuple[str, list]:
+    strategy_config_path: str | None = None,
+) -> tuple[str, list, str | None]:
     """Load stage configurations from model or YAML file with fallback to defaults.
 
     Args:
@@ -520,9 +567,13 @@ def load_and_resolve_stage_configs(
             Mutually exclusive with ``stage_configs_path``.
         stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON.
             Keys are stage_id strings, values are dicts of overrides.
+        strategy_config_path: Optional path to a composable-parallel
+            ``strategy.yaml`` overlaid onto the registry-merged stages.
 
     Returns:
-        Tuple of (config_path, stage_configs)
+        Tuple of ``(config_path, stage_configs, omni_lb_policy)`` — the last is
+        the strategy-derived pipeline-wide load-balancer policy (``None`` when no
+        strategy set one), returned for the engine to apply.
     """
     if stage_configs_path is not None and deploy_config_path is not None:
         raise ValueError(
@@ -550,13 +601,15 @@ def load_and_resolve_stage_configs(
                 stage_configs_path,
             )
 
+    omni_lb_policy: str | None = None
     if deploy_config_path is not None:
         config_path = deploy_config_path
-        stage_configs = load_stage_configs_from_model(
+        stage_configs, omni_lb_policy = load_stage_configs_from_model(
             model,
             base_engine_args=kwargs,
             deploy_config_path=deploy_config_path,
             stage_overrides=stage_overrides,
+            strategy_config_path=strategy_config_path,
         )
         if not stage_configs:
             if default_stage_cfg_factory is not None:
@@ -566,10 +619,11 @@ def load_and_resolve_stage_configs(
                 stage_configs = []
     elif stage_configs_path is None:
         config_path = resolve_model_config_path(model)
-        stage_configs = load_stage_configs_from_model(
+        stage_configs, omni_lb_policy = load_stage_configs_from_model(
             model,
             base_engine_args=kwargs,
             stage_overrides=stage_overrides,
+            strategy_config_path=strategy_config_path,
         )
         if not stage_configs:
             if default_stage_cfg_factory is not None:
@@ -584,7 +638,7 @@ def load_and_resolve_stage_configs(
     stage_configs = filter_stages(config_path, stage_configs, kwargs)
     logger.debug(f"stage_configs: {stage_configs}")
 
-    return config_path, stage_configs
+    return config_path, stage_configs, omni_lb_policy
 
 
 def get_final_stage_id_for_e2e(
diff --git a/vllm_omni/experimental/__init__.py b/vllm_omni/experimental/__init__.py
index f8b58d229..c7d59e2e2 100644
--- a/vllm_omni/experimental/__init__.py
+++ b/vllm_omni/experimental/__init__.py
@@ -1,3 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Experimental vLLM-Omni features (APIs may change without notice)."""
+"""Experimental vLLM-Omni subsystems.
+
+Modules under ``vllm_omni.experimental`` are in active development and not yet
+covered by the stability guarantees of the rest of the package — their APIs,
+config surfaces, and defaults may change without notice. The AR-Diffusion
+Engine (:mod:`vllm_omni.experimental.ar_diffusion`) lives here.
+"""
diff --git a/vllm_omni/experimental/ar_diffusion/__init__.py b/vllm_omni/experimental/ar_diffusion/__init__.py
new file mode 100644
index 000000000..b312660ec
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/__init__.py
@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AR-Diffusion Engine (AR-Diffusion).
+
+The AR-Diffusion engine: a ``DiffusionEngine`` subclass that adds engine-level KV
+cache management for autoregressive / chunked "world-model" diffusion models
+(DreamZero and the AR-DiT family). Selected via
+``OmniDiffusionConfig.engine_backend = "ar_diffusion"``.
+"""
+
+from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
+
+__all__ = ["ARDiffusionEngine"]
diff --git a/vllm_omni/experimental/ar_diffusion/engine.py b/vllm_omni/experimental/ar_diffusion/engine.py
new file mode 100644
index 000000000..6fb6ee7c9
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/engine.py
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The AR-Diffusion Engine (AR-Diffusion)."""
+
+from __future__ import annotations
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+logger = init_logger(__name__)
+
+#: Import path of the runner the AR-Diffusion engine routes its workers to.
+AR_DIFFUSION_MODEL_RUNNER_CLS = "vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner"
+
+
+class ARDiffusionEngine(DiffusionEngine):
+    """AR-Diffusion engine with engine-level KV cache management.
+
+    AR-Diffusion serves autoregressive / chunked blockwise-causal diffusion models
+    (world models, AR-DiT) that materialize persistent attention KV. It reuses
+    vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` / ``BlockTables``)
+    as a library, driven from the engine rather than hand-rolled inside each
+    model. See ``BDE_doc/diffusion_kv_cache_management.md`` for the design and
+    ``BDE_doc/dreamzero_kv_phase1_plan.md`` for the rollout.
+
+    It is selected per model via ``OmniDiffusionConfig.engine_backend = "ar_diffusion"``
+    (resolved by :meth:`DiffusionEngine.make_engine`), so models that do not opt
+    in keep using the base ``DiffusionEngine`` unchanged.
+
+    Architecture note: in the multiproc setup the KV cache lives in the worker /
+    runner process (GPU side), co-located with the model and KV tensors — so the
+    actual KV *body* is :class:`~vllm_omni.experimental.ar_diffusion.kv_cache.manager.ARDiffusionKVCache`, owned
+    by :class:`~vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner`. ``ARDiffusionEngine`` itself is the
+    thin selection / injection seam; it wires the AR-Diffusion executor → worker → runner
+    so DreamZero's rollout runs against the runner-owned KV cache.
+    """
+
+    #: Workers of this engine build the AR-Diffusion runner (unless the config
+    #: sets an explicit ``diffusion_model_runner_cls`` override). Declared on
+    #: the class so routing needs no od_config mutation at engine init.
+    default_diffusion_model_runner_cls = AR_DIFFUSION_MODEL_RUNNER_CLS
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
new file mode 100644
index 000000000..ab29f6093
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py
@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AR-Diffusion engine-level KV cache helpers.
+
+Thin glue over vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` /
+``SlidingWindowManager``), used by the AR-Diffusion Engine to manage KV for
+AR-diffusion models. See ``BDE_doc/dreamzero_kv_phase1_plan.md``.
+
+Layout: ``config`` (public knob) · ``paged`` (engine-generic paging mechanics +
+chunk-window eviction spec) · ``manager`` (the ARDiffusionKVCache orchestrator + its
+request adapter / pool builders) · ``state`` (the model-facing ARDiffusionKVState bridge).
+"""
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.manager import (
+    ARDiffusionKVCache,
+    ARDiffusionRequestAdapter,
+    build_kv_manager,
+    compute_num_blocks,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
+    ChunkWindowManager,
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    chunk_slot_mapping,
+    compute_slot_mapping,
+    pool_write_chunk,
+    resident_block_ids,
+)
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+    ARDiffusionPagedForwardContext,
+    ARDiffusionPagedLayerContext,
+    ARDiffusionPagedLayerInputs,
+    ar_diffusion_paged_attention,
+    paged_write_attn,
+    set_current_paged_kv_cache,
+)
+
+__all__ = [
+    "ARDiffusionKVCache",
+    "ARDiffusionKVConfig",
+    "ARDiffusionPagedForwardContext",
+    "ARDiffusionPagedLayerContext",
+    "ARDiffusionPagedLayerInputs",
+    "ARDiffusionRequestAdapter",
+    "ChunkWindowManager",
+    "ChunkWindowSpec",
+    "allocate_kv_pool_with_views",
+    "ar_diffusion_paged_attention",
+    "build_kv_manager",
+    "chunk_slot_mapping",
+    "compute_num_blocks",
+    "compute_slot_mapping",
+    "paged_write_attn",
+    "pool_write_chunk",
+    "resident_block_ids",
+    "set_current_paged_kv_cache",
+]
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/config.py b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
new file mode 100644
index 000000000..4b3c73775
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py
@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration for AR-Diffusion engine-level KV cache management."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ARDiffusionKVConfig:
+    """Settings for the AR-Diffusion paged KV cache.
+
+    Disabled by default: when ``enable`` is False the AR-Diffusion engine behaves exactly
+    like the base ``DiffusionEngine`` (no pool, no paged KV).
+    """
+
+    enable: bool = False
+    # Persistent KV tokens materialized per AR chunk. For DreamZero this is
+    # ``num_frame_per_block * frame_seqlen``.
+    chunk_size: int = 0
+    # Resident window in chunks. ``None`` means full attention (no eviction).
+    window_chunks: int | None = None
+    # Protected leading chunks (attention sink); never evicted.
+    sink_chunks: int = 0
+    # DreamZero-style window reset vs. VGGT-style sliding replace.
+    reset_at_boundary: bool = False
+    # Fraction of free device memory budgeted for the AR-Diffusion KV pool.
+    gpu_memory_fraction: float = 0.1
+    # When CUDA graph / torch.compile is on (not enforce_eager), pre-capture the
+    # DiT graphs for every window-fill shape at load time via a synthetic rollout,
+    # so the serving run is fast from the first chunk. No effect when eager.
+    warmup_cudagraph: bool = True
+    # Also capture the post-window-boundary (reset-cycle) forward during warm-up.
+    warmup_capture_reset: bool = False
+
+    @property
+    def sliding_window(self) -> int | None:
+        """Window size in tokens, or ``None`` for full attention."""
+        if self.window_chunks is None:
+            return None
+        return self.window_chunks * self.chunk_size
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
new file mode 100644
index 000000000..4b3b5d788
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py
@@ -0,0 +1,497 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ARDiffusionKVCache — the engine-level KV cache orchestrator for one AR-Diffusion model.
+
+This is the *body* of AR-Diffusion's KV management: it owns a vLLM ``KVCacheManager`` (a
+single chunk-window group) and the per-request adapter lifecycle, and exposes the
+per-chunk operations a rollout needs — allocate, slot mapping, commit, window
+lookup, free. It lives in the model runner (worker / GPU side), co-located with
+the model and the KV tensors; the DreamZero pipeline calls these methods during a
+rollout. The main-process ``ARDiffusionEngine`` only selects the engine and is otherwise
+thin.
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+from collections.abc import Sequence
+
+import torch
+from vllm.logger import init_logger
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheSpec,
+    KVCacheTensor,
+)
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import (
+    ChunkWindowSpec,
+    allocate_kv_pool_with_views,
+    chunk_slot_mapping,
+    pool_write_chunk,
+    resident_block_ids,
+)
+
+_log = init_logger(__name__)
+
+
+class ARDiffusionRequestAdapter:
+    """Duck-types the subset of ``vllm.v1.request.Request`` that the
+    ``KVCacheManager`` reads (``allocate_slots`` / ``get_computed_blocks`` /
+    ``free`` and the coordinator they call into).
+
+    It is intentionally NOT a full ``Request``. The conformance test exercises a
+    real ``KVCacheManager`` against this adapter so the surface cannot silently
+    drift across vLLM versions.
+
+    A AR-Diffusion request advances one *chunk* at a time: ``allocate_slots`` is called
+    once per chunk and ``num_computed_tokens`` advances only when a chunk is
+    committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk
+    reuse the same slots.
+    """
+
+    def __init__(
+        self,
+        request_id: str,
+        *,
+        chunk_size: int,
+        prefill_prefix_tokens: int = 0,
+    ) -> None:
+        self.request_id = request_id
+        self._chunk_size = chunk_size
+        self._prefill = prefill_prefix_tokens
+        self._completed_chunks = 0
+        # Filled only when cross-request prefix reuse is enabled (Phase 3).
+        self.block_hashes: list = []
+        self.skip_reading_prefix_cache = True
+        self.num_preemptions = 0
+        # vLLM watermark gate reads this; map the request lifecycle onto it.
+        self.status = RequestStatus.WAITING
+
+    @property
+    def num_computed_tokens(self) -> int:
+        """Persistent KV already materialized (committed chunks + prefill)."""
+        return self._prefill + self._completed_chunks * self._chunk_size
+
+    @property
+    def num_tokens(self) -> int:
+        """Total tokens once the in-flight chunk is committed."""
+        return self._prefill + (self._completed_chunks + 1) * self._chunk_size
+
+    @property
+    def num_prompt_tokens(self) -> int:
+        """The prefill prefix length (read by ``cache_blocks`` when caching)."""
+        return self._prefill
+
+    @property
+    def completed_chunks(self) -> int:
+        return self._completed_chunks
+
+    def on_chunk_committed(self) -> None:
+        """Advance by one chunk. Call once per chunk, not per denoise step."""
+        self._completed_chunks += 1
+
+
+def compute_num_blocks(
+    available_bytes: int,
+    gpu_memory_fraction: float,
+    page_size_bytes: int,
+) -> int:
+    """Number of KV blocks that fit in ``fraction`` of the memory budget."""
+    if page_size_bytes <= 0:
+        raise ValueError(f"page_size_bytes must be positive, got {page_size_bytes}")
+    if not 0.0 < gpu_memory_fraction <= 1.0:
+        raise ValueError(f"gpu_memory_fraction must be in (0, 1], got {gpu_memory_fraction}")
+    budget = int(available_bytes * gpu_memory_fraction)
+    return max(0, budget // page_size_bytes)
+
+
+def build_kv_manager(
+    spec: KVCacheSpec,
+    layer_names: Sequence[str],
+    num_blocks: int,
+    max_model_len: int,
+    *,
+    enable_caching: bool = False,
+) -> KVCacheManager:
+    """Build a ``KVCacheManager`` with a single KV cache group for ``spec``.
+
+    Args:
+        spec: The KV cache spec for the group (e.g. a ``ChunkWindowSpec``).
+        layer_names: Attention layers sharing this group's block table.
+        num_blocks: Total physical blocks in the pool.
+        max_model_len: Upper bound on a request's sequence length.
+        enable_caching: Cross-request prefix caching (Phase 3); off in Phase 1.
+    """
+    layer_names = list(layer_names)
+    group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)
+    tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)]
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=tensors,
+        kv_cache_groups=[group],
+    )
+    kwargs = dict(max_model_len=max_model_len, hash_block_size=spec.block_size, enable_caching=enable_caching)
+    params = inspect.signature(KVCacheManager).parameters
+    if "scheduler_block_size" in params:
+        kwargs["scheduler_block_size"] = spec.block_size
+    if "max_num_batched_tokens" in params:
+        kwargs["max_num_batched_tokens"] = max_model_len
+    return KVCacheManager(config, **kwargs)
+
+
+class ARDiffusionKVCache:
+    """Owns the paged KV pool + per-request lifecycle for a AR-Diffusion model.
+
+    Build once per loaded model (dimensions known); then per request:
+    ``begin_request`` → per chunk (``allocate_chunk`` → ``chunk_write_slots`` →
+    [model writes K/V] → ``commit_chunk``) → ``end_request``.
+    """
+
+    def __init__(
+        self,
+        config: ARDiffusionKVConfig,
+        *,
+        num_layers: int,
+        num_kv_heads: int,
+        head_size: int,
+        dtype: torch.dtype,
+        block_size: int,
+        max_model_len: int,
+        available_bytes: int,
+        cross_attn_length: int = 0,
+        cross_attn_img_length: int = 0,
+        device: torch.device | None = None,
+        local_branches: int = 2,
+        num_frame_per_block: int = 1,
+    ) -> None:
+        if not config.enable:
+            raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig")
+        if config.window_chunks is None:
+            raise ValueError("Phase 1 requires a bounded window (window_chunks)")
+        if config.chunk_size <= 0:
+            raise ValueError("ARDiffusionKVConfig.chunk_size must be set (> 0)")
+        if local_branches not in (1, 2):
+            raise ValueError(f"local_branches must be 1 or 2, got {local_branches}")
+
+        self.config = config
+        # How many CFG branches allocate from THIS rank's pool. Under
+        # CFG-parallel x2 each rank executes exactly one branch (rank0 pos,
+        # rank1 neg) and the other branch's lazy contexts never allocate, so
+        # sizing for both would leave ~half the pool as idle capacity. A
+        # single-process run (TP1 / offline) executes both branches -> 2.
+        self.local_branches = local_branches
+        self.num_frame_per_block = max(1, int(num_frame_per_block))
+        self.block_size = block_size
+        self.num_layers = num_layers
+        self.num_kv_heads = num_kv_heads
+        self.head_size = head_size
+        self.dtype = dtype
+        self.cross_attn_length = cross_attn_length
+        self.device = device or torch.device("cpu")
+        self._adapters: dict[str, ARDiffusionRequestAdapter] = {}
+
+        # -- cross-attention pool (static, one-time-fill) -----------------------
+        # Cross-attn KV is computed once from the text encoder and never changes.
+        # Allocate a per-layer contiguous tensor for each branch: (2, text_len,
+        # kv_heads, head_dim) where dim-0 = [pos, neg].  Allocated separately
+        # from the self-attn block pool (static, small — ~400 MiB for
+        # DreamZero I2V).  Both pools draw from the same GPU free memory but
+        # the cross-attn pool is sized directly rather than via the block pool
+        # budget (no eviction/chunk lifecycle needed).
+        self.cross_attn_img_length = cross_attn_img_length
+        self._cross_k: list[torch.Tensor] = []
+        self._cross_v: list[torch.Tensor] = []
+        # I2V image-token cross-attn pool. Like the text k/v, the image-token
+        # k_img/v_img are session-invariant (the conditioning image doesn't change),
+        # so they are cached once and read every denoise step — see
+        # WanI2VCrossAttention (#4154 caches these model-side too). Empty for T2V.
+        self._cross_k_img: list[torch.Tensor] = []
+        self._cross_v_img: list[torch.Tensor] = []
+
+        # Bytes the cross-attn pools consume directly (K+V, pos+neg, all layers).
+        # Deducted from the self-attn paged-pool budget below so the two
+        # allocations together stay within the GPU free-memory budget.
+        def _cross_pool_bytes(length: int) -> int:
+            return 2 * 2 * length * num_kv_heads * head_size * dtype.itemsize * num_layers
+
+        cross_total_bytes = 0
+        if device is not None and cross_attn_length > 0:
+            cross_shape = (2, cross_attn_length, num_kv_heads, head_size)
+            cross_bytes = _cross_pool_bytes(cross_attn_length)
+            cross_total_bytes += cross_bytes
+            for _ in range(num_layers):
+                self._cross_k.append(torch.empty(cross_shape, dtype=dtype, device=device))
+                self._cross_v.append(torch.empty(cross_shape, dtype=dtype, device=device))
+            _log.info(
+                "AR-Diffusion cross-attn pool: %d layers × (%d tok × %d heads × %d) = %.1f MiB",
+                num_layers,
+                cross_attn_length,
+                num_kv_heads,
+                head_size,
+                cross_bytes / (1024 * 1024),
+            )
+            if cross_attn_img_length > 0:
+                img_shape = (2, cross_attn_img_length, num_kv_heads, head_size)
+                cross_total_bytes += _cross_pool_bytes(cross_attn_img_length)
+                for _ in range(num_layers):
+                    self._cross_k_img.append(torch.empty(img_shape, dtype=dtype, device=device))
+                    self._cross_v_img.append(torch.empty(img_shape, dtype=dtype, device=device))
+                _log.info(
+                    "AR-Diffusion cross-attn IMG pool: %d layers × %d img-tok (I2V)",
+                    num_layers,
+                    cross_attn_img_length,
+                )
+
+        self.spec = ChunkWindowSpec(
+            block_size=block_size,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=dtype,
+            sliding_window=config.window_chunks * config.chunk_size,
+            chunk_size=config.chunk_size,
+            window_chunks=config.window_chunks,
+            sink_chunks=config.sink_chunks,
+            reset_at_boundary=config.reset_at_boundary,
+        )
+        # Each pool block spans all layers' K/V, so size against the per-layer
+        # page size times the layer count.
+        # Size the self-attn pool against the memory left after the cross-attn
+        # pools (allocated above), so cross + self-attn stays within the budget.
+        num_blocks = compute_num_blocks(
+            max(0, available_bytes - cross_total_bytes),
+            config.gpu_memory_fraction,
+            self.spec.page_size_bytes * num_layers,
+        )
+        # Floor: one forward needs the resident window plus the in-flight chunk
+        # (num_frame_per_block frame-blocks) for every branch THIS rank runs,
+        # with a little eviction-transient headroom. The memory-fraction
+        # heuristic can under-size this once block_size grows — e.g. frame-granular
+        # paging at the true frame_seqlen makes each block larger and the pool
+        # fewer-blocks — so guarantee the minimum the rollout cannot run without,
+        # otherwise allocate_chunk hits an exhausted pool mid-forward.
+        min_blocks = self.local_branches * (config.window_chunks + self.num_frame_per_block) + 2
+        if num_blocks < min_blocks:
+            _log.warning(
+                "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block "
+                "floor (%d local CFG branch(es) x (window_chunks=%d + num_frame_per_block=%d) + 2 headroom)",
+                num_blocks,
+                min_blocks,
+                self.local_branches,
+                config.window_chunks,
+                self.num_frame_per_block,
+            )
+            num_blocks = min_blocks
+        layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)]
+        self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len)
+        self.managed_num_blocks = num_blocks
+        self.num_blocks = num_blocks
+        # Scratch blocks are outside KVCacheManager ownership. They hold current
+        # denoise-step video KV when update_kv_cache=False and action/state KV in
+        # both modes; they are reused every forward and never committed.
+        # Default 4: measured high-water is num_frame_per_block video blocks (2)
+        # + 1 action block per branch per forward, +1 margin.
+        scratch_per_branch = int(os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "4"))
+        if scratch_per_branch <= 0:
+            raise ValueError("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH must be positive")
+        self.scratch_blocks_per_branch = scratch_per_branch
+        self.scratch_num_blocks = self.local_branches * scratch_per_branch
+        self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks
+        self.null_block_id = self.manager.block_pool.null_block.block_id
+
+        # Allocate the per-layer paged K/V pools on the given device.
+        self._kv_pools: list[torch.Tensor] = []
+        self._k_pools: list[torch.Tensor] = []
+        self._v_pools: list[torch.Tensor] = []
+        if device is not None:
+            self._kv_pools, self._k_pools, self._v_pools = allocate_kv_pool_with_views(
+                self.num_blocks_total,
+                block_size,
+                num_layers,
+                num_kv_heads,
+                head_size,
+                dtype,
+                device,
+            )
+
+    # -- cross-attention pool access -------------------------------------------
+    # Cross-attn KV is static once populated — write once (from text encoder),
+    # read many (every denoising step). Not managed through the paged block pool.
+
+    def write_cross_kv(
+        self,
+        layer_idx: int,
+        is_negative: bool,
+        k: torch.Tensor | None,
+        v: torch.Tensor | None,
+        k_img: torch.Tensor | None = None,
+        v_img: torch.Tensor | None = None,
+    ) -> None:
+        """Write one layer's cross-attn K/V into the pool.
+
+        ``k`` / ``v``: ``(B, text_len, tp_num_heads, head_dim)``; only batch-0
+        is copied (B=1 for inference). ``None`` skips the text half (window
+        restarts keep the still-valid text K/V and rewrite only the image half).
+        The ``is_negative`` flag selects the correct CFG branch slot. ``k_img`` /
+        ``v_img`` (I2V image tokens, ``(B, 257, ...)``) are written when the
+        image pool is allocated.
+        """
+        branch = 1 if is_negative else 0
+        if k is not None:
+            self._cross_k[layer_idx][branch].copy_(k[0])
+            self._cross_v[layer_idx][branch].copy_(v[0])
+        if k_img is not None and self._cross_k_img:
+            self._cross_k_img[layer_idx][branch].copy_(k_img[0])
+            self._cross_v_img[layer_idx][branch].copy_(v_img[0])
+
+    def read_cross_kv(self, layer_idx: int, is_negative: bool) -> dict:
+        """Return a pool-backed cross-attn cache dict for one layer.
+
+        The dict matches the ``{"is_init": True, "k": Tensor, "v": Tensor}``
+        convention the cross-attention module expects — it reads from the pool
+        slice rather than from the lazy-initialised model-local dict. For I2V,
+        ``k_img`` / ``v_img`` are added so the image-token cache (added by #4154)
+        reads from the pool too.
+        """
+        branch = 1 if is_negative else 0
+        k = self._cross_k[layer_idx][branch].unsqueeze(0)  # (1, L, heads, dim)
+        v = self._cross_v[layer_idx][branch].unsqueeze(0)
+        cache = {"is_init": True, "k": k, "v": v}
+        if self._cross_k_img:
+            cache["k_img"] = self._cross_k_img[layer_idx][branch].unsqueeze(0)
+            cache["v_img"] = self._cross_v_img[layer_idx][branch].unsqueeze(0)
+        return cache
+
+    # -- request lifecycle ---------------------------------------------------
+
+    def begin_request(self, request_id: str, *, prefill_prefix_tokens: int = 0) -> ARDiffusionRequestAdapter:
+        adapter = ARDiffusionRequestAdapter(
+            request_id,
+            chunk_size=self.spec.chunk_size,
+            prefill_prefix_tokens=prefill_prefix_tokens,
+        )
+        self._adapters[request_id] = adapter
+        _log.debug("AR-Diffusion begin_request: req=%s prefill=%d", request_id, prefill_prefix_tokens)
+        return adapter
+
+    def end_request(self, adapter: ARDiffusionRequestAdapter) -> None:
+        _log.debug(
+            "AR-Diffusion end_request: req=%s chunks=%d free=%d",
+            adapter.request_id,
+            adapter.completed_chunks,
+            self.manager.block_pool.get_num_free_blocks(),
+        )
+        self.manager.free(adapter)
+        self._adapters.pop(adapter.request_id, None)
+
+    # -- per-chunk operations ------------------------------------------------
+
+    def allocate_chunk(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
+        """Allocate a chunk's blocks (evicting out-of-window blocks first).
+
+        Returns the request's full block table (incl. null_block placeholders).
+        """
+        blocks = self.manager.allocate_slots(adapter, num_new_tokens=self.spec.chunk_size)
+        if blocks is None:
+            raise RuntimeError("AR-Diffusion KV pool exhausted while allocating a chunk")
+        table = self.block_table(adapter)
+        resident = resident_block_ids(table, self.null_block_id)
+        _log.debug(
+            "AR-Diffusion allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d",
+            adapter.request_id,
+            adapter.completed_chunks,
+            len(table),
+            len(resident),
+            self.manager.block_pool.get_num_free_blocks(),
+        )
+        return table
+
+    def allocate_token_slots(self, adapter: ARDiffusionRequestAdapter, num_tokens: int) -> list[int]:
+        """Allocate managed blocks for an in-flight video span without committing it."""
+        if num_tokens <= 0:
+            raise ValueError(f"num_tokens must be positive, got {num_tokens}")
+        blocks = self.manager.allocate_slots(adapter, num_new_tokens=num_tokens)
+        if blocks is None:
+            raise RuntimeError("AR-Diffusion KV pool exhausted while allocating paged attention slots")
+        return self.block_table(adapter)
+
+    def block_table(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
+        return list(self.manager.get_block_ids(adapter.request_id)[0])
+
+    def chunk_write_slots(self, adapter: ARDiffusionRequestAdapter) -> torch.Tensor:
+        """Slot mapping for the in-flight chunk — the K/V write target."""
+        return chunk_slot_mapping(
+            self.block_table(adapter),
+            adapter.num_computed_tokens,
+            self.spec.chunk_size,
+            self.block_size,
+        )
+
+    def scratch_block_ids(self, is_negative: bool, start: int, count: int) -> list[int]:
+        """Return branch-local scratch block ids outside manager ownership."""
+        if count < 0 or start < 0:
+            raise ValueError(f"scratch start/count must be non-negative, got start={start}, count={count}")
+        if start + count > self.scratch_blocks_per_branch:
+            raise RuntimeError(
+                "AR-Diffusion paged attention scratch blocks exhausted: "
+                f"need [{start}, {start + count}) of {self.scratch_blocks_per_branch}. "
+                "Increase AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH."
+            )
+        # With one local branch (CFG-parallel x2) the rank's only branch maps to
+        # slot 0 whichever CFG side it is; the other branch never allocates here.
+        branch_offset = self.scratch_blocks_per_branch if (is_negative and self.local_branches == 2) else 0
+        base = self.managed_num_blocks + branch_offset + start
+        return list(range(base, base + count))
+
+    def key_cache(self, layer_idx: int) -> torch.Tensor:
+        return self._kv_pools[layer_idx][0]
+
+    def value_cache(self, layer_idx: int) -> torch.Tensor:
+        return self._kv_pools[layer_idx][1]
+
+    def window_block_ids(self, adapter: ARDiffusionRequestAdapter) -> list[int]:
+        """Resident (non-null) managed blocks visible to paged attention."""
+        return resident_block_ids(self.block_table(adapter), self.null_block_id)
+
+    def commit_chunk(self, adapter: ARDiffusionRequestAdapter) -> None:
+        """Advance the adapter by one chunk after its K/V is written.
+
+        This standalone primitive is used by low-level manager tests. In the
+        DreamZero paged-attention path, :meth:`ARDiffusionKVState.commit_paged_context`
+        advances the adapter only after the forward succeeds. Call once per
+        committed chunk, not per denoise step.
+        """
+        _log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks)
+        adapter.on_chunk_committed()
+        _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks)
+
+    # -- pool-backed K/V access --------------------------------------------
+
+    def write_chunk_kv(
+        self,
+        layer_index: int,
+        new_k: torch.Tensor,
+        new_v: torch.Tensor,
+        adapter: ARDiffusionRequestAdapter,
+    ) -> None:
+        """Write one layer's committed-chunk K/V into the pool."""
+        slots = self.chunk_write_slots(adapter)
+        _log.debug(
+            "AR-Diffusion write: req=%s layer=%d chunk=%d shapes=%s dev=%s",
+            adapter.request_id,
+            layer_index,
+            adapter.completed_chunks,
+            (tuple(new_k.shape), tuple(new_v.shape)),
+            slots.device,
+        )
+        pool_write_chunk(
+            self._k_pools[layer_index],
+            self._v_pools[layer_index],
+            new_k,
+            new_v,
+            slots,
+        )
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py
new file mode 100644
index 000000000..a39a66b91
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py
@@ -0,0 +1,208 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generic paging mechanics + chunk-window eviction spec for the AR-Diffusion engine.
+
+Engine-generic, model-agnostic primitives — the layer a second model (e.g. the
+Cosmos port) reuses unchanged. Three concerns live here:
+
+* **Slot mapping** — absolute token positions → physical KV-cache slots, the
+  standard PagedAttention layout ``slot(pos) = block_id(pos) * block_size +
+  (pos % block_size)``.
+* **Pool I/O** — allocate per-layer FlashAttention-compatible paged K/V pools
+  plus flat compatibility views for direct slot writes.
+* **Chunk-window eviction** — a ``SlidingWindowSpec`` subclass whose unit is a
+  *chunk* (``sliding_window = window_chunks * chunk_size``) plus a manager that
+  evicts at chunk boundaries. Memory policy / refcounting / ``null_block``
+  replacement stay in vLLM's ``BlockPool``; only the token-skip math is here.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import torch
+from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
+from vllm.v1.kv_cache_interface import SlidingWindowSpec
+from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec
+
+# ── Slot mapping ────────────────────────────────────────────────────────────
+
+
+def compute_slot_mapping(
+    block_ids: Sequence[int],
+    positions: torch.Tensor | Sequence[int],
+    block_size: int,
+) -> torch.Tensor:
+    """Map absolute token positions to physical KV-cache slots.
+
+    Args:
+        block_ids: physical block id per block index (the request's block table).
+        positions: absolute token positions to map (1-D).
+        block_size: tokens per block.
+
+    Returns:
+        ``LongTensor`` of physical slots, one per position.
+    """
+    if block_size <= 0:
+        raise ValueError(f"block_size must be positive, got {block_size}")
+    table = torch.as_tensor(block_ids, dtype=torch.long)
+    pos = torch.as_tensor(positions, dtype=torch.long)
+    block_index = torch.div(pos, block_size, rounding_mode="floor")
+    offset = pos % block_size
+    return table[block_index] * block_size + offset
+
+
+def chunk_slot_mapping(
+    block_ids: Sequence[int],
+    num_computed_tokens: int,
+    chunk_size: int,
+    block_size: int,
+) -> torch.Tensor:
+    """Slot mapping for the in-flight chunk's tokens (the commit write target).
+
+    The chunk occupies absolute positions
+    ``[num_computed_tokens, num_computed_tokens + chunk_size)``.
+    """
+    positions = torch.arange(
+        num_computed_tokens,
+        num_computed_tokens + chunk_size,
+        dtype=torch.long,
+    )
+    return compute_slot_mapping(block_ids, positions, block_size)
+
+
+def resident_block_ids(block_ids: Sequence[int], null_block_id: int) -> list[int]:
+    """Real (non-null) blocks currently resident, in table order.
+
+    These are the blocks the read path gathers the attention window from;
+    out-of-window positions are the shared ``null_block`` and are excluded.
+    """
+    return [int(b) for b in block_ids if int(b) != null_block_id]
+
+
+# ── Pool I/O (allocate / write) ─────────────────────────────────────────────
+
+
+def allocate_kv_pool_with_views(
+    num_blocks: int,
+    block_size: int,
+    num_layers: int,
+    num_kv_heads: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    """Allocate vLLM-style paged KV pools plus flat slot-write views.
+
+    The owning tensor follows FlashAttention's block-table cache layout:
+    ``(2, num_blocks, block_size, num_kv_heads, head_dim)``, where dim-0 is
+    ``[K, V]``.  The flat K/V views keep ``slot = block_id * block_size +
+    offset`` writes simple without changing the kernel-facing cache layout.
+    """
+    kv_pools: list[torch.Tensor] = []
+    k_pools: list[torch.Tensor] = []
+    v_pools: list[torch.Tensor] = []
+    for _ in range(num_layers):
+        kv = torch.empty(2, num_blocks, block_size, num_kv_heads, head_dim, dtype=dtype, device=device)
+        kv_pools.append(kv)
+        k_pools.append(kv[0].reshape(num_blocks * block_size, num_kv_heads, head_dim))
+        v_pools.append(kv[1].reshape(num_blocks * block_size, num_kv_heads, head_dim))
+    return kv_pools, k_pools, v_pools
+
+
+def pool_write_chunk(
+    k_pool: torch.Tensor,
+    v_pool: torch.Tensor,
+    new_k: torch.Tensor,
+    new_v: torch.Tensor,
+    slot_mapping: torch.Tensor,
+) -> None:
+    """Write the committed chunk's K/V into pool slots (in place).
+
+    ``new_k`` / ``new_v`` shape: ``(batch, chunk_size, num_kv_heads, head_dim)``.
+    ``slot_mapping`` shape: ``(chunk_size,)`` — one physical slot per token position.
+    The batch dim is written to the same slots (multi-batch writes the same
+    positions; the caller is responsible for per-sequence bookkeeping).
+
+    This is the *commit* write — called once per chunk, after the last denoise step.
+    """
+    k_pool[slot_mapping] = new_k[0]  # batch index 0
+    v_pool[slot_mapping] = new_v[0]
+
+
+def _drop_batch(t: torch.Tensor) -> torch.Tensor:
+    """``(1, L, n_heads, head_dim)`` -> ``(L, n_heads, head_dim)``; pass 3-D through."""
+    if t.dim() == 4 and t.shape[0] == 1:
+        return t[0]
+    return t
+
+
+# ── Chunk-window eviction (spec + manager) ──────────────────────────────────
+
+
+def chunk_window_skipped_tokens(
+    num_computed_tokens: int,
+    *,
+    chunk_size: int,
+    sliding_window: int,
+    sink_chunks: int,
+    reset_at_boundary: bool,
+) -> int:
+    """Tokens outside the resident chunk window, snapped to a chunk boundary.
+
+    Pure function so the eviction policy is unit-testable without constructing a
+    manager. Two strategies:
+
+    - ``reset_at_boundary`` (DreamZero): at each chunk boundary everything past
+      the sink is dropped.
+    - otherwise (VGGT-style sliding replace): keep the last ``window`` tokens
+      (plus the sink); the skip count snaps down to a chunk boundary so a chunk
+      is never half-evicted.
+    """
+    sink = sink_chunks * chunk_size
+    if reset_at_boundary:
+        completed = (num_computed_tokens // chunk_size) * chunk_size
+        return max(0, completed - sink)
+    skipped = max(0, num_computed_tokens - sliding_window - sink)
+    return (skipped // chunk_size) * chunk_size
+
+
+class ChunkWindowManager(SlidingWindowManager):
+    """``SlidingWindowManager`` that evicts at chunk boundaries.
+
+    ``self.sliding_window`` is set by the base ``__init__``; the chunk fields are
+    read from ``self.kv_cache_spec`` (a :class:`ChunkWindowSpec`).
+    """
+
+    def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
+        spec = self.kv_cache_spec
+        return chunk_window_skipped_tokens(
+            num_computed_tokens,
+            chunk_size=spec.chunk_size,
+            sliding_window=self.sliding_window,
+            sink_chunks=spec.sink_chunks,
+            reset_at_boundary=spec.reset_at_boundary,
+        )
+
+
+# Register so KVCacheManager resolves ChunkWindowSpec to ChunkWindowManager.
+# Dispatch walks the spec's MRO, so without explicit registration the subclass
+# would silently fall back to the parent SlidingWindowManager (override ignored).
+# uniform_type_base_spec=None => its own KV cache group.
+@register_kv_cache_spec(manager_class=ChunkWindowManager, uniform_type_base_spec=None)
+@dataclass(frozen=True, kw_only=True)
+class ChunkWindowSpec(SlidingWindowSpec):
+    # sliding_window (inherited) MUST equal window_chunks * chunk_size.
+    chunk_size: int
+    window_chunks: int
+    sink_chunks: int = 0
+    reset_at_boundary: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.sliding_window != self.window_chunks * self.chunk_size:
+            raise ValueError(
+                "ChunkWindowSpec.sliding_window must equal "
+                f"window_chunks * chunk_size ({self.window_chunks} * "
+                f"{self.chunk_size}), got {self.sliding_window}"
+            )
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
new file mode 100644
index 000000000..2038e496d
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
@@ -0,0 +1,550 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Paged self-attention helpers for AR-Diffusion DreamZero KV reuse."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, ClassVar, NamedTuple
+
+import torch
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping
+
+# Set by ARDiffusionPagedForwardContext.prepare() before each branch forward and
+# read by the fused write+attend custom op below. The pools are process-lifetime
+# allocations owned by one ARDiffusionKVCache per worker process (CFG-parallel /
+# TP ranks are separate processes), mirroring vLLM's forward-context pattern for
+# keeping multi-GiB cache tensors out of the compiled graph's inputs.
+_CURRENT_PAGED_KV_CACHE: Any = None
+
+
+def set_current_paged_kv_cache(kv_cache: Any) -> None:
+    global _CURRENT_PAGED_KV_CACHE
+    _CURRENT_PAGED_KV_CACHE = kv_cache
+
+
+_LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {}
+
+
+def _layer_idx_tensor(layer_idx: int) -> torch.Tensor:
+    t = _LAYER_IDX_TENSORS.get(layer_idx)
+    if t is None:
+        t = torch.tensor(layer_idx, dtype=torch.int64)
+        _LAYER_IDX_TENSORS[layer_idx] = t
+    return t
+
+
+class ARDiffusionPagedLayerInputs(NamedTuple):
+    """Compiled-region payload for one layer's paged self-attention.
+
+    A NamedTuple of plain tensors + ints so ``torch.compile`` treats every field
+    as a pytree graph input (no object-attribute guards, no recompiles when only
+    tensor *values* change). All layers of one branch forward share the same
+    metadata tensor objects, built once by ``prepare()``.
+
+    ``layer_idx`` is a 0-dim CPU tensor, NOT a python int: all 40 DiT blocks
+    share one compiled code object, and an int here becomes a per-layer dynamo
+    value guard (``layer_idx == k``) — 40 cache variants that blow the
+    recompile limit. A tensor input guards on shape/dtype only, so one graph
+    serves every layer.
+    """
+
+    layer_idx: torch.Tensor
+    seq_len: int
+    video_slots: torch.Tensor
+    action_slots: torch.Tensor
+    block_table: torch.Tensor
+    query_start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    max_query_len: int
+    max_seq_len: int
+
+
+@dataclass
+class ARDiffusionPagedForwardContext:
+    """Mutable branch-level state shared by all layer contexts in one forward."""
+
+    kv_cache: Any
+    adapter: Any
+    is_negative: bool
+    history_block_ids: list[int]
+    seq_len: int
+    commit_current: bool
+    max_video_tokens: int
+    current_video_block_ids: list[int] = field(default_factory=list)
+    current_video_slot_mapping: torch.Tensor | None = None
+    action_scratch_block_ids: list[int] = field(default_factory=list)
+    action_slot_mapping: torch.Tensor | None = None
+    query_len: int = 0
+    kv_len: int = 0
+    _allocated_video: bool = False
+    _committed: bool = False
+    _action_len: int = 0
+    # Set once by prepare(); shared by all layers of the branch forward.
+    block_table: torch.Tensor | None = None
+    query_start_loc: torch.Tensor | None = None
+    seq_lens: torch.Tensor | None = None
+    max_query_len: int = 0
+    max_seq_len: int = 0
+    _prepared: bool = False
+
+    @property
+    def block_size(self) -> int:
+        return int(self.kv_cache.block_size)
+
+    @property
+    def num_current_video_blocks(self) -> int:
+        if self.seq_len % self.block_size != 0:
+            raise AssertionError(
+                "AR-Diffusion paged attention expects frame-aligned seq_len "
+                f"(multiple of block_size={self.block_size}), got {self.seq_len}"
+            )
+        return self.seq_len // self.block_size
+
+    def ensure_video_slots(self, device: torch.device) -> None:
+        """Allocate/write targets for the current video tokens, once per branch."""
+        if self._allocated_video:
+            return
+
+        n_blocks = self.num_current_video_blocks
+        if self.commit_current:
+            start = int(self.adapter.num_computed_tokens)
+            self.kv_cache.allocate_token_slots(self.adapter, self.seq_len)
+            table = self.kv_cache.block_table(self.adapter)
+            start_block = start // self.block_size
+            self.current_video_block_ids = [int(b) for b in table[start_block : start_block + n_blocks]]
+            positions = torch.arange(start, start + self.seq_len, dtype=torch.long)
+            self.current_video_slot_mapping = compute_slot_mapping(table, positions, self.block_size).to(device=device)
+        else:
+            self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.is_negative, 0, n_blocks)
+            positions = torch.arange(self.seq_len, dtype=torch.long)
+            self.current_video_slot_mapping = compute_slot_mapping(
+                self.current_video_block_ids,
+                positions,
+                self.block_size,
+            ).to(device=device)
+        self._allocated_video = True
+
+    def ensure_action_slots(self, action_len: int, device: torch.device) -> None:
+        """Reserve scratch slots for action/state K/V, if present."""
+        if action_len <= 0:
+            self.action_scratch_block_ids = []
+            self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
+            self._action_len = 0
+            return
+
+        self.ensure_video_slots(device)
+        if self.action_slot_mapping is not None and self._action_len == action_len:
+            return
+
+        action_blocks = (action_len + self.block_size - 1) // self.block_size
+        scratch_offset = 0 if self.commit_current else len(self.current_video_block_ids)
+        self.action_scratch_block_ids = self.kv_cache.scratch_block_ids(
+            self.is_negative,
+            scratch_offset,
+            action_blocks,
+        )
+        positions = torch.arange(action_len, dtype=torch.long)
+        self.action_slot_mapping = compute_slot_mapping(
+            self.action_scratch_block_ids,
+            positions,
+            self.block_size,
+        ).to(device=device)
+        self._action_len = action_len
+
+    def video_block_table(self, device: torch.device) -> tuple[list[int], int]:
+        self.ensure_video_slots(device)
+        if self.max_video_tokens % self.block_size != 0:
+            raise AssertionError(
+                "AR-Diffusion paged attention requires max_video_tokens to be block-aligned, "
+                f"got max_video_tokens={self.max_video_tokens}, block_size={self.block_size}"
+            )
+        all_video_blocks = self.history_block_ids + self.current_video_block_ids
+        max_video_blocks = self.max_video_tokens // self.block_size
+        visible_video_blocks = all_video_blocks[-max_video_blocks:] if max_video_blocks else []
+        video_len = min(len(all_video_blocks) * self.block_size, self.max_video_tokens)
+        return visible_video_blocks, video_len
+
+    def build_block_table(
+        self,
+        *,
+        action_len: int,
+        query_len: int,
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
+        """Build FlashAttention block-table metadata for one self-attn call.
+
+        The block table is tail-padded to a fixed width and ``max_seq_len`` is a
+        constant upper bound, so across window growth only tensor *values* change
+        — tensor shapes and the int consts stay stable for ``torch.compile``
+        (and, later, CUDA-graph capture). The kernel only dereferences the first
+        ``ceil(seq_lens/block_size)`` entries, so padding is never read.
+        """
+        video_blocks, video_len = self.video_block_table(device)
+        self.ensure_action_slots(action_len, device)
+        action_blocks = self.action_scratch_block_ids if action_len > 0 else []
+        block_ids = video_blocks + action_blocks
+        if not block_ids:
+            raise RuntimeError("AR-Diffusion paged attention needs at least current video KV blocks")
+
+        # Fixed capacity: full visible video window + one action-capacity block.
+        action_capacity_blocks = max(1, (action_len + self.block_size - 1) // self.block_size)
+        width = max(self.max_video_tokens // self.block_size + action_capacity_blocks, len(block_ids))
+        padded = block_ids + [0] * (width - len(block_ids))
+
+        self.query_len = int(query_len)
+        self.kv_len = int(video_len + action_len)
+        max_seq_len = int(self.max_video_tokens + action_capacity_blocks * self.block_size)
+        block_table = torch.tensor([padded], dtype=torch.int32, device=device)
+        query_start_loc = torch.tensor([0, self.query_len], dtype=torch.int32, device=device)
+        seq_lens = torch.tensor([self.kv_len], dtype=torch.int32, device=device)
+        return block_table, query_start_loc, seq_lens, self.query_len, max_seq_len
+
+    def prepare(self, device: torch.device, action_len: int, query_len: int) -> None:
+        """Host-side, once-per-branch-forward setup (called OUTSIDE torch.compile).
+
+        Allocates the current video/action slots (still lazy: only the branch a
+        CFG-parallel rank actually runs reaches its ``_forward_blocks``), builds
+        the padded block-table metadata ONCE for all layers, and publishes the
+        pool registry for the fused custom op. The compiled per-layer code then
+        only consumes prebuilt tensors via ``ARDiffusionPagedLayerInputs``.
+        """
+        if getattr(self, "_prepared", False):
+            return
+        self.ensure_video_slots(device)
+        (
+            self.block_table,
+            self.query_start_loc,
+            self.seq_lens,
+            self.max_query_len,
+            self.max_seq_len,
+        ) = self.build_block_table(action_len=action_len, query_len=query_len, device=device)
+        if self.action_slot_mapping is None:
+            self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device)
+        set_current_paged_kv_cache(self.kv_cache)
+        self._prepared = True
+
+    def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs:
+        if not getattr(self, "_prepared", False):
+            raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()")
+        return ARDiffusionPagedLayerInputs(
+            layer_idx=_layer_idx_tensor(layer_idx),
+            seq_len=int(self.seq_len),
+            video_slots=self.current_video_slot_mapping,
+            action_slots=self.action_slot_mapping,
+            block_table=self.block_table,
+            query_start_loc=self.query_start_loc,
+            seq_lens=self.seq_lens,
+            max_query_len=int(self.max_query_len),
+            max_seq_len=int(self.max_seq_len),
+        )
+
+    def mark_committed(self) -> None:
+        self._committed = True
+
+
+@dataclass
+class ARDiffusionPagedLayerContext:
+    """Layer-specific handle passed through DreamZero's existing ``kv_cache`` slot."""
+
+    is_ar_diffusion_paged_context: ClassVar[bool] = True
+    layer_idx: int
+    forward_ctx: ARDiffusionPagedForwardContext
+
+    @property
+    def kv_cache(self):
+        return self.forward_ctx.kv_cache
+
+    @property
+    def adapter(self):
+        return self.forward_ctx.adapter
+
+    @property
+    def is_negative(self) -> bool:
+        return self.forward_ctx.is_negative
+
+    @property
+    def history_block_ids(self) -> list[int]:
+        return self.forward_ctx.history_block_ids
+
+    @property
+    def current_video_block_ids(self) -> list[int]:
+        return self.forward_ctx.current_video_block_ids
+
+    @property
+    def current_video_slot_mapping(self) -> torch.Tensor | None:
+        return self.forward_ctx.current_video_slot_mapping
+
+    @property
+    def action_scratch_block_ids(self) -> list[int]:
+        return self.forward_ctx.action_scratch_block_ids
+
+    @property
+    def action_slot_mapping(self) -> torch.Tensor | None:
+        return self.forward_ctx.action_slot_mapping
+
+    @property
+    def seq_len(self) -> int:
+        return self.forward_ctx.seq_len
+
+    @property
+    def query_len(self) -> int:
+        return self.forward_ctx.query_len
+
+    @property
+    def kv_len(self) -> int:
+        return self.forward_ctx.kv_len
+
+    @property
+    def commit_current(self) -> bool:
+        return self.forward_ctx.commit_current
+
+    def to_layer_inputs(self) -> ARDiffusionPagedLayerInputs:
+        """Compiled-region payload; requires ``forward_ctx.prepare()`` first."""
+        return self.forward_ctx.layer_inputs(self.layer_idx)
+
+
+def is_ar_diffusion_paged_context(value: object) -> bool:
+    return isinstance(value, ARDiffusionPagedLayerContext)
+
+
+def _reference_paged_attention(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    softmax_scale: float,
+    *,
+    causal: bool,
+) -> torch.Tensor:
+    if causal:
+        raise NotImplementedError("AR-Diffusion paged self-attention uses causal=False")
+    outs: list[torch.Tensor] = []
+    block_size = key_cache.shape[1]
+    for i in range(seq_lens.shape[0]):
+        q_start = int(query_start_loc[i].item())
+        q_end = int(query_start_loc[i + 1].item())
+        kv_len = int(seq_lens[i].item())
+        q = query[q_start:q_end]
+        positions = torch.arange(kv_len, device=query.device)
+        logical_blocks = torch.div(positions, block_size, rounding_mode="floor")
+        offsets = positions % block_size
+        physical_blocks = block_table[i, logical_blocks].long()
+        k = key_cache[physical_blocks, offsets]
+        v = value_cache[physical_blocks, offsets]
+        scores = torch.einsum("qhd,khd->hqk", q.float(), k.float()) * float(softmax_scale)
+        probs = torch.softmax(scores, dim=-1).to(v.dtype)
+        outs.append(torch.einsum("hqk,khd->qhd", probs, v))
+    return torch.cat(outs, dim=0)
+
+
+_FA_VERSION_BY_HEAD_SIZE: dict[int, int] = {}
+
+
+def _resolve_fa_version(head_size: int) -> int:
+    # get_flash_attn_version -> current_platform.get_device_capability() is not
+    # dynamo-traceable, and the answer is fixed per head size for the process.
+    version = _FA_VERSION_BY_HEAD_SIZE.get(head_size)
+    if version is None:
+        try:
+            from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
+
+            version = int(get_flash_attn_version(requires_alibi=False, head_size=head_size) or 2)
+        except Exception:
+            version = 2
+        _FA_VERSION_BY_HEAD_SIZE[head_size] = version
+    return version
+
+
+def ar_diffusion_paged_attention(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    *,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    max_query_len: int,
+    max_seq_len: int,
+    softmax_scale: float,
+    causal: bool = False,
+) -> torch.Tensor:
+    """Run non-causal paged attention over a vLLM block table.
+
+    ``query`` may be ``(B, L, H, D)`` or already flattened as ``(T, H, D)``.
+    ``key_cache`` / ``value_cache`` are ``(num_blocks, block_size, H, D)``.
+    """
+    batched = query.dim() == 4
+    if batched:
+        batch, q_len = query.shape[:2]
+        query_flat = query.reshape(batch * q_len, *query.shape[2:])
+    else:
+        query_flat = query
+
+    if not query_flat.is_cuda:
+        out = _reference_paged_attention(
+            query_flat,
+            key_cache,
+            value_cache,
+            block_table,
+            query_start_loc,
+            seq_lens,
+            softmax_scale,
+            causal=causal,
+        )
+    else:
+        from vllm.vllm_flash_attn import flash_attn_varlen_func
+
+        fa_version = _resolve_fa_version(query_flat.shape[-1])
+
+        out = torch.empty_like(query_flat)
+        flash_attn_varlen_func(
+            q=query_flat,
+            k=key_cache,
+            v=value_cache,
+            out=out,
+            cu_seqlens_q=query_start_loc,
+            max_seqlen_q=int(max_query_len),
+            seqused_k=seq_lens,
+            max_seqlen_k=int(max_seq_len),
+            softmax_scale=float(softmax_scale),
+            causal=causal,
+            block_table=block_table,
+            fa_version=fa_version,
+        )
+
+    if batched:
+        return out.reshape(query.shape)
+    return out
+
+
+# ── Fused write+attend custom op (torch.compile-safe) ──────────────────────
+#
+# One opaque op per layer keeps the compiled DiT block fullgraph: dynamo treats
+# it as a single graph node (no eager island, no graph breaks), and the K/V slot
+# writes happen inside the op so write→read ordering with the block-table kernel
+# is internal — no hidden-mutation ordering hazards against a separate reader op
+# and no multi-GiB pool tensors as graph inputs (the flat and paged pool views
+# alias the same storage, which AOTAutograd handles poorly as inputs). Pattern
+# follows sage_attn3.py in-repo and vLLM's own unified attention ops.
+def _paged_write_attn_impl(
+    query: torch.Tensor,
+    k_curr: torch.Tensor,
+    v_curr: torch.Tensor,
+    k_act: torch.Tensor | None,
+    v_act: torch.Tensor | None,
+    layer_idx: torch.Tensor,
+    video_slots: torch.Tensor,
+    action_slots: torch.Tensor,
+    block_table: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    seq_lens: torch.Tensor,
+    max_query_len: int,
+    max_seq_len: int,
+    softmax_scale: float,
+) -> torch.Tensor:
+    kv = _CURRENT_PAGED_KV_CACHE
+    if kv is None:
+        raise RuntimeError("ar_diffusion_paged_write_attn called before prepare() set the KV pool registry")
+    layer_idx = int(layer_idx)
+    k_pool = kv._k_pools[layer_idx]
+    v_pool = kv._v_pools[layer_idx]
+    k_pool[video_slots] = k_curr.to(k_pool.dtype)
+    v_pool[video_slots] = v_curr.to(v_pool.dtype)
+    if k_act is not None and v_act is not None and k_act.shape[0] > 0:
+        k_pool[action_slots] = k_act.to(k_pool.dtype)
+        v_pool[action_slots] = v_act.to(v_pool.dtype)
+    return ar_diffusion_paged_attention(
+        query,
+        kv.key_cache(layer_idx),
+        kv.value_cache(layer_idx),
+        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=softmax_scale,
+        causal=False,
+    )
+
+
+# hasattr guard keeps registration idempotent across test re-imports that pop
+# the module from sys.modules (same as sage_attn3.py).
+if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"):
+
+    @torch.library.custom_op("vllm_omni::ar_diffusion_paged_write_attn", mutates_args=())
+    def _paged_write_attn_op(
+        query: torch.Tensor,
+        k_curr: torch.Tensor,
+        v_curr: torch.Tensor,
+        k_act: torch.Tensor | None,
+        v_act: torch.Tensor | None,
+        layer_idx: torch.Tensor,
+        video_slots: torch.Tensor,
+        action_slots: torch.Tensor,
+        block_table: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        max_query_len: int,
+        max_seq_len: int,
+        softmax_scale: float,
+    ) -> torch.Tensor:
+        return _paged_write_attn_impl(
+            query,
+            k_curr,
+            v_curr,
+            k_act,
+            v_act,
+            layer_idx,
+            video_slots,
+            action_slots,
+            block_table,
+            query_start_loc,
+            seq_lens,
+            max_query_len,
+            max_seq_len,
+            softmax_scale,
+        )
+
+    @_paged_write_attn_op.register_fake
+    def _(
+        query,
+        k_curr,
+        v_curr,
+        k_act,
+        v_act,
+        layer_idx,
+        video_slots,
+        action_slots,
+        block_table,
+        query_start_loc,
+        seq_lens,
+        max_query_len,
+        max_seq_len,
+        softmax_scale,
+    ):
+        return torch.empty_like(query)
+
+
+def paged_write_attn(
+    inputs: ARDiffusionPagedLayerInputs, query, k_curr, v_curr, k_act, v_act, softmax_scale: float
+) -> torch.Tensor:
+    """Model-facing entry: routes through the custom op (traceable in fullgraph)."""
+    return torch.ops.vllm_omni.ar_diffusion_paged_write_attn(
+        query,
+        k_curr,
+        v_curr,
+        k_act,
+        v_act,
+        inputs.layer_idx,
+        inputs.video_slots,
+        inputs.action_slots,
+        inputs.block_table,
+        inputs.query_start_loc,
+        inputs.seq_lens,
+        inputs.max_query_len,
+        inputs.max_seq_len,
+        softmax_scale,
+    )
diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
new file mode 100644
index 000000000..510bbe043
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py
@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DreamZero-facing AR-Diffusion KV bridge backed by paged self-attention."""
+
+from __future__ import annotations
+
+from vllm.logger import init_logger
+
+from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+    ARDiffusionPagedForwardContext,
+    ARDiffusionPagedLayerContext,
+)
+
+_log = init_logger(__name__)
+
+
+class ARDiffusionKVState:
+    """Per-request bridge between a DreamZero rollout and the AR-Diffusion KV pool.
+
+    The self-attention path is paged-only: ``get_kv_caches`` returns lightweight
+    per-layer contexts, and each attention layer writes current K/V into the
+    pool before calling the FlashAttention block-table kernel.  No resident
+    window is gathered into a contiguous tensor.
+    """
+
+    def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
+        """One state per DreamZero session, holding BOTH CFG branches' adapters.
+
+        ``pos``/``neg`` are the classifier-free-guidance branches: ``pos`` is
+        the conditional pass (real prompt), ``neg`` the unconditional pass
+        (negative prompt). Each owns an independent request adapter — and
+        therefore independent pool blocks — because the two branches produce
+        different K/V for the same frames. Every ``is_negative`` flag in this
+        module selects between them (``False`` -> ``pos``, ``True`` -> ``neg``).
+        """
+        self.kv_cache = kv_cache
+        self.pos = pos_adapter
+        self.neg = neg_adapter
+        self.num_layers = num_layers
+        self._committed: dict[bool, int] = {False: 0, True: 0}
+        self._paged_pending: dict[bool, ARDiffusionPagedForwardContext | None] = {False: None, True: None}
+        # Cross-attn KV is populated once after text encoding (eager), then
+        # read every denoising step. Tracked per half: the text K/V depend only
+        # on the session prompt (survive window-boundary resets), while the I2V
+        # image-token K/V come from the current observation's CLIP features and
+        # must be re-projected on every window restart.
+        self._cross_text_populated: dict[bool, bool] = {False: False, True: False}
+        self._cross_img_populated: dict[bool, bool] = {False: False, True: False}
+
+    def _adapter(self, is_negative: bool):
+        return self.neg if is_negative else self.pos
+
+    def get_kv_caches(
+        self,
+        is_negative: bool,
+        seq_len: int | None = None,
+        commit_current: bool = False,
+    ) -> list[ARDiffusionPagedLayerContext]:
+        if seq_len is None:
+            raise ValueError("AR-Diffusion paged self-attention requires seq_len in get_kv_caches()")
+        return self.prepare_paged_context(is_negative, seq_len, commit_current)
+
+    def prepare_paged_context(
+        self,
+        is_negative: bool,
+        seq_len: int,
+        commit_current: bool,
+    ) -> list[ARDiffusionPagedLayerContext]:
+        """Return per-layer paged attention contexts for one branch forward.
+
+        Allocation is lazy: constructing both CFG branches' kwargs must not
+        allocate manager blocks for a branch that a CFG-parallel rank will not run.
+        The first self-attn layer that consumes the context allocates the current
+        video slots, then all layers share that branch-level state.
+        """
+        cs = self.kv_cache.spec.chunk_size
+        if int(seq_len) % cs != 0:
+            raise AssertionError(
+                f"AR-Diffusion expects frame-aligned seq_len (multiple of chunk_size={cs}), got {seq_len}"
+            )
+
+        pending = self._paged_pending.get(is_negative)
+        if pending is not None and pending.commit_current and pending._allocated_video and not pending._committed:
+            raise RuntimeError("AR-Diffusion paged context replaced before its managed current chunk was committed")
+
+        adapter = self._adapter(is_negative)
+        branch = "neg" if is_negative else "pos"
+        forward_ctx = ARDiffusionPagedForwardContext(
+            kv_cache=self.kv_cache,
+            adapter=adapter,
+            is_negative=is_negative,
+            history_block_ids=self.kv_cache.window_block_ids(adapter),
+            seq_len=int(seq_len),
+            commit_current=bool(commit_current),
+            max_video_tokens=int(self.kv_cache.spec.sliding_window),
+        )
+        self._paged_pending[is_negative] = forward_ctx
+        _log.debug(
+            "AR-Diffusion GET   [%s] source=paged-attn layers=%d history_blocks=%d seq_len=%d commit_current=%s",
+            branch,
+            self.num_layers,
+            len(forward_ctx.history_block_ids),
+            int(seq_len),
+            bool(commit_current),
+        )
+        return [ARDiffusionPagedLayerContext(layer_idx=i, forward_ctx=forward_ctx) for i in range(self.num_layers)]
+
+    def commit_paged_context(self, is_negative: bool) -> None:
+        """Commit the managed current video blocks after a successful forward."""
+        ctx = self._paged_pending.get(is_negative)
+        if ctx is None:
+            return
+        branch = "neg" if is_negative else "pos"
+        if ctx.commit_current and ctx._allocated_video:
+            n_chunks = ctx.seq_len // self.kv_cache.spec.chunk_size
+            for _ in range(n_chunks):
+                ctx.adapter.on_chunk_committed()
+            self._committed[is_negative] += ctx.seq_len
+            _log.debug(
+                "AR-Diffusion COMMIT [%s] paged-attn new_tokens=%d chunks=%d resident=%d/%d",
+                branch,
+                ctx.seq_len,
+                n_chunks,
+                len(self.kv_cache.window_block_ids(ctx.adapter)),
+                self.kv_cache.spec.window_chunks,
+            )
+        ctx.mark_committed()
+        self._paged_pending[is_negative] = None
+
+    def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
+        """Return pool-backed cross-attn cache dicts.
+
+        Under AR-Diffusion the cross-attn pool is always populated by ``_kv_populate_cross``
+        before the first read, so this must never be reached unpopulated — if it
+        were, the model would lazily project and write cross KV itself, violating
+        engine ownership. Fail loud rather than fall back to the model.
+        """
+        img_ok = not self.kv_cache._cross_k_img or self._cross_img_populated[is_negative]
+        if not (self._cross_text_populated[is_negative] and img_ok and self.kv_cache.cross_attn_length > 0):
+            raise RuntimeError(
+                f"AR-Diffusion cross-attn read before _kv_populate_cross (neg={is_negative}, "
+                f"cross_attn_length={self.kv_cache.cross_attn_length}, "
+                f"text={self._cross_text_populated[is_negative]} "
+                f"img={self._cross_img_populated[is_negative]}) — the engine must own all cross KV"
+            )
+        return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
+
+    def close(self) -> None:
+        """Final teardown: free both branches' resident pool blocks."""
+        self.kv_cache.end_request(self.pos)
+        self.kv_cache.end_request(self.neg)
+        self._paged_pending = {False: None, True: None}
+
+    def reset(self, *, keep_cross_text: bool = False) -> None:
+        """Drop this session's KV — mirrors the model-local ``state.reset()``.
+
+        DreamZero resets at the attention-window boundary (``should_reset``); the
+        model starts a fresh sliding window, so AR-Diffusion must free the resident pool
+        blocks and start a fresh adapter for each branch. The ``ARDiffusionKVState``
+        object (and thus ``pipeline._ar_diffusion_kv_state``) is preserved across the reset
+        so the runner's session mapping stays valid — only the pool-backed state
+        is recycled.
+
+        Args:
+            keep_cross_text: On window ("inference") resets the session prompt is
+                unchanged, so the text cross-attn K/V in the static pool buffers are
+                still valid — keep them and only force the image half to repopulate.
+                Session resets (prompt change / explicit reset) pass ``False``.
+        """
+        pos_id, neg_id = self.pos.request_id, self.neg.request_id
+        self.close()
+        self.pos = self.kv_cache.begin_request(pos_id)
+        self.neg = self.kv_cache.begin_request(neg_id)
+        self._committed = {False: 0, True: 0}
+        self._paged_pending = {False: None, True: None}
+        if not keep_cross_text:
+            self._cross_text_populated = {False: False, True: False}
+        self._cross_img_populated = {False: False, True: False}
+        _log.info(
+            "AR-Diffusion RESET [%s/%s] session KV cleared (window boundary; cross text %s)",
+            pos_id,
+            neg_id,
+            "kept" if keep_cross_text else "cleared",
+        )
diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py
new file mode 100644
index 000000000..862b3d9e8
--- /dev/null
+++ b/vllm_omni/experimental/ar_diffusion/runner.py
@@ -0,0 +1,409 @@
+# SPDX-License-Identifier: Apache-2.0
+"""ARDiffusionModelRunner — the diffusion model runner for the AR-Diffusion engine.
+
+Subclasses ``DiffusionModelRunner`` and owns the engine-level KV cache
+(:class:`ARDiffusionKVCache`). It brackets a request's rollout with the KV lifecycle
+(``begin_request`` / ``end_request``) and exposes the live ``ARDiffusionKVCache`` so the
+model pipeline (DreamZero) can drive the per-chunk allocate / slot-mapping /
+gather / commit operations during ``pipeline.forward``.
+
+Selecting the AR-Diffusion engine enables KV management (no env gate); optional KV
+parameters come from the deploy yaml's ``model_config["ar_diffusion_kv_config"]``.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import math
+import time
+from collections import OrderedDict
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import MAX_DREAMZERO_SESSIONS
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig
+from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache
+from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
+
+logger = init_logger(__name__)
+
+
+def resolve_ar_diffusion_kv_config(od_config: OmniDiffusionConfig) -> ARDiffusionKVConfig:
+    """Resolve the AR-Diffusion KV config from the deploy/command arg.
+
+    Selecting the AR-Diffusion engine implies KV management, so the result is always
+    enabled. Optional overrides (``window_chunks`` / ``gpu_memory_fraction`` /
+    ``sink_chunks`` / ``reset_at_boundary``) may be supplied via the deploy yaml as
+    ``model_config["ar_diffusion_kv_config"]``. ``chunk_size`` / ``window_chunks`` are
+    finalized from the model geometry at load (see ``_preallocate_kv_cache``).
+    """
+    raw = getattr(od_config, "ar_diffusion_kv_config", None)
+    if raw is None:
+        model_config = getattr(od_config, "model_config", None)
+        if isinstance(model_config, dict):
+            raw = model_config.get("ar_diffusion_kv_config")
+    if isinstance(raw, ARDiffusionKVConfig):
+        return dataclasses.replace(raw, enable=True)
+    if isinstance(raw, dict):
+        return ARDiffusionKVConfig(**{**raw, "enable": True})
+    return ARDiffusionKVConfig(enable=True)
+
+
+class ARDiffusionModelRunner(DiffusionModelRunner):
+    def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None:
+        super().__init__(vllm_config, od_config, device)
+        self.ar_diffusion_kv_config = resolve_ar_diffusion_kv_config(od_config)
+        # Built after the model is loaded (dimensions known); stays None while KV
+        # management is disabled.
+        self.kv_cache: ARDiffusionKVCache | None = None
+        # DreamZero KV is session-scoped (persists across a session's forwards),
+        # so AR-Diffusion KV state is keyed by session_id and reused, not created per request.
+        # Bounded by the same LRU cap as DreamZeroPipeline._states: an evicted
+        # session's ARDiffusionKVState owns pool blocks (two CFG adapters), so without a
+        # bound, session-id churn in a long-running server would leak pool ownership
+        # until the KV pool is exhausted. Evicting frees those blocks (state.close()).
+        self._ar_diffusion_states: OrderedDict[str, ARDiffusionKVState] = OrderedDict()
+        self._max_ar_diffusion_states = MAX_DREAMZERO_SESSIONS
+        # Per-request server-side forward E2E times (seconds), recorded in
+        # execute_model. This is the worker-side compute time — the analog of the
+        # upstream server's per-request E2E — and excludes the engine<->worker IPC
+        # that the client's omni.generate wall time includes. Read (and cleared)
+        # by the ar_diffusion_perf_stats worker RPC for the perf-compare summary.
+        self._perf_e2e_times: list[float] = []
+
+    def build_kv_cache(
+        self,
+        *,
+        num_layers: int,
+        num_kv_heads: int,
+        head_size: int,
+        dtype,
+        block_size: int,
+        max_model_len: int,
+        available_bytes: int,
+        cross_attn_length: int = 0,
+        cross_attn_img_length: int = 0,
+        local_branches: int = 2,
+        num_frame_per_block: int = 1,
+    ) -> None:
+        """Construct the ARDiffusionKVCache (preallocating GPU pools on ``self.device``).
+
+        No-op when KV management is disabled.
+        """
+        if not self.ar_diffusion_kv_config.enable:
+            return
+        self.kv_cache = ARDiffusionKVCache(
+            self.ar_diffusion_kv_config,
+            num_layers=num_layers,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=dtype,
+            block_size=block_size,
+            max_model_len=max_model_len,
+            available_bytes=available_bytes,
+            cross_attn_length=cross_attn_length,
+            cross_attn_img_length=cross_attn_img_length,
+            local_branches=local_branches,
+            num_frame_per_block=num_frame_per_block,
+            device=self.device,
+        )
+        logger.info(
+            "AR-Diffusion KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s, "
+            "layers=%d kv_heads=%d head_dim=%d block_size=%d cross_attn_len=%d device=%s",
+            self.kv_cache.num_blocks,
+            self.ar_diffusion_kv_config.chunk_size,
+            self.ar_diffusion_kv_config.window_chunks,
+            num_layers,
+            num_kv_heads,
+            head_size,
+            block_size,
+            cross_attn_length,
+            self.device,
+        )
+
+    # -- preallocation at load -------------------------------------------------
+
+    def load_model(self, *args, **kwargs):
+        super().load_model(*args, **kwargs)
+        if self.ar_diffusion_kv_config.enable and self.pipeline is not None:
+            self._preallocate_kv_cache()
+            # Pre-capture the CUDA graphs for every window-fill shape at load time
+            # (only when compiling/cuda-graph is on) so serving is fast from chunk 0.
+            if not self.od_config.enforce_eager and self.ar_diffusion_kv_config.warmup_cudagraph:
+                self._warmup_ar_rollout()
+
+    def _infer_frame_seqlen(self) -> int:
+        """frame_seqlen = (H//8)*(W//8)//4 from the configured image_resolution."""
+        mc = getattr(self.od_config, "model_config", None) or {}
+        psc = (mc.get("policy_server_config") if isinstance(mc, dict) else None) or {}
+        res = psc.get("image_resolution", [180, 320])
+        h, w = int(res[0]), int(res[1])
+        return (h // 8) * (w // 8) // 4
+
+    def _preallocate_kv_cache(self) -> None:
+        """Build the KV pool once, from the loaded DreamZero transformer geometry."""
+        # The paged self-attention path carries one rollout per request: slot
+        # mappings/adapters are single-sequence and K/V writes address batch
+        # index 0 only. Reject a multi-sequence config at init instead of
+        # failing (or silently dropping sequences) mid-forward.
+        max_num_seqs = int(getattr(self.od_config, "max_num_seqs", 1) or 1)
+        if max_num_seqs > 1:
+            raise ValueError(
+                "AR-Diffusion paged KV supports max_num_seqs=1 (single-sequence "
+                f"rollouts); got max_num_seqs={max_num_seqs}. Per-sequence "
+                "adapters/slot maps are not implemented."
+            )
+        t = self.pipeline.transformer
+        num_layers = int(t.num_layers)
+        num_kv_heads = int(getattr(t.blocks[0].self_attn, "tp_num_heads", t.num_heads))
+        head_size = int(t.dim // t.num_heads)
+        num_frame_per_block = int(t.num_frame_per_block)
+        local_attn_size = int(t.local_attn_size)
+        # The loaded transformer is the source of truth for frame_seqlen (it derives
+        # max_attention_size from it). The deploy-config inference is only a
+        # cross-check — warn on drift rather than silently sizing chunks wrong.
+        frame_seqlen = int(getattr(t, "frame_seqlen", 0)) or self._infer_frame_seqlen()
+        inferred = self._infer_frame_seqlen()
+        if frame_seqlen != inferred:
+            logger.warning(
+                "AR-Diffusion frame_seqlen mismatch: transformer=%d deploy-config=%d; using transformer",
+                frame_seqlen,
+                inferred,
+            )
+        # The model's own attention window, in tokens (= the [-max_attention_size:]
+        # slice it applies). Read it directly rather than recomputing.
+        max_attention_size = int(t.blocks[0].self_attn.max_attention_size)
+        # Cross-attn text length — static pool for the text-encoder output.
+        # The cached k/v spans the full text sequence (the image tokens
+        # prepended in _forward_blocks are stripped inside the cross-attn
+        # forward but are not part of the cached k/v).
+        cross_attn_length = int(getattr(t, "text_len", 0))
+        # I2V image-token cross-attn length: the model splits context[:, :257] as the
+        # image tokens (the img_emb(clip_feature) output), caching k_img/v_img. T2V: 0.
+        cross_attn_img_length = 257 if getattr(t, "model_type", "t2v") == "i2v" else 0
+
+        # Frame-granular paging: 1 block = 1 frame = frame_seqlen tokens, so the
+        # resident window matches max_attention_size exactly (it need not be a whole
+        # number of num_frame_per_block causal blocks).
+        chunk_size = frame_seqlen
+        window_chunks = self.ar_diffusion_kv_config.window_chunks or (max_attention_size // frame_seqlen)
+
+        self.ar_diffusion_kv_config = dataclasses.replace(
+            self.ar_diffusion_kv_config, chunk_size=chunk_size, window_chunks=window_chunks
+        )
+        free_bytes = torch.cuda.mem_get_info(self.device)[0]
+        # Under CFG-parallel each rank executes exactly ONE branch (rank0 pos,
+        # rank1 neg; the other branch's lazy contexts never allocate on this
+        # rank), so its pool only needs one branch's capacity. A single-process
+        # run (cfg world 1) executes both branches from one pool.
+        try:
+            from vllm_omni.diffusion.distributed.parallel_state import (
+                get_classifier_free_guidance_world_size,
+            )
+
+            cfg_world = int(get_classifier_free_guidance_world_size())
+        except Exception:
+            cfg_world = 1
+        local_branches = 1 if cfg_world >= 2 else 2
+        logger.info(
+            "AR-Diffusion preallocating (paged): frame_seqlen=%d num_frame_per_block=%d "
+            "local_attn_size=%d -> chunk_size=%d window_chunks=%d (window=%d tokens)",
+            frame_seqlen,
+            num_frame_per_block,
+            local_attn_size,
+            chunk_size,
+            window_chunks,
+            window_chunks * chunk_size,
+        )
+        self.build_kv_cache(
+            num_layers=num_layers,
+            num_kv_heads=num_kv_heads,
+            head_size=head_size,
+            dtype=self.od_config.dtype,
+            block_size=chunk_size,  # one pool block per chunk
+            max_model_len=1 << 20,
+            available_bytes=free_bytes,
+            cross_attn_length=cross_attn_length,
+            cross_attn_img_length=cross_attn_img_length,
+            local_branches=local_branches,
+            num_frame_per_block=num_frame_per_block,
+        )
+
+    def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
+        # KV disabled -> base behavior, unchanged.
+        if self.kv_cache is None:
+            return super().execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs)
+
+        kv = self.kv_cache
+        # DreamZero KV is session-scoped (the model state persists across a
+        # session's forwards), so the AR-Diffusion KV state is keyed by session_id and
+        # reused — matching how pipeline.forward resolves the model-local state.
+        extra_args = req.sampling_params.extra_args or {}
+        session_id = str(extra_args.get("session_id") or "default")
+        state = self._ar_diffusion_states.get(session_id)
+        if state is None:
+            pos = kv.begin_request(f"bde__{session_id}")
+            neg = kv.begin_request(f"bde__{session_id}__neg")
+            state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers)
+            self._ar_diffusion_states[session_id] = state
+            # Evict the least-recently-used session(s) past the cap, freeing their
+            # pool blocks — mirrors DreamZeroPipeline._states so the AR-Diffusion pool can't
+            # outlive the model-local state map it shadows.
+            while len(self._ar_diffusion_states) > self._max_ar_diffusion_states:
+                old_id, old_state = self._ar_diffusion_states.popitem(last=False)
+                old_state.close()
+                logger.debug("AR-Diffusion evicted session=%s (LRU); freed pool blocks", old_id)
+        # Track recency on every access (hit or miss) so the LRU order is correct.
+        self._ar_diffusion_states.move_to_end(session_id)
+        logger.debug(
+            "AR-Diffusion execute_model: req=%s session=%s chunk_size=%d window_chunks=%s num_blocks=%d",
+            req.request_id,
+            session_id,
+            kv.spec.chunk_size,
+            kv.config.window_chunks,
+            kv.num_blocks,
+        )
+        self.pipeline._ar_diffusion_kv_state = state
+        # Time the worker-side forward (true per-request compute E2E). The CUDA
+        # sync below makes the stop reflect completed GPU work, not just dispatch.
+        _e2e_t0 = time.perf_counter()
+        try:
+            out = super().execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs)
+        except Exception:
+            # Transactional containment: a forward that died partway may have
+            # written some layers' K/V into allocated-but-uncommitted blocks
+            # and advanced model-local state. Tear the whole session down
+            # (engine pool blocks freed via close(); model-local state dropped)
+            # so the next request with this session id starts clean instead of
+            # tripping the pending-commit guard or reading half-written KV.
+            broken = self._ar_diffusion_states.pop(session_id, None)
+            if broken is not None:
+                broken.close()
+            states = getattr(self.pipeline, "_states", None)
+            if states is not None:
+                states.pop(session_id, None)
+            logger.warning(
+                "AR-Diffusion execute_model failed for session=%s; session state torn down "
+                "(pool blocks freed) — the next request starts a fresh session",
+                session_id,
+            )
+            raise
+        finally:
+            self.pipeline._ar_diffusion_kv_state = None
+        if self.device is not None and torch.cuda.is_available():
+            torch.accelerator.synchronize(self.device)
+        self._perf_e2e_times.append(time.perf_counter() - _e2e_t0)
+        return out
+
+    # -- cuda-graph warm-up ----------------------------------------------------
+
+    #: Dedicated session id for the warm-up rollout; cleaned from both state maps.
+    _WARMUP_SID = "__ardiffusion_warmup__"
+
+    def _synth_robot_obs(self, h: int, w: int, n_frames: int) -> dict:
+        """Zero-filled roboarena observation matching the client's obs layout.
+
+        Mirrors ``examples/.../client_schedule.make_obs_from_video`` (2 ext + wrist
+        cameras, 7/6/1 state vectors) so it survives ``_transform_robot_obs`` exactly
+        like a real request. One frame -> ``(H,W,3)``; chunk -> ``(n_frames,H,W,3)``.
+        """
+        img = np.zeros((h, w, 3), dtype=np.uint8) if n_frames == 1 else np.zeros((n_frames, h, w, 3), dtype=np.uint8)
+        return {
+            "observation/exterior_image_0_left": img,
+            "observation/exterior_image_1_left": img,
+            "observation/wrist_image_left": img,
+            "observation/joint_position": np.zeros(7, dtype=np.float32),
+            "observation/cartesian_position": np.zeros(6, dtype=np.float32),
+            "observation/gripper_position": np.zeros(1, dtype=np.float32),
+            "prompt": "warmup",
+            "session_id": self._WARMUP_SID,
+        }
+
+    def _warmup_ar_rollout(self) -> None:
+        """Drive a synthetic rollout through one window-fill so every DiT graph is
+        captured at load (off the serving hot path). Frees the synthetic session and
+        leaves the KV pool exactly as found. Never raises — falls back to lazy capture.
+        """
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        kv = self.kv_cache
+        if kv is None:
+            return
+        sid = self._WARMUP_SID
+        t = self.pipeline.transformer
+        nfpb = int(t.num_frame_per_block)
+        window_chunks = int(kv.config.window_chunks or 0)
+        if window_chunks <= 0:
+            return
+        # Distinct DiT window shapes are walked by one fill: prefill + ceil((W-1)/nfpb)
+        # chunk forwards. +1 optionally to also capture the reset-cycle forward.
+        n_forwards = 1 + math.ceil(max(0, window_chunks - 1) / nfpb)
+        if self.ar_diffusion_kv_config.warmup_capture_reset:
+            n_forwards += 1
+        res = self._image_resolution()
+        h, w = int(res[0]), int(res[1])
+
+        free_before = kv.manager.block_pool.get_num_free_blocks()
+        # The step cache is left active: the DiT input shape is per-chunk (the resident
+        # window length), and the first denoise step of every chunk always computes, so
+        # every window shape is captured regardless of how many later steps are skipped.
+        logger.info(
+            "AR-Diffusion cudagraph warm-up: up to %d synthetic forwards (window_chunks=%d, nfpb=%d)",
+            n_forwards,
+            window_chunks,
+            nfpb,
+        )
+        try:
+            for i in range(n_forwards):
+                n_frames = 1 if i == 0 else 4  # client convention: 1-frame prefill, 4-frame chunks
+                obs = self._synth_robot_obs(h, w, n_frames)
+                sp = OmniDiffusionSamplingParams(extra_args={"reset": i == 0, "session_id": sid, "robot_obs": obs})
+                req = OmniDiffusionRequest(prompts=["warmup"], sampling_params=sp, request_id=f"ardiffusion-warmup-{i}")
+                self.execute_model(req)
+                # Stop once the resident window is full — the remaining shapes are
+                # already captured (the window caps/resets through the same set).
+                ar_state = self._ar_diffusion_states.get(sid)
+                if (
+                    not self.ar_diffusion_kv_config.warmup_capture_reset
+                    and ar_state is not None
+                    and len(kv.window_block_ids(ar_state.pos)) >= window_chunks
+                ):
+                    break
+        except Exception as e:  # noqa: BLE001 — warm-up must never break model load
+            logger.warning("AR-Diffusion cudagraph warm-up failed (%s); using lazy capture.", e)
+        finally:
+            # Free the synthetic session from BOTH maps (KV pool + pipeline state).
+            ar_state = self._ar_diffusion_states.pop(sid, None)
+            if ar_state is not None:
+                try:
+                    ar_state.close()  # return KV blocks to the BlockPool
+                except Exception as e:  # noqa: BLE001
+                    logger.warning("AR-Diffusion warm-up: state.close() failed (%s)", e)
+            try:
+                self.pipeline._states.pop(sid, None)
+            except Exception:  # noqa: BLE001
+                pass
+            # Drop any per-forward timings the warm-up recorded so serving stats start clean.
+            self._perf_e2e_times.clear()
+            free_after = kv.manager.block_pool.get_num_free_blocks()
+            if free_after != free_before:
+                logger.warning(
+                    "AR-Diffusion warm-up: KV pool not restored (free %d -> %d) — possible leak",
+                    free_before,
+                    free_after,
+                )
+            else:
+                logger.info(
+                    "AR-Diffusion cudagraph warm-up complete; KV pool restored (%d free blocks)",
+                    free_after,
+                )
+
+    def _image_resolution(self) -> list[int]:
+        mc = getattr(self.od_config, "model_config", None) or {}
+        psc = (mc.get("policy_server_config") if isinstance(mc, dict) else None) or {}
+        return psc.get("image_resolution", [180, 320])
diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py
index 3b7e6e8d8..9660e4f34 100644
--- a/vllm_omni/metrics/definitions.py
+++ b/vllm_omni/metrics/definitions.py
@@ -11,11 +11,14 @@ time-bearing metrics use the ``_s`` suffix (values in seconds), counters use
 ``_total`` (auto-suffixed by the prometheus client), sizes use ``_bytes``.
 """
 
+import logging
+import os
 from typing import Any
 
 # vllm_omni: namespace for omni-specific Prometheus families, distinct from
 # the upstream vllm:* families.
 METRIC_PREFIX = "vllm_omni:"
+logger = logging.getLogger(__name__)
 
 
 # ============================================================================
@@ -280,6 +283,9 @@ def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int
 # ming_flash 16000 — these models populate multimodal_output["audio_sample_rate"]
 # at runtime so this default only kicks in when the field is missing.
 DEFAULT_AUDIO_SAMPLE_RATE = 24000
+DEFAULT_AUDIO_CHANNELS = 1
+AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE"
+AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS"
 
 _SAMPLE_RATE_KEYS = ("output_sample_rate", "audio_sample_rate", "sample_rate", "sampling_rate", "sr")
 
@@ -306,3 +312,26 @@ def resolve_audio_sample_rate(source: dict[str, Any] | Any | None) -> int:
         if value > 0:
             return value
     return DEFAULT_AUDIO_SAMPLE_RATE
+
+
+def stream_pcm_format_from_env(
+    *,
+    default_sample_rate: int = DEFAULT_AUDIO_SAMPLE_RATE,
+    default_channels: int = DEFAULT_AUDIO_CHANNELS,
+) -> tuple[int, int]:
+    """Return the sample rate and channel count for streamed raw PCM."""
+    sample_rate = default_sample_rate
+    channels = default_channels
+    raw_sr = os.environ.get(AUDIO_SAMPLE_RATE_ENV)
+    if raw_sr:
+        try:
+            sample_rate = int(raw_sr)
+        except ValueError:
+            logger.warning("Invalid %s=%r; using default %d", AUDIO_SAMPLE_RATE_ENV, raw_sr, sample_rate)
+    raw_ch = os.environ.get(AUDIO_CHANNELS_ENV)
+    if raw_ch:
+        try:
+            channels = int(raw_ch)
+        except ValueError:
+            logger.warning("Invalid %s=%r; using default %d", AUDIO_CHANNELS_ENV, raw_ch, channels)
+    return max(sample_rate, 1), max(channels, 1)
diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py
index 282abed23..a3dfd7bf1 100644
--- a/vllm_omni/model_executor/models/bagel/bagel.py
+++ b/vllm_omni/model_executor/models/bagel/bagel.py
@@ -587,9 +587,9 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration):
         input_ids: torch.Tensor | None,
         positions: torch.Tensor | None,
         inputs_embeds: torch.Tensor | None,
-        req_ids: list[str],
-        num_computed_tokens: list[int],
-        num_scheduled_tokens: list[int],
+        req_ids: Sequence[str],
+        num_computed_tokens: Sequence[int],
+        num_scheduled_tokens: Sequence[int],
         input_ids_buffer: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """Restore input_ids so _adjust_positions_for_img2img can locate
@@ -599,7 +599,7 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration):
             input_ids = input_ids_buffer
         return input_ids, positions
 
-    def flush_pending_metadata(self, req_ids: list[str]) -> None:
+    def flush_pending_metadata(self, req_ids: Sequence[str]) -> None:
         """Map pending metadata (batch order) to req_ids after forward().
 
         Guard: if a request already has metadata with ``image_shape``
diff --git a/vllm_omni/model_executor/models/common/ming/audio_vae.py b/vllm_omni/model_executor/models/common/ming/audio_vae.py
index 89520f515..aecc35da9 100644
--- a/vllm_omni/model_executor/models/common/ming/audio_vae.py
+++ b/vllm_omni/model_executor/models/common/ming/audio_vae.py
@@ -138,8 +138,8 @@ class Encoder(nn.Module):
         self.norm = nn.LayerNorm(config.hidden_size)
         self.patch_size = patch_size
         if patch_size != -1:
-            aggregator_config = _qwen2_config({**encoder_args, "num_hidden_layers": 4})
-            self.aggregator = Qwen2Model(aggregator_config)
+            config.num_hidden_layers = 4
+            self.aggregator = Qwen2Model(config)
             self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size))
             self.cls_embed.data.normal_(0, 0.02)
 
@@ -346,8 +346,9 @@ class AudioVAE(PreTrainedModel):
         h, _ = self.encoder(waveform)
         h = h.transpose(1, 2)
 
-        mean = torch.chunk(h, 2, dim=1)[0]
-        latent = mean.transpose(1, 2)
+        mean, scale = torch.chunk(h, 2, dim=1)
+        std = F.softplus(scale) + 1e-4
+        latent = (mean + std * torch.randn_like(mean)).transpose(1, 2)
         return latent, frame_num
 
     @torch.inference_mode()
diff --git a/vllm_omni/model_executor/models/common/ming/dit.py b/vllm_omni/model_executor/models/common/ming/dit.py
index 3b4ab1122..56e65f47f 100644
--- a/vllm_omni/model_executor/models/common/ming/dit.py
+++ b/vllm_omni/model_executor/models/common/ming/dit.py
@@ -6,6 +6,19 @@ from torch import nn
 from x_transformers.x_transformers import apply_rotary_pos_emb
 
 
+def _apply_rope_cached(t, cos, sin):
+    """Bit-exact RoPE for ``scale==1`` / interleaved-half.
+    Uses precomputed ``cos``/``sin`` per seq_len to skip per-step trig and cat."""
+    rot_dim = cos.shape[-1]
+    tr = t[..., :rot_dim]
+    x = tr.reshape(*tr.shape[:-1], -1, 2)
+    rot_half = torch.stack((-x[..., 1], x[..., 0]), dim=-1).reshape_as(tr)
+    out = tr * cos + rot_half * sin
+    if rot_dim < t.shape[-1]:  # partial rotary (unused by ming heads); keep general
+        out = torch.cat((out, t[..., rot_dim:]), dim=-1)
+    return out.type(t.dtype)
+
+
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-6):
         super().__init__()
@@ -66,12 +79,33 @@ class Attention(nn.Module):
         self.to_out.append(nn.Dropout(dropout))
         self.pe_attn_head = pe_attn_head
         self.attn_mask_enabled = attn_mask_enabled
+        # cos/sin cache keyed by seq_len (freqs are deterministic in seq_len).
+        self._rope_cache: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
+        # TODO: materialize fused weights eagerly outside any torch.compile
+        # scope (e.g. post-weight-load);
+        self._qkv_w: torch.Tensor | None = None
+        self._qkv_b: torch.Tensor | None = None
+
+    def _qkv(self, x):
+        if self._qkv_w is None:
+            self._qkv_w = torch.cat([self.to_q.weight, self.to_k.weight, self.to_v.weight], dim=0)
+            if self.to_q.bias is not None:
+                self._qkv_b = torch.cat([self.to_q.bias, self.to_k.bias, self.to_v.bias], dim=0)
+        qkv = F.linear(x, self._qkv_w, self._qkv_b)
+        return qkv.split(self.inner_dim, dim=-1)
+
+    def _rope_cos_sin(self, freqs):
+        seq_len = freqs.shape[-2]
+        cached = self._rope_cache.get(seq_len)
+        if cached is None:
+            f = freqs.unsqueeze(1) if freqs.ndim == 3 else freqs  # b n d -> b 1 n d for 4D q/k
+            cached = (f.cos(), f.sin())
+            self._rope_cache[seq_len] = cached
+        return cached
 
     def forward(self, x, mask=None, rope=None):
         batch_size = x.shape[0]
-        query = self.to_q(x)
-        key = self.to_k(x)
-        value = self.to_v(x)
+        query, key, value = self._qkv(x)
 
         inner_dim = key.shape[-1]
         head_dim = inner_dim // self.heads
@@ -86,14 +120,21 @@ class Attention(nn.Module):
 
         if rope is not None:
             freqs, xpos_scale = rope
-            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
-            if self.pe_attn_head is not None:
-                on = self.pe_attn_head
-                query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
-                key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
+            scale_is_one = xpos_scale is None or (not torch.is_tensor(xpos_scale) and xpos_scale == 1.0)
+            if scale_is_one and self.pe_attn_head is None:
+                # Fast path: cached, cat-free RoPE (bit-exact for scale==1, full rotary).
+                cos, sin = self._rope_cos_sin(freqs)
+                query = _apply_rope_cached(query, cos, sin)
+                key = _apply_rope_cached(key, cos, sin)
             else:
-                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
-                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+                q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+                if self.pe_attn_head is not None:
+                    on = self.pe_attn_head
+                    query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale)
+                    key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale)
+                else:
+                    query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                    key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
 
         if self.attn_mask_enabled and mask is not None:
             valid_sample_indices = mask.any(dim=1)
diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py
index eca5addba..7ac2ba123 100644
--- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py
+++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py
@@ -3,7 +3,7 @@
 Shared by Qwen3-Omni and Qwen3-TTS talker models.
 
 * SDPA attention (F.scaled_dot_product_attention) with native GQA support
-* HF-compatible numerics (float32 RMSNorm, float32 RoPE, separate linear layers)
+* HF-compatible CPU/CUDA numerics with NPU-only fused norm/RoPE fast paths
 * Per-call embedding buffer to avoid cross-request aliasing
 * Pre-allocated position_ids (read-only, safe to persist)
 * torch.compile (epilogue_fusion=False) on inner transformer by default
@@ -14,7 +14,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models.
 from __future__ import annotations
 
 import dataclasses
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 
 import torch
 import torch.nn as nn
@@ -24,37 +24,56 @@ from vllm.logger import init_logger
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+from vllm_omni.diffusion.layers.custom_op import CustomOp
 from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
+_GeneratorLike = torch.Generator | Sequence[torch.Generator | None] | None
+_UNIFORM_EPS = 1e-20
+
+if current_omni_platform.is_npu():
+    import torch_npu
+
 
 # ===================================================================
-# HF-numerics-compatible layers for code predictor
+# Portable layers for code predictor
 # ===================================================================
 #
 # These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32,
 # rotate_half RoPE) to produce outputs numerically identical to the
-# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel,
-# get_rope) introduce small precision differences that compound across
-# the autoregressive steps of the code predictor, causing severe
-# audio quality degradation.
+# HuggingFace reference on CPU/CUDA. vLLM's fused kernels (RMSNorm,
+# QKVParallel, get_rope) introduce small precision differences that compound
+# across the autoregressive steps of the code predictor, causing severe audio
+# quality degradation. The Ascend fused norm/RoPE kernels below are dispatched
+# by the current device platform.
 #
 # See: https://github.com/vllm-project/vllm-omni/issues/2274
 
 
-class _RMSNorm(nn.Module):
-    """RMSNorm matching HuggingFace's implementation exactly.
-
-    Computes variance in float32 to avoid bfloat16 precision loss.
-    """
+class _RMSNorm(CustomOp):
+    """RMSNorm with HuggingFace-compatible CPU/CUDA math and an NPU fast path."""
 
     def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
         super().__init__()
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward_npu(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, _ = torch_npu.npu_rms_norm(
+            hidden_states,
+            self.weight,
+            self.variance_epsilon,
+        )
+        return hidden_states
+
+    def forward_cuda(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(hidden_states)
+
+    def forward_xpu(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.forward_native(hidden_states)
+
+    def forward_native(self, hidden_states: torch.Tensor) -> torch.Tensor:
         input_dtype = hidden_states.dtype
         hidden_states = hidden_states.to(torch.float32)
         variance = hidden_states.pow(2).mean(-1, keepdim=True)
@@ -69,11 +88,8 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor:
     return torch.cat((-x2, x1), dim=-1)
 
 
-class _RotaryEmbedding(nn.Module):
-    """RoPE matching HuggingFace's implementation exactly.
-
-    Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False).
-    """
+class _RotaryEmbedding(CustomOp):
+    """RoPE with HuggingFace-compatible CPU/CUDA math and cached NPU tables."""
 
     def __init__(self, config) -> None:
         super().__init__()
@@ -85,8 +101,24 @@ class _RotaryEmbedding(nn.Module):
         rope_theta = getattr(config, "rope_theta", 10000.0)
         inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
         self.register_buffer("inv_freq", inv_freq, persistent=False)
+        if current_omni_platform.is_npu():
+            max_seq = int(getattr(config, "num_code_groups", 0) or 0) + 1
+            positions = torch.arange(max_seq, dtype=torch.float32)
+            freqs = torch.outer(positions, inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            self.register_buffer("cos_cached", emb.cos(), persistent=False)
+            self.register_buffer("sin_cached", emb.sin(), persistent=False)
 
-    def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward_npu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.cos_cached[position_ids].to(dtype=x.dtype), self.sin_cached[position_ids].to(dtype=x.dtype)
+
+    def forward_cuda(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.forward_native(x, position_ids)
+
+    def forward_xpu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.forward_native(x, position_ids)
+
+    def forward_native(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         # position_ids: [batch, seq_len]
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
         position_ids_expanded = position_ids[:, None, :].float()
@@ -142,7 +174,6 @@ class CodePredictorAttention(nn.Module):
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
         self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps)
         self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-
         if current_omni_platform.is_npu():
             if self.max_seq > 2048:
                 raise ValueError(
@@ -165,8 +196,6 @@ class CodePredictorAttention(nn.Module):
         bsz: int,
         seq_len: int,
     ) -> torch.Tensor:
-        import torch_npu
-
         q_f, k_f, v_f = q, k, v
         if self.is_gqa:
             k_f = (
@@ -227,10 +256,13 @@ class CodePredictorAttention(nn.Module):
         # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads
         cos = cos.unsqueeze(1)  # [batch, 1, seq_len, head_dim]
         sin = sin.unsqueeze(1)
-        q = (q * cos) + (_rotate_half(q) * sin)
-        k = (k * cos) + (_rotate_half(k) * sin)
-
-        if not current_omni_platform.is_npu():
+        if current_omni_platform.is_npu():
+            q = torch_npu.npu_rotary_mul(q, cos, sin)
+            k = torch_npu.npu_rotary_mul(k, cos, sin)
+            attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len)
+        else:
+            q = (q * cos) + (_rotate_half(q) * sin)
+            k = (k * cos) + (_rotate_half(k) * sin)
             attn_out = F.scaled_dot_product_attention(
                 q,
                 k,
@@ -239,8 +271,6 @@ class CodePredictorAttention(nn.Module):
                 is_causal=True,
                 enable_gqa=self.is_gqa,
             )
-        else:
-            attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len)
 
         attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
         return self.o_proj(attn_out)
@@ -287,10 +317,17 @@ class CodePredictorDecoderLayer(nn.Module):
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
         hidden_states = self.self_attn(hidden_states, position_embeddings)
-        hidden_states = residual + hidden_states
-
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        if current_omni_platform.is_npu():
+            hidden_states, _, residual = torch_npu.npu_add_rms_norm(
+                hidden_states,
+                residual,
+                self.post_attention_layernorm.weight,
+                self.post_attention_layernorm.variance_epsilon,
+            )
+        else:
+            hidden_states = residual + hidden_states
+            residual = hidden_states
+            hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
         return hidden_states
@@ -605,6 +642,30 @@ class CodePredictorWrapper(nn.Module):
                 values.add(parsed)
         return values
 
+    @staticmethod
+    def _normalize_generators(
+        generator: _GeneratorLike, batch_size: int
+    ) -> torch.Generator | list[torch.Generator | None] | None:
+        if generator is None or isinstance(generator, torch.Generator):
+            return generator
+
+        row_generators = list(generator)
+        if len(row_generators) != batch_size:
+            raise ValueError(f"Expected {batch_size} per-row generators, but got {len(row_generators)}.")
+        return row_generators
+
+    @classmethod
+    def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor:
+        """Sample ``logits`` via Gumbel-max with optional per-row generators."""
+        row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
+        u = torch.empty_like(logits, dtype=torch.float32)
+        if isinstance(row_generators, list):
+            for row, row_generator in enumerate(row_generators):
+                u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator)
+        else:
+            u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators)
+        return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True)
+
     def _prefix_seq_lens(self, max_seq: int) -> list[int]:
         all_seq_lens = list(range(2, max_seq))
         if not self._prefix_graph_seq_lens:
@@ -742,6 +803,27 @@ class CodePredictorWrapper(nn.Module):
     #  Forward -- re-prefill + inline sampling
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _multinomial(
+        probs: torch.Tensor,
+        generator: torch.Generator | None,
+        generators: Sequence[torch.Generator | None] | None,
+    ) -> torch.Tensor:
+        """Sample one code per row, optionally with per-row generators.
+
+        Per-row generators keep explicitly-seeded requests deterministic in a
+        multi-row batch: each row consumes draws only from its own generator,
+        so the transformer forward can stay batched (#4883).
+        """
+        if generators is None:
+            return torch.multinomial(probs, num_samples=1, generator=generator)
+        return torch.cat(
+            [
+                torch.multinomial(probs[row : row + 1], num_samples=1, generator=row_generator)
+                for row, row_generator in enumerate(generators)
+            ]
+        )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -753,9 +835,13 @@ class CodePredictorWrapper(nn.Module):
         top_k: int = 50,
         top_p: float = 1.0,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """Predict residual codebooks 1..G-1 autoregressively via re-prefill."""
         bsz = int(layer0_code.shape[0])
+        if generators is not None and len(generators) != bsz:
+            raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
+        sample_generator: _GeneratorLike = generators if generators is not None else generator
         num_groups = self._num_groups
         device = layer0_code.device
 
@@ -823,7 +909,6 @@ class CodePredictorWrapper(nn.Module):
             # Use captured device graph if available, otherwise call compiled fn.
             device_graph_entry = self._device_graphs.get(graph_key)
 
-            # Run transformer (device graph replay or compiled forward)
             if device_graph_entry is not None:
                 device_graph_entry[0].replay()
                 hidden_out = device_graph_entry[1]
@@ -832,9 +917,20 @@ class CodePredictorWrapper(nn.Module):
 
             logits = lm_heads[step - 1](hidden_out[:bsz, step, :])
 
-            # Sample next code
+            # Sample next code via Gumbel-max.
+            #
+            # ``argmax_i(logits_i + Gumbel_i)`` with
+            # ``Gumbel_i = -log(-log(u_i)), u_i ~ Uniform(0, 1)`` is
+            # distributionally identical to sampling from ``softmax(logits)``.
+            # In this file the motivations are practical rather than graph
+            # related: it is measurably cheaper than ``softmax + multinomial``
+            # on the B x 2048 shapes used here, it stays well-defined for
+            # degenerate masked rows with a surviving finite entry (and is more
+            # defensive than ``multinomial`` around fully-masked/NaN inputs),
+            # and the helper below can honor either one batch generator or one
+            # generator per seeded row.
             if stored_mode:
-                # "stored" mode: top-k -> top-p -> softmax -> multinomial
+                # "stored" mode: top-k -> top-p -> Gumbel-max
                 if s_top_k > 0:
                     topk_vals, _ = logits.topk(s_top_k, dim=-1)
                     logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf"))
@@ -845,17 +941,15 @@ class CodePredictorWrapper(nn.Module):
                     remove_mask = (cumulative_probs - sorted_probs) >= s_top_p
                     sorted_logits[remove_mask] = float("-inf")
                     logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
-                probs = F.softmax(logits, dim=-1, dtype=torch.float32)
-                code = torch.multinomial(probs, num_samples=1, generator=generator)
+                code = self._sample_codes_gumbel(logits, generator=sample_generator)
             else:
-                # "per_call" mode: temperature-scaled + top-k
+                # "per_call" mode: temperature-scaled + top-k -> Gumbel-max
                 if use_sampling:
                     scaled = logits * inv_temperature
                     if top_k > 0:
                         topk_vals, _ = scaled.topk(top_k, dim=-1)
                         scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
-                    probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
-                    code = torch.multinomial(probs, num_samples=1, generator=generator)
+                    code = self._sample_codes_gumbel(scaled, generator=sample_generator)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 
@@ -878,6 +972,18 @@ class CodePredictorWrapper(nn.Module):
     #  Weight loading
     # ------------------------------------------------------------------
 
+    def _prepare_npu_weights(self) -> None:
+        from vllm_ascend.utils import maybe_trans_nz
+
+        linear_count = 0
+        with torch.no_grad():
+            # Pack linear weights once for NPU matmul.
+            for module in self.modules():
+                if isinstance(module, nn.Linear):
+                    module.weight.data = maybe_trans_nz(module.weight.data)
+                    linear_count += 1
+        logger.info("Prepared NPU code predictor weights: linear=%d", linear_count)
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights directly (no fused projection remapping needed)."""
         loaded: set[str] = set()
@@ -904,4 +1010,7 @@ class CodePredictorWrapper(nn.Module):
             weight_loader(param, w)
             loaded.add(name)
 
+        if current_omni_platform.is_npu():
+            self._prepare_npu_weights()
+
         return loaded
diff --git a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
index 32db8ef5b..a8d51fdfe 100644
--- a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
+++ b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py
@@ -5,9 +5,8 @@
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
 Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
   * ``sync_process_input_func`` (``text2flow_token_only``) runs when
-    ``deploy.async_chunk=false``: stage 1 builds full-sequence flow input
-    from the emitted speech tokens. ``text2flow`` is the
-    ``custom_process_input_func`` variant.
+    ``deploy.async_chunk=false``: stage 1 allocates placeholder slots while
+    bulk tensors arrive via ``text2flow_full_payload`` on the connector.
   * ``async_chunk_process_next_stage_input_func`` runs when
     ``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1
     through the shared-memory connector.
@@ -47,7 +46,6 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="latent",
-            custom_process_input_func=f"{_PROC}.text2flow",
             sync_process_input_func=f"{_PROC}.text2flow_token_only",
         ),
     ),
diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py
old mode 100644
new mode 100755
index 960f6cfef..c3f5024b3
--- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py
+++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py
@@ -22,6 +22,7 @@ generation.
 
 from __future__ import annotations
 
+import json
 import os
 import re
 from typing import Any
@@ -170,14 +171,25 @@ _AUDIO_TOKENIZER_PATH_ENVS = (
 )
 
 
+def _is_higgs_audio_tokenizer_config(config_path: str) -> bool:
+    try:
+        with open(config_path, encoding="utf-8") as f:
+            config = json.load(f)
+    except (OSError, json.JSONDecodeError):
+        return False
+    return config.get("model_type") == "higgs_audio_v2_tokenizer"
+
+
 def _normalize_audio_tokenizer_dir(path: str) -> str | None:
     if not path:
         return None
     expanded = os.path.abspath(os.path.expanduser(path))
-    if os.path.isfile(os.path.join(expanded, "config.json")):
+    config_path = os.path.join(expanded, "config.json")
+    if os.path.isfile(config_path) and _is_higgs_audio_tokenizer_config(config_path):
         return expanded
     nested = os.path.join(expanded, _K2_OMNIVOICE_SUBDIR)
-    if os.path.isfile(os.path.join(nested, "config.json")):
+    nested_config_path = os.path.join(nested, "config.json")
+    if os.path.isfile(nested_config_path) and _is_higgs_audio_tokenizer_config(nested_config_path):
         return nested
     return None
 
@@ -201,7 +213,9 @@ def _resolve_audio_tokenizer_dir() -> str | None:
         filename=f"{_K2_OMNIVOICE_SUBDIR}/config.json",
     )
     if isinstance(cached_config, str) and os.path.isfile(cached_config):
-        return os.path.dirname(cached_config)
+        candidate = _normalize_audio_tokenizer_dir(os.path.dirname(cached_config))
+        if candidate is not None:
+            return candidate
 
     from huggingface_hub.constants import HF_HUB_CACHE
 
@@ -240,7 +254,9 @@ def _load_audio_tokenizer():
             _K2_OMNIVOICE_REPO,
             allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"],
         )
-        audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR)
+        audio_tokenizer_dir = _normalize_audio_tokenizer_dir(repo_path)
+        if audio_tokenizer_dir is None:
+            raise RuntimeError(f"Downloaded {_K2_OMNIVOICE_REPO} does not contain a valid Higgs audio tokenizer")
     device = current_omni_platform.get_torch_device()
     model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device)
     return model.eval()
diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
index d9efccd9a..17bbad001 100755
--- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
+++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py
@@ -222,14 +222,8 @@ class HunyuanModel(HunYuanModel):
             if self.config.tie_word_embeddings and "lm_head.weight" in name:
                 continue
 
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache scales for compressed-tensors quantization
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                continue
-
+            # KV-cache scales are renamed to .attn.{k,v,q}_scale by the
+            # quant_config mapper in the outer AutoWeightsLoader.
             is_found = False
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
diff --git a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py
index 6311216dc..2fe616cfd 100644
--- a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py
+++ b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py
@@ -483,19 +483,15 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
         ]
         params_dict = dict(self.named_parameters(remove_duplicate=False))
         loaded_params: set[str] = set()
+        if self.quant_config is not None:
+            cache_scale_mapper = self.quant_config.get_cache_scale_mapper()
+            if cache_scale_mapper is not None:
+                weights = cache_scale_mapper.apply(weights)
 
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
 
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                loaded_params.add(scale_name)
-                continue
-
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
diff --git a/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py b/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py
new file mode 100644
index 000000000..be35eb562
--- /dev/null
+++ b/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MammothModa2 pipeline topology (frozen).
+
+Stage 0: AR  — multimodal understanding + latent generation
+Stage 1: DiT — latent → image
+
+For text/image understanding tasks (text output only), use MAMMOTH_MODA2_AR_PIPELINE.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.mammoth_moda2"
+
+MAMMOTH_MODA2_PIPELINE = PipelineConfig(
+    model_type="mammoth_moda2",
+    model_arch="MammothModa2ForConditionalGeneration",
+    hf_architectures=("Mammothmoda2Model", "MammothModa2ForConditionalGeneration"),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=False,
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="dit",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            owns_tokenizer=False,
+            requires_multimodal_data=False,
+            engine_output_type="image",
+            custom_process_input_func=f"{_PROC}.ar2dit",
+        ),
+    ),
+)
+
+# Single-stage AR variant for understanding / summarization tasks
+MAMMOTH_MODA2_AR_PIPELINE = PipelineConfig(
+    model_type="mammoth_moda2_ar",
+    model_arch="MammothModa2ForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+        ),
+    ),
+)
diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
index a640ade3a..48715fb98 100644
--- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
+++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py
@@ -1153,17 +1153,14 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
         params_dict = dict(self.named_parameters(remove_duplicate=False))
 
         loaded_params: set[str] = set()
+        if self.quant_config is not None:
+            cache_scale_mapper = self.quant_config.get_cache_scale_mapper()
+            if cache_scale_mapper is not None:
+                weights = cache_scale_mapper.apply(weights)
+
         for name, loaded_weight in weights:
             if name.startswith("model."):
                 name = "model." + name
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache quantization scales
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                loaded_params.add(scale_name)
-                continue
 
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if name.startswith("input_local_transformer.") or name.startswith("local_transformer."):
diff --git a/vllm_omni/model_executor/models/mimo_audio/pipeline.py b/vllm_omni/model_executor/models/mimo_audio/pipeline.py
index 126c90176..9de6575d9 100644
--- a/vllm_omni/model_executor/models/mimo_audio/pipeline.py
+++ b/vllm_omni/model_executor/models/mimo_audio/pipeline.py
@@ -56,7 +56,6 @@ MIMO_AUDIO_PIPELINE = PipelineConfig(
             final_output=True,
             final_output_type="audio",
             engine_output_type="audio",
-            custom_process_input_func=f"{_PROC}.llm2code2wav",
             sync_process_input_func=f"{_PROC}.llm2code2wav_token_only",
             sampling_constraints={"detokenize": False},
         ),
diff --git a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
index a83989a03..fdac2a30e 100644
--- a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
+++ b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py
@@ -22,6 +22,7 @@ from .constants import (
     DEFAULT_SIGMA,
     DEFAULT_TEMPERATURE,
     HISTORY_PATCH_SIZE,
+    INITIAL_LATENT_CHUNK_SIZE,
     KEY_CFG,
     KEY_CHUNK_ID,
     KEY_DECODE_STEP,
@@ -239,6 +240,7 @@ class MingTTSConfig:
     max_decode_steps: int = MAX_DECODE_STEPS
 
     latent_chunk_size: int = LATENT_CHUNK_SIZE
+    initial_latent_chunk_size: int = INITIAL_LATENT_CHUNK_SIZE
     latent_left_context: int = LATENT_LEFT_CONTEXT
 
     text_eos_token_id: int = TEXT_EOS_TOKEN_ID
diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py
index a189e4285..e0508bda7 100644
--- a/vllm_omni/model_executor/models/ming_tts/constants.py
+++ b/vllm_omni/model_executor/models/ming_tts/constants.py
@@ -58,6 +58,7 @@ DEFAULT_TEMPERATURE = 0.0
 
 # Connector / Stage-2 streaming defaults (runtime tuning)
 LATENT_CHUNK_SIZE = 25
+INITIAL_LATENT_CHUNK_SIZE = 4
 LATENT_LEFT_CONTEXT = 0
 MAX_DECODE_STEPS = 200
 
diff --git a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
index ae4d61f68..39ddd7b9f 100644
--- a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
+++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py
@@ -326,6 +326,28 @@ class MingLLMModel(nn.Module):
         try:
             from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler
 
+            # TODO(perf):
+            #   (1) compile per-block (dit_model.blocks) instead of the
+            #   whole forward for tighter fusion / fewer graph breaks;
+            #   (2) move the compile + a synthetic warmup forward to a post-weight-load hook
+            dit_model = self.flowloss.cfm.model
+            if not getattr(dit_model, "_ming_compiled", False):
+                try:
+                    # Scoped options (avoid mutating global inductor config).
+                    dit_model.forward = torch.compile(
+                        dit_model.forward,
+                        fullgraph=False,
+                        dynamic=False,
+                        options={
+                            "triton.cudagraphs": False,
+                            "triton.cudagraph_trees": False,
+                        },
+                    )
+                    dit_model._ming_compiled = True
+                    logger.info("Ming CFM DiT torch.compile enabled (inductor cudagraphs off).")
+                except Exception as exc:
+                    logger.warning("Ming CFM torch.compile failed (%s); using uncompiled DiT.", exc)
+
             sampler = CFMSampler(self.flowloss.cfm.model, steps=_CFM_STEPS)
             self._cfm_graph = CFMGraphExecutor(
                 sampler,
diff --git a/vllm_omni/model_executor/models/ming_tts/patch_emission.py b/vllm_omni/model_executor/models/ming_tts/patch_emission.py
index c0a25a1c1..89acdf231 100644
--- a/vllm_omni/model_executor/models/ming_tts/patch_emission.py
+++ b/vllm_omni/model_executor/models/ming_tts/patch_emission.py
@@ -127,7 +127,7 @@ def _validate_ming_decode_window(
             continue
         max_steps = _resolve_runtime_int(info, KEY_MAX_DECODE_STEPS, default_max_decode_steps)
         min_steps = _resolve_optional_runtime_int(info, KEY_MIN_DECODE_STEPS, 0)
-        min_required = max(min_stop_step + 1, min_steps)
+        min_required = max(min_stop_step + 2, min_steps)
         if max_steps < min_required:
             req_id = info.get(KEY_REQUEST_ID, f"idx={i}")
             raise ValueError(
@@ -148,7 +148,7 @@ def _resolve_ming_stop_decision(
     audio_dummy_token_id: int,
     text_eos_token_id: int,
 ) -> tuple[str, bool, bool, int, int]:
-    min_required_decode_steps = max(min_stop_step + 1, min_decode_steps)
+    min_required_decode_steps = max(min_stop_step + 2, min_decode_steps)
     if max_decode_steps < min_required_decode_steps:
         raise RuntimeError(
             "Invalid Ming decode window: "
diff --git a/vllm_omni/model_executor/models/ming_tts/validation.py b/vllm_omni/model_executor/models/ming_tts/validation.py
index c4ccf9dcd..3f3da6cf0 100644
--- a/vllm_omni/model_executor/models/ming_tts/validation.py
+++ b/vllm_omni/model_executor/models/ming_tts/validation.py
@@ -172,6 +172,9 @@ def validate_ming_tts_config(cfg: Any) -> None:
 
     if cfg.latent_chunk_size <= 0:
         raise ValueError(f"latent_chunk_size must be > 0, got {cfg.latent_chunk_size}.")
+    initial_latent_chunk_size = getattr(cfg, "initial_latent_chunk_size", 0)
+    if initial_latent_chunk_size < 0:
+        raise ValueError(f"initial_latent_chunk_size must be >= 0, got {initial_latent_chunk_size}.")
     if cfg.latent_left_context < 0:
         raise ValueError(f"latent_left_context must be >= 0, got {cfg.latent_left_context}.")
     if cfg.max_decode_steps <= 0:
diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py
new file mode 100644
index 000000000..f3ba18c36
--- /dev/null
+++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py
@@ -0,0 +1,1959 @@
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch MossAudioTokenizer model."""
+
+from __future__ import annotations
+
+import copy
+import math
+from contextlib import ExitStack, contextmanager
+from dataclasses import dataclass
+from typing import cast
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.modeling_utils import PreTrainedAudioTokenizerBase
+from transformers.utils import ModelOutput, auto_docstring, logging
+
+from .configuration_moss_audio_tokenizer_v2 import MossAudioTokenizerConfig
+
+logger = logging.get_logger(__name__)
+
+
+# =============================================================================
+# Output Classes
+# =============================================================================
+
+
+@dataclass
+@auto_docstring
+class MossAudioTokenizerEncoderOutput(ModelOutput):
+    r"""
+    audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*):
+        Discrete audio codes computed using the encoder and quantizer.
+    audio_codes_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+        Valid lengths for each sample's audio codes.
+    encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, hidden_size, sequence_length)`, *optional*):
+        Hidden states from the encoder before quantization.
+    """
+
+    audio_codes: torch.Tensor | None = None
+    audio_codes_lengths: torch.Tensor | None = None
+    encoder_hidden_states: torch.Tensor | None = None
+
+
+@dataclass
+@auto_docstring
+class MossAudioTokenizerDecoderOutput(ModelOutput):
+    r"""
+    audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*):
+        Decoded audio waveform.
+    audio_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+        Valid lengths for each sample's audio.
+    """
+
+    audio: torch.Tensor | None = None
+    audio_lengths: torch.Tensor | None = None
+
+
+@dataclass
+@auto_docstring
+class MossAudioTokenizerOutput(ModelOutput):
+    r"""
+    audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*):
+        Decoded audio waveform.
+    audio_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+        Valid lengths for each sample's audio.
+    audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*):
+        Discrete audio codes computed using the encoder and quantizer.
+    audio_codes_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+        Valid lengths for each sample's audio codes.
+    """
+
+    audio: torch.Tensor | None = None
+    audio_lengths: torch.Tensor | None = None
+    audio_codes: torch.Tensor | None = None
+    audio_codes_lengths: torch.Tensor | None = None
+
+
+# =============================================================================
+# Streaming Module Base Classes
+# =============================================================================
+
+
+@dataclass
+class StreamingState:
+    """Base state for streaming modules."""
+
+    batch_size: int
+    device: torch.device
+
+    def __post_init__(self):
+        self.exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=self.device)
+
+    def set_exec_mask(self, exec_mask: torch.Tensor):
+        self.exec_mask[:] = exec_mask
+
+    def reset(self, reset_mask: torch.Tensor) -> None:
+        self.exec_mask[:] = torch.where(reset_mask, torch.ones_like(self.exec_mask), self.exec_mask)
+
+    def __enter__(self):
+        # ExitStack expects a context manager; returning self is conventional and useful for debugging.
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        pass
+
+
+class StreamingModule(nn.Module):
+    """Base class for streaming components."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._streaming_state: StreamingState | None = None
+        self._streaming_detached: bool = False
+        self._cached_children: list[tuple[str, StreamingModule]] | None = None
+
+    @property
+    def is_streaming(self):
+        return self._streaming_state is not None
+
+    def _apply_named_streaming(self, fn):
+        def _handle_module(prefix: str, module: nn.Module):
+            if isinstance(module, StreamingModule):
+                if module._streaming_detached and prefix != "":
+                    return
+                if self._cached_children is None:
+                    raise RuntimeError("Internal error: _cached_children should be initialized before traversal.")
+                self._cached_children.append((prefix, module))
+            for name, child in module.named_children():
+                new_prefix = f"{prefix}.{name}" if prefix else name
+                _handle_module(new_prefix, child)
+
+        if self._cached_children is None:
+            self._cached_children = []
+            _handle_module("", self)
+        for name, child in self._cached_children:
+            fn(name, child)
+
+    def _start_streaming(self, batch_size: int, exit_stack: ExitStack):
+        def _start_streaming_fn(name: str, module: StreamingModule):
+            if module._streaming_state is not None:
+                raise RuntimeError(f"{name} is already streaming!")
+            state = module._init_streaming_state(batch_size)
+            exit_stack.enter_context(state)
+            module._streaming_state = state
+
+        self._apply_named_streaming(_start_streaming_fn)
+
+    def _stop_streaming(self) -> None:
+        def _stop_streaming_fn(name: str, module: StreamingModule):
+            module._streaming_state = None
+
+        self._apply_named_streaming(_stop_streaming_fn)
+
+    def _init_streaming_state(self, batch_size: int) -> StreamingState:
+        device = next(iter(self.parameters())).device
+        return StreamingState(batch_size, device)
+
+    def streaming(self, batch_size: int) -> ExitStack:
+        """Context manager to enter streaming mode."""
+        exit_stack = ExitStack()
+        self._start_streaming(batch_size, exit_stack)
+        exit_stack.callback(self._stop_streaming)
+        return exit_stack
+
+    def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None:
+        """Set the active slot mask on every live streaming module."""
+
+        def _set_exec_mask(_: str, module: StreamingModule) -> None:
+            state = module._streaming_state
+            if state is not None:
+                state.set_exec_mask(exec_mask.to(state.device))
+
+        self._apply_named_streaming(_set_exec_mask)
+
+
+class StreamingContainer(StreamingModule):
+    """Container for streaming modules."""
+
+    pass
+
+
+# =============================================================================
+# Normalization Layers
+# =============================================================================
+
+
+class MossAudioTokenizerRMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization."""
+
+    def __init__(
+        self,
+        dim: int,
+        eps: float = 1e-5,
+        dtype: torch.dtype | None = None,
+        device=None,
+    ):
+        super().__init__()
+        self.eps = eps
+        self.dtype = dtype
+        self.alpha = nn.Parameter(torch.full((1, 1, dim), 1.0, requires_grad=True, device=device, dtype=dtype))
+
+    def forward(self, x: torch.Tensor):
+        x_dtype = x.dtype
+        if self.dtype is not None:
+            x = x.to(self.dtype)
+        var = self.eps + torch.mean(x**2, dim=2, keepdim=True)
+        y = (x * (self.alpha.to(var) * torch.rsqrt(var))).to(x_dtype)
+        return y
+
+
+class MossAudioTokenizerLayerScale(nn.Module):
+    """Layer scale from Touvron et al. 2021."""
+
+    def __init__(
+        self,
+        channels: int,
+        init: float = 1e-4,
+        channel_last: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.channel_last = channel_last
+        self.scale = nn.Parameter(torch.full((channels,), init, requires_grad=True, device=device, dtype=dtype))
+
+    def forward(self, x: torch.Tensor):
+        if self.channel_last:
+            return self.scale * x
+        else:
+            return self.scale[:, None] * x
+
+
+def create_norm_fn(norm_type: str, dim: int, **kwargs) -> nn.Module:
+    """Create normalization module."""
+    if norm_type == "layer_norm":
+        return nn.LayerNorm(dim, eps=1e-5, **kwargs)
+    elif norm_type in {"rms_norm"}:
+        return MossAudioTokenizerRMSNorm(dim, eps=1e-5, **kwargs)
+    elif norm_type in {"rms_norm_f32"}:
+        kwargs.pop("dtype", None)
+        return MossAudioTokenizerRMSNorm(dim, eps=1e-8, dtype=torch.float, **kwargs)
+    else:
+        raise ValueError(f"Unknown norm type: {norm_type}")
+
+
+# =============================================================================
+# Rotary Position Embedding
+# =============================================================================
+
+
+def apply_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    offset: torch.Tensor,
+    max_period: float = 10_000,
+    time_before_heads: bool = False,
+):
+    """Apply rotary position embedding."""
+    if time_before_heads:
+        B, T, H, D = q.shape
+    else:
+        B, H, T, D = q.shape
+    if k.shape != q.shape:
+        raise ValueError(f"Expected k.shape == q.shape, got k={tuple(k.shape)} q={tuple(q.shape)}")
+    if D <= 0 or (D % 2) != 0:
+        raise ValueError(f"RoPE requires an even last dimension, got D={D}")
+
+    ds = torch.arange(D // 2, device=q.device, dtype=torch.float32)
+    freqs = torch.exp(ds * (-math.log(max_period) * 2 / D))
+    ts = offset.float().view(-1, 1) + torch.arange(T, device=q.device, dtype=torch.float32)
+
+    if time_before_heads:
+        ts = ts.view(B, -1, 1, 1)
+    else:
+        ts = ts.view(B, 1, -1, 1)
+
+    dims = q.shape[:-1]
+    q = q.view(*dims, D // 2, 2)
+    k = k.view(*dims, D // 2, 2)
+
+    qr, qi = q[..., 0].float(), q[..., 1].float()
+    kr, ki = k[..., 0].float(), k[..., 1].float()
+
+    rotr = torch.cos(freqs * ts)
+    roti = torch.sin(freqs * ts)
+
+    qor = qr * rotr - qi * roti
+    qoi = qr * roti + qi * rotr
+    kor = kr * rotr - ki * roti
+    koi = kr * roti + ki * rotr
+
+    dtype = q.dtype
+    qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1)
+    ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1)
+
+    return qo.view(*dims, D), ko.view(*dims, D)
+
+
+class MossAudioTokenizerRotaryEmbedding(nn.Module):
+    """Rotary positional embedding (RoPE)."""
+
+    def __init__(self, max_period: float = 10000.0):
+        super().__init__()
+        self.max_period = max_period
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        offset: torch.Tensor,
+        time_before_heads: bool = False,
+    ):
+        return apply_rope(q, k, offset, self.max_period, time_before_heads)
+
+
+# =============================================================================
+# Gating Modules
+# =============================================================================
+
+
+class MossAudioTokenizerActivationGating(nn.Module):
+    """Gating FFN layer with activation."""
+
+    def __init__(self, dim: int, dim_feedforward: int, activation, **factory_kwargs):
+        super().__init__()
+        if dim_feedforward == 4 * dim:
+            hidden = (21 * dim) // 8
+        else:
+            hidden = (2 * dim_feedforward) // 3
+
+        self.linear_in = nn.Linear(dim, 2 * hidden, bias=False, **factory_kwargs)
+        self.linear_out = nn.Linear(hidden, dim, bias=False, **factory_kwargs)
+        self.activation = activation
+
+    def forward(self, x: torch.Tensor):
+        x = self.linear_in(x)
+        B, T, _ = x.shape
+        x = x.view(B, T, 2, -1)
+        x = self.activation(x[..., 0, :]) * x[..., 1, :]
+        x = self.linear_out(x)
+        return x
+
+
+def _get_activation(name: str):
+    if name in ["sigmoid", "tanh", "relu"]:
+        return getattr(torch, name)
+    elif name in ["leaky_relu", "elu", "gelu", "silu", "mish", "softsign"]:
+        return getattr(F, name)
+    elif name == "identity":
+        return nn.Identity()
+    else:
+        raise ValueError(f"Unknown activation {name}")
+
+
+def make_gating(name: str, dim: int, dim_feedforward: int, **factory_kwargs) -> nn.Module:
+    return MossAudioTokenizerActivationGating(dim, dim_feedforward, _get_activation(name), **factory_kwargs)
+
+
+# =============================================================================
+# Positional Embeddings
+# =============================================================================
+
+
+def create_sin_embedding(
+    positions: torch.Tensor,
+    dim: int,
+    max_period: float = 10000,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Create sinusoidal positional embedding with shape [B, T, C]."""
+    if dim % 2 != 0:
+        raise ValueError(f"Sinusoidal embedding requires even dim, got dim={dim}")
+    half_dim = dim // 2
+    if half_dim <= 1:
+        raise ValueError(f"Sinusoidal embedding requires dim >= 4, got dim={dim}")
+    positions = positions.to(dtype)
+    adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1)
+    max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype)
+    phase = positions / (max_period_tensor ** (adim / (half_dim - 1)))
+    return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
+
+
+# =============================================================================
+# KV Cache for Attention
+# =============================================================================
+
+
+class KVCacheResult:
+    """Container for KV cache results that supports tuple unpacking."""
+
+    __slots__ = ("keys", "values", "positions")
+
+    def __init__(self, keys: torch.Tensor, values: torch.Tensor, positions: torch.Tensor):
+        self.keys = keys
+        self.values = values
+        self.positions = positions
+
+    def __iter__(self):
+        """Allow unpacking as (keys, values, positions)."""
+        return iter((self.keys, self.values, self.positions))
+
+    @staticmethod
+    def from_kv(keys: torch.Tensor, values: torch.Tensor) -> KVCacheResult:
+        B, H, T, D = keys.shape
+        positions = torch.arange(T, device=keys.device, dtype=torch.long)
+        return KVCacheResult(keys, values, positions.expand(B, -1))
+
+
+class RingKVCache:
+    """Efficient streaming KVCache compatible with CUDA Graph."""
+
+    def __init__(
+        self,
+        batch_size: int,
+        num_heads: int,
+        dim_per_head: int,
+        capacity: int,
+        respect_exec_mask: bool = True,
+        device: torch.device = torch.device("cuda"),
+        dtype: torch.dtype = torch.bfloat16,
+    ):
+        self.capacity = capacity
+        self.cache = torch.zeros(
+            (2, batch_size, num_heads, capacity, dim_per_head),
+            device=device,
+            dtype=dtype,
+        )
+        self.respect_exec_mask = respect_exec_mask
+        if self.respect_exec_mask:
+            self.end_offset = torch.zeros(batch_size, device=device, dtype=torch.long)
+        else:
+            self.end_offset = torch.zeros(1, device=device, dtype=torch.long)
+
+    def reset(self, reset_mask: torch.Tensor) -> None:
+        self.end_offset[:] = torch.where(reset_mask, torch.zeros_like(self.end_offset), self.end_offset)
+
+    def complete(self, k: torch.Tensor, v: torch.Tensor, exec_mask: torch.Tensor) -> KVCacheResult:
+        B, H, T, D = k.shape
+        if T <= 0:
+            raise ValueError(f"Expected T > 0, got T={T}")
+
+        indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype)
+        indexes = indexes + self.end_offset.view(-1, 1)
+        indexes = indexes % self.capacity
+
+        if self.respect_exec_mask:
+            this_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D)
+            self.cache[0].scatter_(2, this_indexes, k)
+            self.cache[1].scatter_(2, this_indexes, v)
+        else:
+            self.cache[0].index_copy_(2, indexes[0], k)
+            self.cache[1].index_copy_(2, indexes[0], v)
+
+        keys = self.cache[0]
+        values = self.cache[1]
+
+        indexes = torch.arange(self.capacity, device=self.end_offset.device, dtype=torch.long)
+        last_offset = self.end_offset.view(-1, 1) + T - 1
+        end_index = last_offset % self.capacity
+        delta = indexes - end_index
+
+        positions = torch.where(
+            delta <= 0,
+            last_offset + delta,
+            last_offset + delta - self.capacity,
+        )
+
+        if self.respect_exec_mask:
+            self.end_offset[:] = torch.where(exec_mask, self.end_offset + T, self.end_offset)
+        else:
+            self.end_offset.add_(T)
+
+        invalid = indexes >= self.end_offset.view(-1, 1)
+        positions = torch.where(invalid, torch.full_like(positions, -1), positions)
+
+        return KVCacheResult(keys, values, positions)
+
+
+# =============================================================================
+# Multi-Head Attention
+# =============================================================================
+
+
+@dataclass
+class MHAState(StreamingState):
+    kv_cache: RingKVCache | None
+    offset: torch.Tensor
+    offset_cpu: int
+
+    def reset(self, reset_mask: torch.Tensor):
+        super().reset(reset_mask)
+        self.offset[:] = torch.where(reset_mask, torch.zeros_like(self.offset), self.offset)
+        if self.kv_cache is not None:
+            self.kv_cache.reset(reset_mask)
+        self.offset_cpu = 0
+
+
+def apply_weights_per_step(
+    modules: nn.ModuleList,
+    schedule: list[int] | None,
+    x: torch.Tensor,
+    offset: int | None,
+) -> torch.Tensor:
+    """Apply different weights for each time step."""
+    if len(modules) == 1:
+        return modules[0](x)
+
+    if offset is None:
+        raise ValueError("offset must be provided when using per-step weights (len(modules) > 1).")
+    ys = []
+    B, T, C = x.shape
+    for t in range(T):
+        module_index = t + offset
+        if schedule is not None:
+            if module_index >= len(schedule) or module_index < 0:
+                raise ValueError(
+                    f"weights_per_step_schedule is too short for module_index={module_index} (len={len(schedule)})."
+                )
+            module_index = schedule[module_index]
+        if module_index >= len(modules) or module_index < 0:
+            raise ValueError(f"module_index={module_index} out of range for len(modules)={len(modules)}.")
+        y = modules[module_index](x[:, t : t + 1])
+        ys.append(y)
+    return torch.cat(ys, 1)
+
+
+class MossAudioTokenizerMultiheadAttention(StreamingModule):
+    """Multi-head attention with streaming support."""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        causal: bool = False,
+        context: int | None = None,
+        rope: MossAudioTokenizerRotaryEmbedding | None = None,
+        weights_per_step: int = 0,
+        weights_per_step_schedule: list[int] | None = None,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        self.embed_dim = embed_dim
+        self.causal = causal
+        self.context = context
+        self.rope = rope
+        self.num_heads = num_heads
+        self.weights_per_step = weights_per_step
+        self.weights_per_step_schedule = weights_per_step_schedule
+
+        out_dim = 3 * embed_dim
+        mult = 1
+        if weights_per_step:
+            mult = max(weights_per_step_schedule) + 1 if weights_per_step_schedule else weights_per_step
+        self.mult = mult
+
+        self.out_projs = nn.ModuleList(
+            [nn.Linear(embed_dim, embed_dim, bias=False, **factory_kwargs) for _ in range(mult)]
+        )
+        self.in_projs = nn.ModuleList(
+            [nn.Linear(embed_dim, out_dim, bias=False, **factory_kwargs) for _ in range(mult)]
+        )
+
+        self._register_load_state_dict_pre_hook(self._load_hook, with_module=True)
+
+    @staticmethod
+    def _load_hook(module, state_dict, prefix, *_):
+        mappings = {
+            "in_proj_weight": "in_projs.{i}.weight",
+            "in_proj.weight": "in_projs.{i}.weight",
+            "out_proj.weight": "out_projs.{i}.weight",
+        }
+        mult = module.mult
+        for suffix in ["", "_scb"]:
+            for source, target in mappings.items():
+                this_source = prefix + source + suffix
+                if this_source in state_dict:
+                    weight = state_dict[this_source]
+                    _, *OD = weight.shape
+                    weight = weight.view(mult, -1, *OD)
+                    for i in range(mult):
+                        state_dict[prefix + target.format(i=i) + suffix] = weight[i]
+                    state_dict.pop(this_source)
+
+    def _init_streaming_state(self, batch_size: int) -> MHAState:
+        in_proj = cast(nn.Linear, self.in_projs[0])
+        device = cast(torch.device, in_proj.weight.device)
+        dtype = cast(torch.dtype, in_proj.weight.dtype)
+
+        dim_per_head = self.embed_dim // self.num_heads
+        if self.context is None:
+            capacity = self.weights_per_step if self.weights_per_step else 1024
+        else:
+            capacity = self.context
+
+        kv_cache = RingKVCache(
+            batch_size,
+            self.num_heads,
+            dim_per_head,
+            capacity,
+            respect_exec_mask=not self.weights_per_step,
+            device=cast(torch.device, device),
+            dtype=cast(torch.dtype, dtype),
+        )
+        return MHAState(
+            batch_size,
+            cast(torch.device, device),
+            kv_cache,
+            offset=torch.zeros(batch_size, device=cast(torch.device, device), dtype=torch.long),
+            offset_cpu=0,
+        )
+
+    def _complete_kv(self, k, v) -> KVCacheResult:
+        state = cast(MHAState | None, self._streaming_state)
+        if state is None:
+            return KVCacheResult.from_kv(k, v)
+        if state.kv_cache is None:
+            return KVCacheResult.from_kv(k, v)
+        return state.kv_cache.complete(k, v, state.exec_mask)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
+        state = cast(MHAState | None, self._streaming_state)
+        B, T = query.shape[:2]
+
+        if state is None:
+            offset = torch.zeros(B, device=query.device, dtype=torch.long)
+            offset_cpu = 0
+        else:
+            offset = state.offset
+            offset_cpu = state.offset_cpu
+
+        projected = apply_weights_per_step(self.in_projs, self.weights_per_step_schedule, query, offset_cpu)
+        dim_per_head = self.embed_dim // self.num_heads
+        projected = projected.reshape(B, T, 3, self.num_heads, dim_per_head).permute(2, 0, 3, 1, 4)
+        q, k, v = projected[0], projected[1], projected[2]
+
+        if self.rope:
+            q, k = self.rope(q, k, offset, time_before_heads=False)
+
+        k, v, pos_k = self._complete_kv(k, v)
+        pos_k = pos_k[:, None]
+
+        if self.causal:
+            pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1)
+            delta = pos_q - pos_k
+            attn_bias = (pos_k >= 0) & (delta >= 0)
+            if self.context is not None:
+                attn_bias = attn_bias & (delta < self.context)
+            attn_bias = attn_bias[:, None]
+        else:
+            attn_bias = None
+
+        x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
+        x = x.transpose(1, 2).reshape(B, T, self.embed_dim)
+        x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu)
+
+        if state is not None:
+            state.offset[:] = torch.where(state.exec_mask, state.offset + T, state.offset)
+            state.offset_cpu += T
+        return x
+
+
+# =============================================================================
+# Transformer Layer
+# =============================================================================
+
+
+@dataclass
+class LayerState(StreamingState):
+    offset_cpu: int = 0
+
+    def reset(self, reset_mask: torch.Tensor):
+        super().reset(reset_mask)
+        self.offset_cpu = 0
+
+
+class MossAudioTokenizerTransformerLayer(StreamingModule):
+    """Transformer layer with streaming support."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dim_feedforward: int = 2048,
+        causal: bool = False,
+        context: int | None = None,
+        rope: MossAudioTokenizerRotaryEmbedding | None = None,
+        norm: str = "layer_norm",
+        layer_scale: float | None = None,
+        gating: str = "none",
+        weights_per_step: int = 0,
+        weights_per_step_schedule: list[int] | None = None,
+        activation=F.gelu,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+
+        self.self_attn = MossAudioTokenizerMultiheadAttention(
+            embed_dim=d_model,
+            num_heads=num_heads,
+            causal=causal,
+            context=context,
+            rope=rope,
+            weights_per_step=weights_per_step,
+            weights_per_step_schedule=weights_per_step_schedule,
+            **factory_kwargs,
+        )
+        self.norm1 = create_norm_fn(norm, d_model, **factory_kwargs)
+        self.norm2 = create_norm_fn(norm, d_model, **factory_kwargs)
+
+        self.weights_per_step = weights_per_step
+        self.weights_per_step_schedule = weights_per_step_schedule
+        self.gating: nn.Module | nn.ModuleList | None = None
+        self.linear1: nn.Module | None = None
+        self.linear2: nn.Module | None = None
+        self.activation = activation
+
+        num_weights = 1
+        if weights_per_step:
+            num_weights = max(weights_per_step_schedule) + 1 if weights_per_step_schedule else weights_per_step
+
+        if gating == "none":
+            self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False, **factory_kwargs)
+            self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False, **factory_kwargs)
+        else:
+            if weights_per_step:
+                dim_ff_list = [dim_feedforward] * num_weights if isinstance(dim_feedforward, int) else dim_feedforward
+                self.gating = nn.ModuleList(
+                    [make_gating(gating, d_model, dim, **factory_kwargs) for dim in dim_ff_list]
+                )
+            else:
+                self.gating = make_gating(gating, d_model, dim_feedforward, **factory_kwargs)
+
+        if layer_scale is None:
+            self.layer_scale_1 = nn.Identity()
+            self.layer_scale_2 = nn.Identity()
+        else:
+            self.layer_scale_1 = MossAudioTokenizerLayerScale(
+                channels=d_model, init=layer_scale, channel_last=True, **cast(dict[str, object], factory_kwargs)
+            )
+            self.layer_scale_2 = MossAudioTokenizerLayerScale(
+                channels=d_model, init=layer_scale, channel_last=True, **cast(dict[str, object], factory_kwargs)
+            )
+
+    def _init_streaming_state(self, batch_size: int) -> LayerState:
+        device = next(iter(self.parameters())).device
+        return LayerState(batch_size, device, offset_cpu=0)
+
+    def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
+        state = self._streaming_state
+        offset = state.offset_cpu if isinstance(state, LayerState) else 0
+
+        x_orig = x
+        x = self.norm2(x)
+
+        if self.gating is None:
+            assert self.linear1 is not None
+            assert self.linear2 is not None
+            update = self.linear2(self.activation(self.linear1(x)))
+        else:
+            if self.weights_per_step:
+                assert isinstance(self.gating, nn.ModuleList)
+                update = apply_weights_per_step(self.gating, self.weights_per_step_schedule, x, offset)
+            else:
+                update = self.gating(x)
+        return x_orig.to(update) + self.layer_scale_2(update)
+
+    def _sa_block(self, x: torch.Tensor):
+        x_orig = x
+        x = self.norm1(x)
+        update = self.self_attn(x, x, x)
+        return x_orig.to(update) + self.layer_scale_1(update)
+
+    def forward(self, x: torch.Tensor):
+        x = self._sa_block(x)
+        x = self._ff_block(x)
+        state = self._streaming_state
+        if state is not None:
+            assert isinstance(state, LayerState)
+            state.offset_cpu += x.shape[1]
+        return x
+
+
+# =============================================================================
+# Streaming Transformer
+# =============================================================================
+
+
+@dataclass
+class TransformerState(StreamingState):
+    offsets: torch.Tensor
+
+    def reset(self, reset_mask: torch.Tensor):
+        super().reset(reset_mask)
+        self.offsets[:] = torch.where(reset_mask, torch.zeros_like(self.offsets), self.offsets)
+
+
+class MossAudioTokenizerTransformer(StreamingModule):
+    """Transformer with streaming/causal support."""
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        num_layers: int,
+        dim_feedforward: int = 2048,
+        causal: bool = False,
+        context: int | None = None,
+        positional_embedding: str = "sin",
+        max_period: float = 10_000,
+        positional_scale: float = 1.0,
+        device=None,
+        dtype=None,
+        **kwargs,
+    ):
+        super().__init__()
+        if d_model % num_heads != 0:
+            raise ValueError(f"d_model must be divisible by num_heads, got d_model={d_model}, num_heads={num_heads}")
+
+        self.positional_embedding = positional_embedding
+        self.max_period = max_period
+        self.positional_scale = positional_scale
+
+        self.rope: MossAudioTokenizerRotaryEmbedding | None = None
+        if positional_embedding in {"rope", "sin_rope"}:
+            self.rope = MossAudioTokenizerRotaryEmbedding(max_period=max_period)
+
+        self.layers = nn.ModuleList()
+        for _ in range(num_layers):
+            self.layers.append(
+                MossAudioTokenizerTransformerLayer(
+                    d_model=d_model,
+                    num_heads=num_heads,
+                    dim_feedforward=dim_feedforward,
+                    causal=causal,
+                    context=context,
+                    rope=self.rope,
+                    device=device,
+                    dtype=dtype,
+                    **kwargs,
+                )
+            )
+
+    def _init_streaming_state(self, batch_size: int) -> TransformerState:
+        device = next(self.parameters()).device
+        return TransformerState(
+            batch_size,
+            device,
+            offsets=torch.zeros(batch_size, device=device, dtype=torch.long),
+        )
+
+    def forward(self, x: torch.Tensor, *args, **kwargs):
+        B, T, C = x.shape
+        state = self._streaming_state
+        offsets = (
+            torch.zeros(1, dtype=torch.long, device=x.device)
+            if state is None
+            else (
+                state.offsets
+                if isinstance(state, TransformerState)
+                else torch.zeros(1, dtype=torch.long, device=x.device)
+            )
+        )
+
+        if self.positional_embedding in {"sin", "sin_rope"}:
+            positions = torch.arange(T, device=x.device).view(1, -1, 1)
+            positions = positions + offsets.view(-1, 1, 1)
+            pos_emb = create_sin_embedding(positions, C, max_period=self.max_period, dtype=x.dtype)
+            x = x + self.positional_scale * pos_emb
+
+        for layer in self.layers:
+            x = layer(x, *args, **kwargs)
+
+        if state is not None:
+            assert isinstance(state, TransformerState)
+            state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets)
+        return x
+
+
+class MossAudioTokenizerProjectedTransformer(StreamingContainer):
+    """Transformer with input/output projections."""
+
+    def __init__(
+        self,
+        input_dimension: int,
+        output_dimension: int,
+        d_model: int,
+        *,
+        conv_layout: bool = False,
+        module_type: str,
+        **kwargs,
+    ):
+        super().__init__()
+        self.module_type = module_type
+        self.downsample_ratio: int = 1
+        self.input_dimension = input_dimension
+        self.output_dimension = output_dimension
+
+        self.input_proj = nn.Linear(input_dimension, d_model, bias=False)
+        self.transformer = MossAudioTokenizerTransformer(d_model=d_model, **kwargs)
+        self.conv_layout = conv_layout
+        self.output_proj = nn.Linear(d_model, output_dimension, bias=False)
+
+    def forward(self, x, input_lengths, *args, **kwargs):
+        x = self.input_proj(x.transpose(1, 2))  # (B, D, T) -> (B, T, D)
+        x = self.transformer(x, *args, **kwargs)
+        x = self.output_proj(x).transpose(1, 2)  # (B, T, D) -> (B, D, T)
+        return x, input_lengths
+
+
+# =============================================================================
+# Patched Pretransform Module
+# =============================================================================
+
+
+class MossAudioTokenizerPatchedPretransform(nn.Module):
+    """Patching module for downsampling/upsampling."""
+
+    def __init__(self, patch_size: int, is_downsample: bool, module_type: str, **kwargs):
+        super().__init__()
+        self.patch_size = patch_size
+        self.downsample_ratio: int = patch_size
+        self.is_downsample = is_downsample
+        self.module_type = module_type
+
+    def encode(self, x, input_lengths):
+        b, d, _ = x.shape
+        h = self.patch_size
+        x = x.reshape(b, d, -1, h).permute(0, 1, 3, 2).reshape(b, d * h, -1)
+        # We pad the input waveform to a multiple of `downsample_rate` before applying the encoder.
+        # Use a ceil division to match that padding and avoid dropping the last (partially padded) frame.
+        output_lengths = input_lengths // self.patch_size
+        return x, output_lengths
+
+    def decode(self, x, input_lengths):
+        b, dh, length = x.shape
+        h = self.patch_size
+        d = dh // h
+        x = x.reshape(b, d, h, length).permute(0, 1, 3, 2).reshape(b, d, length * h)
+        output_lengths = input_lengths * self.patch_size
+        return x, output_lengths
+
+    def forward(self, x, input_lengths):
+        if self.is_downsample:
+            return self.encode(x, input_lengths)
+        else:
+            return self.decode(x, input_lengths)
+
+
+# =============================================================================
+# Vector Quantization
+# =============================================================================
+
+
+def WNConv1d(*args, **kwargs):
+    """Weight-normalized Conv1d."""
+    return nn.utils.parametrizations.weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+class MossAudioTokenizerVectorQuantize(nn.Module):
+    """Single codebook vector quantization (inference only)."""
+
+    def __init__(
+        self,
+        input_dim: int,
+        codebook_size: int,
+        codebook_dim: int,
+        **kwargs,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        if input_dim != codebook_dim:
+            self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
+            self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
+        else:
+            self.in_proj = nn.Identity()
+            self.out_proj = nn.Identity()
+
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    @torch.no_grad()
+    def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            z: Input tensor of shape (B, D, T)
+        Returns:
+            z_q: Quantized tensor of shape (B, D, T)
+            indices: Code indices of shape (B, T)
+            z_e: Encoded tensor before quantization
+        """
+        z = z.float()
+        z_e = self.in_proj(z).float()
+
+        encodings = z_e.transpose(1, 2).reshape(-1, z_e.shape[1])
+
+        codebook_weight = self.codebook.weight
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook_weight.float().t()
+            + codebook_weight.float().pow(2).sum(1, keepdim=True).t()
+        )
+
+        indices = (-dist).max(1)[1]
+        indices = indices.reshape(z.size(0), -1)
+
+        z_q = self.decode_code(indices)
+        z_q = self.out_proj(z_q).float()
+
+        return z_q, indices, z_e
+
+    def decode_code(self, embed_id: torch.Tensor) -> torch.Tensor:
+        """Decode code indices to embeddings."""
+        return self.codebook(embed_id).transpose(1, 2).float()
+
+
+class MossAudioTokenizerLFQ(nn.Module):
+    """LFQ (inference-only) used by ResidualLFQ."""
+
+    def __init__(
+        self,
+        input_dim: int,
+        codebook_size: int,
+        codebook_dim: int,
+        **kwargs,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        if self.input_dim != self.codebook_dim:
+            self.in_proj = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1)
+            self.out_proj = WNConv1d(self.codebook_dim, self.input_dim, kernel_size=1)
+        else:
+            self.in_proj = nn.Identity()
+            self.out_proj = nn.Identity()
+
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    @torch.no_grad()
+    def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Quantize z into codebook vectors."""
+        z = z.float()
+        z_e = self.in_proj(z).float()
+        z_q, indices = self.decode_latents(z_e)
+        z_q = (z_e + (z_q - z_e).detach()).float()
+        z_q = self.out_proj(z_q).float()
+        return z_q, indices, z_e
+
+    def embed_code(self, embed_id: torch.Tensor) -> torch.Tensor:
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code_wo_out_proj(self, embed_id: torch.Tensor) -> torch.Tensor:
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_code(self, embed_id: torch.Tensor) -> torch.Tensor:
+        z_q = self.decode_code_wo_out_proj(embed_id).float()
+        z_q = self.out_proj(z_q).float()
+        return z_q
+
+    def decode_latents(self, latents: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Match training LFQ: L2-normalize then argmin squared distance."""
+        encodings = latents.transpose(1, 2).reshape(-1, latents.shape[1]).float()
+        codebook = self.codebook.weight.float()
+
+        encodings = F.normalize(encodings)
+        codebook = F.normalize(codebook)
+
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = (-dist).max(1)[1]
+        indices = indices.reshape(latents.size(0), -1)
+        z_q = self.decode_code_wo_out_proj(indices).float()
+        return z_q, indices
+
+
+class MossAudioTokenizerResidualVQ(nn.Module):
+    """Residual Vector Quantization (inference only)."""
+
+    def __init__(
+        self,
+        input_dim: int = 1024,
+        rvq_dim: int | None = None,
+        output_dim: int | None = None,
+        num_quantizers: int = 32,
+        codebook_size: int = 1024,
+        codebook_dim: int = 8,
+        **kwargs,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.rvq_dim = rvq_dim or input_dim
+        self.output_dim = output_dim or input_dim
+        self.num_quantizers = num_quantizers
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        self.input_proj = (
+            WNConv1d(input_dim, self.rvq_dim, kernel_size=1) if input_dim != self.rvq_dim else nn.Identity()
+        )
+        self.output_proj = (
+            WNConv1d(self.rvq_dim, self.output_dim, kernel_size=1) if self.rvq_dim != self.output_dim else nn.Identity()
+        )
+
+        self.quantizers = nn.ModuleList(
+            [
+                MossAudioTokenizerVectorQuantize(
+                    input_dim=self.rvq_dim,
+                    codebook_size=codebook_size,
+                    codebook_dim=codebook_dim,
+                    **kwargs,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+
+    @torch.no_grad()
+    def forward(
+        self,
+        z: torch.Tensor,
+        input_length: torch.Tensor,
+        n_quantizers: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            z: Input tensor of shape (B, D, T)
+            input_length: Valid lengths for each sample (B,)
+            n_quantizers: Number of quantizers to use
+        Returns:
+            quantized_out: Quantized output (B, D, T)
+            all_indices: All code indices (N, B, T)
+            output_length: Output lengths (B,)
+        """
+        z = self.input_proj(z)
+
+        batch_size, _, max_time = z.shape
+        mask = torch.arange(max_time, device=z.device).expand(batch_size, max_time) < input_length.unsqueeze(1)
+
+        quantized_out = torch.zeros_like(z, dtype=torch.float32)
+        residual = z.clone().float()
+        all_indices = []
+
+        n_quantizers = n_quantizers or self.num_quantizers
+
+        for i, quantizer in enumerate(self.quantizers):
+            if i >= n_quantizers:
+                break
+
+            masked_residual = residual * mask.unsqueeze(1)
+            z_q_i, indices_i, _ = quantizer(masked_residual)
+
+            update_mask = mask.unsqueeze(1)
+            quantized_out = quantized_out + z_q_i * update_mask
+            residual = residual - z_q_i * update_mask
+            all_indices.append(indices_i)
+
+        all_indices = torch.stack(all_indices)  # (N, B, T)
+        quantized_out = self.output_proj(quantized_out)
+
+        return quantized_out, all_indices, input_length
+
+    def decode_codes(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode codes from multiple quantizers to embeddings."""
+        nq, B, T = codes.shape
+        emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32)
+
+        for i, quantizer in enumerate(self.quantizers[:nq]):
+            quantizer = cast(MossAudioTokenizerVectorQuantize, quantizer)
+            quantized_i = quantizer.decode_code(codes[i])
+            emb += quantized_i
+
+        emb = self.output_proj(emb)
+        return emb
+
+
+class MossAudioTokenizerResidualLFQ(nn.Module):
+    """Residual LFQ (inference only)."""
+
+    def __init__(
+        self,
+        input_dim: int = 1024,
+        rvq_dim: int | None = None,
+        output_dim: int | None = None,
+        num_quantizers: int = 32,
+        codebook_size: int = 1024,
+        codebook_dim: int = 8,
+        **kwargs,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.rvq_dim = rvq_dim or input_dim
+        self.output_dim = output_dim or input_dim
+        self.num_quantizers = num_quantizers
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        self.input_proj = (
+            WNConv1d(input_dim, self.rvq_dim, kernel_size=1) if input_dim != self.rvq_dim else nn.Identity()
+        )
+        self.output_proj = (
+            WNConv1d(self.rvq_dim, self.output_dim, kernel_size=1) if self.rvq_dim != self.output_dim else nn.Identity()
+        )
+
+        self.quantizers = nn.ModuleList(
+            [
+                MossAudioTokenizerLFQ(
+                    input_dim=self.rvq_dim,
+                    codebook_size=codebook_size,
+                    codebook_dim=codebook_dim,
+                    **kwargs,
+                )
+                for _ in range(num_quantizers)
+            ]
+        )
+
+    @torch.no_grad()
+    def forward(
+        self,
+        z: torch.Tensor,
+        input_length: torch.Tensor,
+        n_quantizers: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Inference quantization."""
+        z = self.input_proj(z).float()
+
+        batch_size, _, max_time = z.shape
+        mask = torch.arange(max_time, device=z.device).expand(batch_size, max_time) < input_length.unsqueeze(1)
+
+        quantized_out = torch.zeros_like(z, dtype=torch.float32)
+        residual = z.clone().float()
+        all_indices = []
+
+        n_quantizers = n_quantizers or self.num_quantizers
+        for i, quantizer in enumerate(self.quantizers):
+            if i >= n_quantizers:
+                break
+
+            masked_residual = residual * mask.unsqueeze(1)
+            z_q_i, indices_i, _ = quantizer(masked_residual)
+
+            update_mask = mask.unsqueeze(1)
+            quantized_out = quantized_out + z_q_i * update_mask
+            residual = residual - z_q_i * update_mask
+            all_indices.append(indices_i)
+
+        all_indices = (
+            torch.stack(all_indices)
+            if all_indices
+            else torch.empty(0, batch_size, max_time, device=z.device, dtype=torch.long)
+        )
+        quantized_out = self.output_proj(quantized_out)
+        return quantized_out, all_indices, input_length
+
+    def decode_codes(self, codes: torch.Tensor) -> torch.Tensor:
+        nq, B, T = codes.shape
+        emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32)
+        for i, quantizer in enumerate(self.quantizers[:nq]):
+            quantizer = cast(MossAudioTokenizerLFQ, quantizer)
+            emb += quantizer.decode_code(codes[i]).float()
+        emb = self.output_proj(emb)
+        return emb
+
+
+# =============================================================================
+# Main Model Classes
+# =============================================================================
+
+
+@auto_docstring
+class MossAudioTokenizerPreTrainedModel(PreTrainedAudioTokenizerBase):
+    """Base class for MossAudioTokenizer models."""
+
+    config_class = MossAudioTokenizerConfig
+    base_model_prefix = ""
+    main_input_name = "input_values"
+    input_modalities = "audio"
+    supports_gradient_checkpointing = False
+    _no_split_modules = [
+        "MossAudioTokenizerTransformerLayer",
+        "MossAudioTokenizerResidualVQ",
+        "MossAudioTokenizerResidualLFQ",
+    ]
+
+
+@auto_docstring(
+    custom_intro="""
+    The MossAudioTokenizer neural audio codec model for audio tokenization and synthesis.
+    """
+)
+class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel):
+    """
+    MossAudioTokenizer model for audio tokenization and synthesis.
+
+    This model can encode audio waveforms into discrete tokens and decode
+    tokens back into audio waveforms.
+    """
+
+    def __init__(self, config: MossAudioTokenizerConfig):
+        super().__init__(config)
+
+        self.config = config
+        _ = config.version
+        self.sampling_rate = config.sampling_rate
+        self.downsample_rate = config.downsample_rate
+        self.number_channels = getattr(config, "number_channels", getattr(config, "num_channels", 1))
+        self.enable_channel_interleave = getattr(config, "enable_channel_interleave", True)
+        self.causal_transformer_context_duration = config.causal_transformer_context_duration
+
+        # Build encoder
+        channel_interleave_factor = (
+            self.number_channels if self.enable_channel_interleave and self.number_channels > 1 else 1
+        )
+        current_frame_rate: float = float(self.sampling_rate * channel_interleave_factor)
+        self.encoder = nn.ModuleList()
+
+        for encoder_kwargs_i in config.encoder_kwargs:
+            encoder_kwargs_i = dict(encoder_kwargs_i)  # Make a copy
+            if encoder_kwargs_i["module_type"] == "PatchedPretransform":
+                self.encoder.append(MossAudioTokenizerPatchedPretransform(**encoder_kwargs_i, is_downsample=True))
+            elif encoder_kwargs_i["module_type"] == "Transformer":
+                context_duration = float(
+                    encoder_kwargs_i.pop("context_duration", self.causal_transformer_context_duration)
+                )
+                self.encoder.append(
+                    MossAudioTokenizerProjectedTransformer(
+                        **encoder_kwargs_i,
+                        context=int(round(current_frame_rate * context_duration)),
+                    )
+                )
+            current_frame_rate /= self.encoder[-1].downsample_ratio
+
+        # Build quantizer
+        quantizer_kwargs = dict(config.quantizer_kwargs)
+        quantizer_type = quantizer_kwargs.get("quantizer_type", getattr(config, "quantizer_type", "rvq"))
+        if quantizer_type in {"rvq", "spec_rvq"}:
+            self.quantizer = MossAudioTokenizerResidualVQ(**quantizer_kwargs)
+        elif quantizer_type in {"rlfq", "random_prefix_rlfq"}:
+            self.quantizer = MossAudioTokenizerResidualLFQ(**quantizer_kwargs)
+        else:
+            raise ValueError(f"Unsupported quantizer_type: {quantizer_type}")
+
+        # Build decoder
+        decoder_kwargs_list = copy.deepcopy(config.decoder_kwargs)
+        self.decoder = nn.ModuleList()
+
+        for decoder_kwargs_i in decoder_kwargs_list:
+            decoder_kwargs_i = dict(decoder_kwargs_i)
+            if decoder_kwargs_i["module_type"] == "PatchedPretransform":
+                self.decoder.append(MossAudioTokenizerPatchedPretransform(**decoder_kwargs_i, is_downsample=False))
+            elif decoder_kwargs_i["module_type"] == "Transformer":
+                context_duration = float(
+                    decoder_kwargs_i.pop("context_duration", self.causal_transformer_context_duration)
+                )
+                self.decoder.append(
+                    MossAudioTokenizerProjectedTransformer(
+                        **decoder_kwargs_i,
+                        context=int(round(current_frame_rate * context_duration)),
+                    )
+                )
+            current_frame_rate *= self.decoder[-1].downsample_ratio
+
+        expected_output_frame_rate = float(self.sampling_rate * channel_interleave_factor)
+        if int(round(current_frame_rate)) != int(round(expected_output_frame_rate)):
+            raise ValueError(
+                "Decoder stack does not invert the encoder frame rate correctly: "
+                f"got current_frame_rate={current_frame_rate}, expected={expected_output_frame_rate}."
+            )
+
+        self.post_init()
+
+    def _start_streaming(self, batch_size: int):
+        """Start streaming mode for all modules."""
+
+        def _start(module):
+            if isinstance(module, StreamingModule):
+                module._streaming_state = module._init_streaming_state(batch_size)
+
+        self.apply(_start)
+
+    def _stop_streaming(self):
+        """Stop streaming mode for all modules."""
+
+        def _stop(module):
+            if isinstance(module, StreamingModule):
+                module._streaming_state = None
+
+        self.apply(_stop)
+
+    def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None:
+        """Set the active slot mask on every live streaming module."""
+
+        def _set(module):
+            if isinstance(module, StreamingModule) and module._streaming_state is not None:
+                module._streaming_state.set_exec_mask(exec_mask.to(module._streaming_state.device))
+
+        self.apply(_set)
+
+    @contextmanager
+    def streaming(self, batch_size: int = 1):
+        """Context manager for streaming mode."""
+        self._start_streaming(batch_size)
+        try:
+            yield
+        finally:
+            self._stop_streaming()
+
+    def _prepare_waveform_batch(self, wav_list: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+        if len(wav_list) == 0:
+            raise ValueError("`wav_list` must contain at least one waveform.")
+
+        device = wav_list[0].device
+        dtype = wav_list[0].dtype
+        batch_size = len(wav_list)
+        input_lengths = torch.zeros(batch_size, device=device, dtype=torch.long)
+
+        normalized_wavs: list[torch.Tensor] = []
+        for i, wav in enumerate(wav_list):
+            if self.number_channels == 1:
+                if wav.dim() == 1:
+                    wav_i = wav.unsqueeze(0)
+                elif wav.dim() == 2 and wav.shape[0] == 1:
+                    wav_i = wav
+                else:
+                    raise ValueError(
+                        f"Expected wav_list[{i}] to have shape `(T,)` or `(1, T)` for a mono model, "
+                        f"got {tuple(wav.shape)}."
+                    )
+            else:
+                if wav.dim() != 2 or wav.shape[0] != self.number_channels:
+                    raise ValueError(
+                        f"Expected wav_list[{i}] to have shape `({self.number_channels}, T)`, got {tuple(wav.shape)}."
+                    )
+                wav_i = wav
+
+            normalized_wavs.append(wav_i)
+            input_lengths[i] = wav_i.shape[-1]
+
+        max_length = int(input_lengths.max().item()) if batch_size > 0 else 0
+        input_values = torch.zeros(
+            batch_size,
+            self.number_channels,
+            max_length,
+            device=device,
+            dtype=dtype,
+        )
+        for i, wav_i in enumerate(normalized_wavs):
+            input_values[i, :, : wav_i.shape[-1]] = wav_i
+        return input_values, input_lengths
+
+    def _prepare_codes_batch(
+        self,
+        codes_list: list[torch.Tensor],
+        num_quantizers: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, int]:
+        if len(codes_list) == 0:
+            raise ValueError("`codes_list` must contain at least one code tensor.")
+
+        batch_size = len(codes_list)
+        device = codes_list[0].device
+        nqs = [codes.shape[0] for codes in codes_list]
+        if num_quantizers is None:
+            num_quantizers = nqs[0]
+            if any(nq != num_quantizers for nq in nqs):
+                raise ValueError(
+                    "All elements in `codes_list` must have the same number of quantizers when "
+                    "`num_quantizers` is None. Pass `num_quantizers=...` to decode a common prefix."
+                )
+        else:
+            min_nq = min(nqs)
+            if min_nq < num_quantizers:
+                raise ValueError(
+                    "`num_quantizers` must be <= the number of quantizers for every element in `codes_list`. "
+                    f"Got num_quantizers={num_quantizers}, min(codes.shape[0])={min_nq}."
+                )
+
+        max_length = max(codes.shape[-1] for codes in codes_list)
+        audio_codes = torch.zeros(num_quantizers, batch_size, max_length, device=device, dtype=torch.long)
+        audio_codes_lengths = torch.zeros(batch_size, device=device, dtype=torch.long)
+
+        for i, codes in enumerate(codes_list):
+            codes_i = codes[:num_quantizers]
+            audio_codes[:, i, : codes_i.shape[-1]] = codes_i
+            audio_codes_lengths[i] = codes_i.shape[-1]
+        return audio_codes, audio_codes_lengths, num_quantizers
+
+    def _flatten_channels_for_codec(
+        self,
+        input_values: torch.Tensor,
+        input_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if input_values.dim() != 3:
+            raise ValueError(f"Expected `input_values` with shape `(B, C, T)`, got {tuple(input_values.shape)}.")
+        if input_values.shape[1] != self.number_channels:
+            raise ValueError(
+                f"Expected `input_values.shape[1] == {self.number_channels}`, got {input_values.shape[1]}."
+            )
+
+        if input_values.shape[-1] % self.downsample_rate != 0:
+            pad_length = self.downsample_rate - (input_values.shape[-1] % self.downsample_rate)
+            input_values = F.pad(input_values, (0, pad_length))
+
+        if self.number_channels > 1 and self.enable_channel_interleave:
+            input_values = input_values.transpose(1, 2).contiguous().view(input_values.shape[0], 1, -1)
+            input_lengths = input_lengths * self.number_channels
+        return input_values, input_lengths
+
+    def _restore_channels_from_codec(
+        self,
+        output_values: torch.Tensor,
+        output_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.number_channels == 1 or not self.enable_channel_interleave:
+            return output_values.float(), output_lengths
+
+        output_values = (
+            output_values.squeeze(1)
+            .contiguous()
+            .view(output_values.shape[0], -1, self.number_channels)
+            .transpose(1, 2)
+            .contiguous()
+            .float()
+        )
+        output_lengths = torch.div(output_lengths, self.number_channels, rounding_mode="floor")
+        return output_values, output_lengths
+
+    @torch.no_grad()
+    def batch_encode(
+        self, wav_list: list[torch.Tensor], num_quantizers: int | None = None
+    ) -> MossAudioTokenizerEncoderOutput:
+        """Batch encode a list of audio waveforms.
+
+        Args:
+            wav_list: List of audio tensors, each of shape `(num_samples,)`.
+            num_quantizers: Number of quantizers to use. By default, all quantizers are used.
+
+        Returns:
+            [`MossAudioTokenizerEncoderOutput`] with `audio_codes` and `audio_codes_lengths`.
+        """
+        input_values, input_lengths = self._prepare_waveform_batch(wav_list)
+        return self._encode_frame(input_values, input_lengths, n_quantizers=num_quantizers)
+
+    @torch.no_grad()
+    def batch_decode(
+        self, codes_list: list[torch.Tensor], num_quantizers: int | None = None
+    ) -> MossAudioTokenizerDecoderOutput:
+        """Batch decode a list of audio codes.
+
+        Args:
+            codes_list: List of audio code tensors, each of shape `(num_quantizers, codes_length)`.
+            num_quantizers: If provided, decode only the first `num_quantizers` quantizers from each element in
+                `codes_list`. If omitted, all elements in `codes_list` must have the same number of quantizers.
+
+        Returns:
+            [`MossAudioTokenizerDecoderOutput`] with `audio` and `audio_lengths`.
+        """
+        audio_codes, audio_codes_lengths, _ = self._prepare_codes_batch(
+            codes_list,
+            num_quantizers=num_quantizers,
+        )
+        return self._decode_frame(audio_codes, audio_codes_lengths)
+
+    @torch.no_grad()
+    def _encode_frame(
+        self,
+        input_values: torch.Tensor,
+        input_lengths: torch.Tensor | None = None,
+        n_quantizers: int | None = None,
+    ) -> MossAudioTokenizerEncoderOutput:
+        """Tokenize audio waveform into discrete tokens."""
+        # Handle input shape
+        if input_values.dim() == 1:
+            input_values = input_values.view(1, 1, -1)
+        elif input_values.dim() == 2:
+            if self.number_channels == 1:
+                input_values = input_values.unsqueeze(1)
+            else:
+                input_values = input_values.unsqueeze(0)
+
+        B, _, T = input_values.shape
+        device = input_values.device
+
+        if input_lengths is None:
+            input_lengths = torch.full((B,), T, device=device, dtype=torch.long)
+
+        input_values, input_lengths = self._flatten_channels_for_codec(input_values, input_lengths)
+
+        # Encode
+        e, e_lengths = input_values, input_lengths
+        for encoder_module in self.encoder:
+            e, e_lengths = encoder_module(e, e_lengths)
+
+        # Quantize
+        quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer)
+        zq, audio_codes, audio_codes_lengths = quantizer(e, e_lengths, n_quantizers)
+
+        return MossAudioTokenizerEncoderOutput(
+            audio_codes=audio_codes, audio_codes_lengths=audio_codes_lengths, encoder_hidden_states=e
+        )
+
+    @torch.no_grad()
+    def _decode_frame(
+        self,
+        codes: torch.Tensor,
+        codes_lengths: torch.Tensor | None = None,
+    ) -> MossAudioTokenizerDecoderOutput:
+        """Detokenize discrete tokens into audio waveform."""
+        nq, B, T = codes.shape
+        device = codes.device
+
+        if codes_lengths is None:
+            codes_lengths = torch.full((B,), T, device=device, dtype=torch.long)
+
+        # Decode from codes
+        quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer)
+        zq = quantizer.decode_codes(codes)
+
+        d, d_lengths = zq, codes_lengths
+        for decoder_module in self.decoder:
+            d, d_lengths = decoder_module(d, d_lengths)
+
+        d, d_lengths = self._restore_channels_from_codec(d, d_lengths)
+        return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_lengths)
+
+    def encode(  # type: ignore[override]
+        self,
+        input_values: torch.Tensor,
+        padding_mask: torch.Tensor | None = None,
+        num_quantizers: int | None = None,
+        return_dict: bool | None = None,
+        chunk_duration: float | None = None,
+    ) -> MossAudioTokenizerEncoderOutput | tuple[torch.Tensor, torch.Tensor]:
+        """
+        Encodes the input audio waveform into discrete codes.
+
+        Args:
+            input_values (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
+                Float values of the input audio waveform.
+            padding_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to indicate valid audio samples.
+            num_quantizers (`int`, *optional*):
+                Number of quantizers to use. By default, all quantizers are used.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+            chunk_duration (`float`, *optional*):
+                If provided, encode the input waveform in successive chunks of `chunk_duration` seconds while keeping a
+                streaming KV cache for the causal transformers.
+
+                `chunk_duration` must be <= `config.causal_transformer_context_duration`, and
+                `chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`.
+
+        Returns:
+            `MossAudioTokenizerEncoderOutput` or tuple containing audio codes and lengths.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        # Handle input shape
+        if input_values.dim() == 1:
+            input_values = input_values.view(1, 1, -1)
+        elif input_values.dim() == 2:
+            if self.number_channels == 1:
+                input_values = input_values.unsqueeze(1)
+            else:
+                input_values = input_values.unsqueeze(0)
+
+        B, _, T = input_values.shape
+        device = input_values.device
+
+        if padding_mask is not None:
+            input_lengths = padding_mask.sum(dim=-1).long()
+        else:
+            input_lengths = torch.full((B,), T, device=device, dtype=torch.long)
+
+        if chunk_duration is None:
+            encoder_output = self._encode_frame(input_values, input_lengths, num_quantizers)
+        else:
+            if chunk_duration <= 0:
+                raise ValueError("`chunk_duration` must be > 0 when provided.")
+            if chunk_duration > self.causal_transformer_context_duration:
+                raise ValueError(
+                    "`chunk_duration` must be <= `config.causal_transformer_context_duration` "
+                    f"({self.causal_transformer_context_duration}), got {chunk_duration}."
+                )
+            if B != 1:
+                raise ValueError("Streaming encode via `chunk_duration` currently only supports batch_size=1.")
+
+            chunk_length = int(round(chunk_duration * self.sampling_rate))
+            if chunk_length <= 0:
+                raise ValueError("`chunk_duration` is too small and results in chunk_length <= 0.")
+            if chunk_length % self.downsample_rate != 0:
+                raise ValueError(
+                    "`chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. "
+                    f"Got chunk_length={chunk_length}, downsample_rate={self.downsample_rate}."
+                )
+
+            input_length = int(input_lengths[0].item())
+            if input_length <= chunk_length:
+                encoder_output = self._encode_frame(input_values[..., :input_length], input_lengths, num_quantizers)
+            else:
+                codes_chunks: list[torch.Tensor] = []
+                hidden_chunks: list[torch.Tensor] = []
+
+                with ExitStack() as exit_stack:
+                    for encoder_module in self.encoder:
+                        if isinstance(encoder_module, StreamingModule):
+                            exit_stack.enter_context(encoder_module.streaming(batch_size=B))
+
+                    for start_idx in range(0, input_length, chunk_length):
+                        input_length_i = min(chunk_length, input_length - start_idx)
+                        if input_length_i <= 0:
+                            break
+
+                        input_lengths_i = torch.tensor([input_length_i], device=device, dtype=torch.long)
+                        input_values_i = input_values[..., start_idx : start_idx + input_length_i]
+                        result_i = self._encode_frame(input_values_i, input_lengths_i, num_quantizers)
+
+                        if result_i.audio_codes is None or result_i.audio_codes_lengths is None:
+                            raise RuntimeError("Internal error: `_encode_frame` returned empty audio codes.")
+                        if result_i.encoder_hidden_states is None:
+                            raise RuntimeError("Internal error: `_encode_frame` returned empty encoder hidden states.")
+
+                        codes_length_i = result_i.audio_codes_lengths
+                        codes_chunks.append(result_i.audio_codes[:, :, : codes_length_i[0]])
+                        hidden_chunks.append(result_i.encoder_hidden_states[:, :, : codes_length_i[0]])
+
+                audio_codes = torch.cat(codes_chunks, dim=-1)
+                encoder_hidden_states = torch.cat(hidden_chunks, dim=-1)
+                audio_codes_lengths = torch.tensor([audio_codes.shape[-1]], device=device, dtype=torch.long)
+                encoder_output = MossAudioTokenizerEncoderOutput(
+                    audio_codes=audio_codes,
+                    audio_codes_lengths=audio_codes_lengths,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
+
+        if not return_dict:
+            assert encoder_output.audio_codes is not None
+            assert encoder_output.audio_codes_lengths is not None
+            return (
+                cast(torch.Tensor, encoder_output.audio_codes),
+                cast(torch.Tensor, encoder_output.audio_codes_lengths),
+            )
+        return encoder_output
+
+    def decode(  # type: ignore[override]
+        self,
+        audio_codes: torch.Tensor,
+        padding_mask: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+        chunk_duration: float | None = None,
+        num_quantizers: int | None = None,
+    ) -> MossAudioTokenizerDecoderOutput | tuple[torch.Tensor]:
+        """
+        Decodes the given codes into an output audio waveform.
+
+        Args:
+            audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`):
+                Discrete code embeddings computed using `model.encode`.
+            padding_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to indicate valid code positions.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+            chunk_duration (`float`, *optional*):
+                If provided, decode the input codes in successive chunks of `chunk_duration` seconds while keeping a
+                streaming KV cache for the causal transformers.
+
+            num_quantizers (`int`, *optional*):
+                Number of quantizers to use. By default, all quantizers in `audio_codes` are used.
+
+                `chunk_duration` must be <= `config.causal_transformer_context_duration`, and
+                `chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`.
+
+        Returns:
+            `MossAudioTokenizerDecoderOutput` or tuple containing decoded audio.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        if audio_codes.dim() == 2:
+            audio_codes = audio_codes.unsqueeze(1)  # nq, T -> nq, B=1, T
+
+        if num_quantizers is not None:
+            if num_quantizers > audio_codes.shape[0]:
+                raise ValueError(
+                    f"`num_quantizers` ({num_quantizers}) must be <= audio_codes.shape[0] ({audio_codes.shape[0]})."
+                )
+            audio_codes = audio_codes[:num_quantizers]
+
+        _, B, T = audio_codes.shape
+        device = audio_codes.device
+
+        if padding_mask is not None:
+            codes_lengths = padding_mask.sum(dim=-1).long()
+        else:
+            codes_lengths = torch.full((B,), T, device=device, dtype=torch.long)
+
+        if chunk_duration is None:
+            decoder_output = self._decode_frame(audio_codes, codes_lengths)
+        else:
+            if chunk_duration <= 0:
+                raise ValueError("`chunk_duration` must be > 0 when provided.")
+            if chunk_duration > self.causal_transformer_context_duration:
+                raise ValueError(
+                    "`chunk_duration` must be <= `config.causal_transformer_context_duration` "
+                    f"({self.causal_transformer_context_duration}), got {chunk_duration}."
+                )
+            if B != 1:
+                raise ValueError("Streaming decode via `chunk_duration` currently only supports batch_size=1.")
+
+            chunk_length = int(round(chunk_duration * self.sampling_rate))
+            if chunk_length <= 0:
+                raise ValueError("`chunk_duration` is too small and results in chunk_length <= 0.")
+            if chunk_length % self.downsample_rate != 0:
+                raise ValueError(
+                    "`chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. "
+                    f"Got chunk_length={chunk_length}, downsample_rate={self.downsample_rate}."
+                )
+
+            chunk_frame_length = chunk_length // self.downsample_rate
+            codes_length = int(codes_lengths[0].item())
+            if codes_length <= chunk_frame_length:
+                decoder_output = self._decode_frame(audio_codes[..., :codes_length], codes_lengths)
+            else:
+                wav_chunks: list[torch.Tensor] = []
+                with ExitStack() as exit_stack:
+                    for decoder_module in self.decoder:
+                        if isinstance(decoder_module, StreamingModule):
+                            exit_stack.enter_context(decoder_module.streaming(batch_size=B))
+
+                    for start_idx in range(0, codes_length, chunk_frame_length):
+                        codes_length_i = min(chunk_frame_length, codes_length - start_idx)
+                        if codes_length_i <= 0:
+                            break
+
+                        codes_lengths_i = torch.tensor([codes_length_i], device=device, dtype=torch.long)
+                        codes_i = audio_codes[:, :, start_idx : start_idx + codes_length_i]
+                        result_i = self._decode_frame(codes_i, codes_lengths_i)
+
+                        if result_i.audio is None or result_i.audio_lengths is None:
+                            raise RuntimeError("Internal error: `_decode_frame` returned empty audio.")
+
+                        wav_chunks.append(result_i.audio[:, :, : result_i.audio_lengths[0]])
+
+                wav = torch.cat(wav_chunks, dim=-1)
+                audio_lengths = torch.tensor([wav.shape[-1]], device=device, dtype=torch.long)
+                decoder_output = MossAudioTokenizerDecoderOutput(audio=wav, audio_lengths=audio_lengths)
+
+        if not return_dict:
+            assert decoder_output.audio is not None
+            return (cast(torch.Tensor, decoder_output.audio),)
+        return decoder_output
+
+    @auto_docstring
+    def forward(
+        self,
+        input_values: torch.FloatTensor | None = None,
+        padding_mask: torch.BoolTensor | None = None,
+        audio_codes: torch.Tensor | None = None,
+        num_quantizers: int | None = None,
+        return_dict: bool | None = None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None] | MossAudioTokenizerOutput:  # type: ignore[override]
+        r"""
+        input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*):
+            Raw audio input converted to Float.
+        padding_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid computing on padding token indices. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*):
+            Discrete code embeddings computed using `model.encode`.
+        num_quantizers (`int`, *optional*):
+            Number of quantizers (codebooks) to use. By default, all quantizers are used.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+
+        Examples:
+
+        ```python
+        >>> import torch
+        >>> from transformers import MossAudioTokenizerModel
+
+        >>> model = MossAudioTokenizerModel.from_pretrained("moss_audio_tokenizer-model")
+
+        >>> # Create dummy audio input
+        >>> audio = torch.randn(1, 1, 24000)  # 1 second of audio at 24kHz
+
+        >>> outputs = model(input_values=audio)
+        >>> audio_codes = outputs.audio_codes
+        >>> audio_values = outputs.audio
+        ```
+        """
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        output_audio_codes: torch.Tensor | None = None
+        output_audio_codes_lengths: torch.Tensor | None = None
+        output_audio: torch.Tensor | None = None
+        output_audio_lengths: torch.Tensor | None = None
+        decoded_from_encoded_codes = False
+
+        # Encode if input_values provided
+        if input_values is not None:
+            encoder_output = self.encode(input_values, padding_mask, num_quantizers, return_dict=True)
+            encoder_output = cast(MossAudioTokenizerEncoderOutput, encoder_output)
+            output_audio_codes = encoder_output.audio_codes
+            output_audio_codes_lengths = encoder_output.audio_codes_lengths
+
+            # If codes not provided separately, use encoded codes for decoding
+            if audio_codes is None:
+                audio_codes = output_audio_codes
+                decoded_from_encoded_codes = True
+
+        # Decode if codes available
+        if audio_codes is not None:
+            # If we're decoding the codes we just produced, use the computed lengths so we don't decode padded garbage.
+            if decoded_from_encoded_codes and output_audio_codes_lengths is not None:
+                decoder_output = self._decode_frame(audio_codes, output_audio_codes_lengths)
+            else:
+                decoder_output = self.decode(
+                    audio_codes,
+                    padding_mask=padding_mask,
+                    return_dict=True,
+                    num_quantizers=num_quantizers,
+                )
+                decoder_output = cast(MossAudioTokenizerDecoderOutput, decoder_output)
+            output_audio = decoder_output.audio
+            output_audio_lengths = decoder_output.audio_lengths
+
+        if not return_dict:
+            return (output_audio_codes, output_audio, output_audio_lengths)
+
+        return MossAudioTokenizerOutput(
+            audio=output_audio,
+            audio_lengths=output_audio_lengths,
+            audio_codes=output_audio_codes,
+            audio_codes_lengths=output_audio_codes_lengths,
+        )
+
+
+__all__ = ["MossAudioTokenizerModel", "MossAudioTokenizerPreTrainedModel"]
diff --git a/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py
new file mode 100644
index 000000000..c518742a2
--- /dev/null
+++ b/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py
@@ -0,0 +1,359 @@
+# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MossAudioTokenizer model configuration"""
+
+from typing import Any
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class MossAudioTokenizerConfig(PreTrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MossAudioTokenizerModel`]. It is used to
+    instantiate a MossAudioTokenizer model according to the specified arguments, defining the model architecture.
+
+    Instantiating a configuration with the defaults will yield a similar configuration to that of the
+    [VoiceAgentGroup/moss_audio_tokenizer](https://huggingface.co/VoiceAgentGroup/moss_audio_tokenizer) architecture.
+
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
+
+    Args:
+        sampling_rate (`int`, *optional*, defaults to 24000):
+            The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz).
+        downsample_rate (`int`, *optional*, defaults to 1920):
+            Total downsampling rate from waveform to tokens.
+        causal_transformer_context_duration (`float`, *optional*, defaults to 10.0):
+            Context duration in seconds for causal transformer.
+        encoder_kwargs (`list[dict]`, *optional*):
+            List of encoder module configurations. Each dict specifies a module type and its parameters.
+        decoder_kwargs (`list[dict]`, *optional*):
+            List of decoder module configurations in execution order.
+        quantizer_type (`str`, *optional*, defaults to `"rvq"`):
+            Quantizer type. Options include `"rvq"`, `"spec_rvq"`, `"rlfq"`, `"random_prefix_rlfq"`.
+        quantizer_kwargs (`dict`, *optional*):
+            Configuration for the quantizer including `input_dim`, `rvq_dim`, `output_dim`, `num_quantizers`,
+            `codebook_size`, and `codebook_dim`.
+
+    Example:
+
+    ```python
+    >>> from transformers import MossAudioTokenizerModel, MossAudioTokenizerConfig
+
+    >>> # Initializing a MossAudioTokenizer style configuration
+    >>> configuration = MossAudioTokenizerConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = MossAudioTokenizerModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "moss-audio-tokenizer"
+
+    # Backward-compatible alias used by some checkpoints.
+    attribute_map = {"sample_rate": "sampling_rate"}
+
+    sampling_rate: int
+    downsample_rate: int
+    causal_transformer_context_duration: float
+    encoder_kwargs: list[dict[str, Any]]
+    decoder_kwargs: list[dict[str, Any]]
+    number_channels: int
+    enable_channel_interleave: bool
+    attention_implementation: str
+    compute_dtype: str
+    codec_weight_dtype: str
+    quantizer_type: str
+    quantizer_kwargs: dict[str, Any]
+
+    def __init__(
+        self,
+        version: str | None = None,
+        sampling_rate: int = 48000,
+        downsample_rate: int = 3840,
+        causal_transformer_context_duration: float = 10.0,
+        encoder_kwargs: list[dict[str, Any]] | None = None,
+        decoder_kwargs: list[dict[str, Any]] | None = None,
+        number_channels: int = 2,
+        enable_channel_interleave: bool = True,
+        attention_implementation: str = "sdpa",
+        compute_dtype: str = "fp32",
+        codec_weight_dtype: str = "fp32",
+        quantizer_type: str = "rlfq",
+        quantizer_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        # Some checkpoints might include an incorrect/legacy `model_type` (e.g. "speech_tokenizer").
+        # We drop it to avoid overriding the class-level `model_type`.
+        kwargs.pop("model_type", None)
+        if "channels_numbers" in kwargs:
+            number_channels = kwargs.pop("channels_numbers")
+        if "enable_channel_interleave" in kwargs:
+            enable_channel_interleave = kwargs.pop("enable_channel_interleave")
+        if "attention_backend" in kwargs and attention_implementation == "sdpa":
+            attention_implementation = kwargs.pop("attention_backend")
+        if "codec_compute_dtype" in kwargs and compute_dtype == "fp32":
+            compute_dtype = kwargs.pop("codec_compute_dtype")
+        if "codec_load_dtype" in kwargs and codec_weight_dtype == "fp32":
+            codec_weight_dtype = kwargs.pop("codec_load_dtype")
+
+        # `version` is accepted for compatibility but not used in modeling.
+        self.version = version
+        self.sampling_rate = sampling_rate
+        self.downsample_rate = downsample_rate
+        self.causal_transformer_context_duration = causal_transformer_context_duration
+        self.number_channels = number_channels
+        self.enable_channel_interleave = enable_channel_interleave
+        self.attention_implementation = attention_implementation
+        self.compute_dtype = compute_dtype
+        self.codec_weight_dtype = codec_weight_dtype
+        # Default encoder configuration
+        if encoder_kwargs is None:
+            encoder_kwargs = [
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 240,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 240,
+                    "output_dimension": 384,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 384,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 640,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 1280,
+                    "output_dimension": 768,
+                    "d_model": 1280,
+                    "num_heads": 20,
+                    "num_layers": 32,
+                    "dim_feedforward": 5120,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+            ]
+        self.encoder_kwargs = encoder_kwargs
+
+        # Default decoder configuration (execution order)
+        if decoder_kwargs is None:
+            decoder_kwargs = [
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 768,
+                    "output_dimension": 1280,
+                    "d_model": 1280,
+                    "num_heads": 20,
+                    "num_layers": 32,
+                    "dim_feedforward": 5120,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 640,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 768,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 2,
+                },
+                {
+                    "module_type": "Transformer",
+                    "input_dimension": 384,
+                    "output_dimension": 240,
+                    "d_model": 768,
+                    "num_heads": 12,
+                    "num_layers": 12,
+                    "dim_feedforward": 3072,
+                    "causal": True,
+                    "norm": "layer_norm",
+                    "positional_embedding": "rope",
+                    "max_period": 10000,
+                    "gating": "none",
+                    "layer_scale": 0.01,
+                    "conv_layout": True,
+                },
+                {
+                    "module_type": "PatchedPretransform",
+                    "patch_size": 240,
+                },
+            ]
+        self.decoder_kwargs = decoder_kwargs
+
+        # Default quantizer configuration
+        if quantizer_kwargs is None:
+            quantizer_kwargs = {
+                "input_dim": 768,
+                "rvq_dim": 512,
+                "output_dim": 768,
+                "num_quantizers": 32,
+                "codebook_size": 1024,
+                "codebook_dim": 8,
+                "quantizer_type": "rlfq",
+            }
+
+        # Handle quantizer_type from kwargs or config
+        kw_qtype = quantizer_kwargs.get("quantizer_type", None)
+        if kw_qtype is not None:
+            self.quantizer_type = kw_qtype
+        else:
+            self.quantizer_type = quantizer_type
+            quantizer_kwargs["quantizer_type"] = quantizer_type
+
+        self.quantizer_kwargs = quantizer_kwargs
+
+        super().__init__(**kwargs)
+
+    @property
+    def num_quantizers(self) -> int:
+        """Return the number of quantizers from quantizer_kwargs."""
+        return self.quantizer_kwargs.get("num_quantizers", 32)
+
+    @property
+    def codebook_size(self) -> int:
+        """Return the codebook size from quantizer_kwargs."""
+        return self.quantizer_kwargs.get("codebook_size", 4096)
+
+    @property
+    def frame_rate(self) -> float:
+        """Return the frame rate (tokens per second)."""
+        return self.sampling_rate / self.downsample_rate
+
+
+__all__ = ["MossAudioTokenizerConfig"]
diff --git a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py
new file mode 100644
index 000000000..1daa5081e
--- /dev/null
+++ b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py
@@ -0,0 +1,164 @@
+# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""CUDA graph wrapper for MOSS-TTS Local streaming codec decode."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from torch.cuda import CUDAGraph
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _CapturedStreamingDecodeGraph:
+    graph: CUDAGraph
+    static_codes: torch.Tensor
+    static_lengths: torch.Tensor
+    static_audio: torch.Tensor
+    static_audio_lengths: torch.Tensor
+
+
+class CUDAGraphStreamingDecoderWrapper:
+    """Replay fixed-size streaming ``_decode_frame`` steps with CUDA graphs.
+
+    The wrapped codec must already be inside ``codec.streaming(batch_size)``.
+    Graphs are keyed by the streaming step frame count ``T``. The batch width is
+    fixed to the codec stream slot count.
+    """
+
+    def __init__(
+        self,
+        codec: nn.Module,
+        *,
+        batch_size: int,
+        num_quantizers: int,
+        reset_streaming_state: Callable[[], None] | None = None,
+    ) -> None:
+        self.codec = codec
+        self.batch_size = int(batch_size)
+        self.num_quantizers = int(num_quantizers)
+        self.reset_streaming_state = reset_streaming_state
+        self.graphs: dict[int, _CapturedStreamingDecodeGraph] = {}
+        self._pool = None
+        self._warmed_up = False
+
+    @property
+    def is_ready(self) -> bool:
+        return bool(self.graphs)
+
+    @property
+    def capture_sizes(self) -> list[int]:
+        return sorted(self.graphs)
+
+    @torch.no_grad()
+    def warmup(self, device: torch.device, capture_sizes: list[int]) -> None:
+        if self._warmed_up:
+            return
+        self._warmed_up = True
+        if device.type != "cuda":
+            return
+        sizes = sorted({int(size) for size in capture_sizes if int(size) > 0}, reverse=True)
+        if not sizes:
+            return
+
+        logger.info(
+            "MOSS-TTS streaming decoder CUDA graph warmup: batch_size=%d n_vq=%d capture_sizes=%s",
+            self.batch_size,
+            self.num_quantizers,
+            list(reversed(sizes)),
+        )
+        start_s = time.perf_counter()
+        with torch.cuda.device(device):
+            for size in sizes:
+                try:
+                    self._capture(size, device)
+                    logger.info("  Captured MOSS-TTS streaming decoder CUDA graph for size=%d", size)
+                except Exception:
+                    self.graphs.pop(size, None)
+                    logger.warning(
+                        "  Failed to capture MOSS-TTS streaming decoder CUDA graph for size=%d; "
+                        "this size will use eager decode",
+                        size,
+                        exc_info=True,
+                    )
+        logger.info(
+            "MOSS-TTS streaming decoder CUDA graph warmup complete: %d/%d captured in %.1f ms",
+            len(self.graphs),
+            len(sizes),
+            (time.perf_counter() - start_s) * 1000.0,
+        )
+
+    @torch.no_grad()
+    def _capture(self, size: int, device: torch.device) -> None:
+        codes = torch.zeros(
+            self.num_quantizers,
+            self.batch_size,
+            size,
+            dtype=torch.long,
+            device=device,
+        )
+        lengths = torch.full((self.batch_size,), size, dtype=torch.long, device=device)
+        exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=device)
+
+        self.codec._set_streaming_exec_mask(exec_mask)
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            _ = self.codec._decode_frame(codes, lengths)
+        torch.cuda.current_stream().wait_stream(stream)
+        torch.accelerator.synchronize(device)
+
+        if self.reset_streaming_state is not None:
+            self.reset_streaming_state()
+        self.codec._set_streaming_exec_mask(exec_mask)
+        if self._pool is None:
+            self._pool = current_platform.get_global_graph_pool()
+        graph = CUDAGraph()
+        with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"):
+            output = self.codec._decode_frame(codes, lengths)
+
+        if output.audio is None or output.audio_lengths is None:
+            raise RuntimeError("MOSS-TTS streaming decoder CUDA graph capture produced empty audio.")
+
+        self.graphs[size] = _CapturedStreamingDecodeGraph(
+            graph=graph,
+            static_codes=codes,
+            static_lengths=lengths,
+            static_audio=output.audio,
+            static_audio_lengths=output.audio_lengths,
+        )
+
+    @torch.no_grad()
+    def decode(
+        self,
+        codes_step: torch.Tensor,
+        exec_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        if not codes_step.is_cuda or torch.cuda.is_current_stream_capturing():
+            return None
+        n_vq, batch_size, step_t = codes_step.shape
+        if int(n_vq) != self.num_quantizers or int(batch_size) != self.batch_size:
+            return None
+        # The streaming codec advances its internal KV offsets by tensor T, so
+        # bucket-padding a smaller step into a larger captured graph would write
+        # padded frames into state. Capture exact step sizes instead.
+        entry = self.graphs.get(int(step_t))
+        if entry is None:
+            return None
+
+        self.codec._set_streaming_exec_mask(exec_mask)
+        entry.static_codes.copy_(codes_step)
+        entry.graph.replay()
+        return entry.static_audio, entry.static_audio_lengths
+
+
+__all__ = ["CUDAGraphStreamingDecoderWrapper"]
diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py
index 7c5d47653..8f1a72e9b 100644
--- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py
+++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py
@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import contextlib
 from collections.abc import Iterable
 from typing import Any
 
@@ -19,6 +20,15 @@ from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import (
     MossAudioTokenizerConfig,
     MossAudioTokenizerModel,
 )
+from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import (
+    MossAudioTokenizerModel as MossAudioTokenizerV2Model,
+)
+from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import (
+    MossAudioTokenizerConfig as MossAudioTokenizerV2Config,
+)
+from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import (
+    CUDAGraphStreamingDecoderWrapper,
+)
 from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
     MossTTSCUDAGraphCodecWrapper,
 )
@@ -27,6 +37,121 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 logger = init_logger(__name__)
 
 
+class _MossCodecStreamSession:
+    """Persistent streaming decode session for vendored MOSS-Audio-Tokenizer-v2."""
+
+    def __init__(
+        self,
+        codec: nn.Module,
+        *,
+        stream_slots: int,
+        n_vq: int,
+        cudagraph_capture_sizes: list[int] | None = None,
+    ) -> None:
+        self._codec = codec
+        self._stream_slots = int(stream_slots)
+        self._batch_size = self._stream_slots
+        self._n_vq = int(n_vq)
+        self._device = next(codec.parameters()).device
+        self._free_stream_slots = list(range(self._stream_slots))
+        self._exit_stack = contextlib.ExitStack()
+        self._closed = False
+        with torch.no_grad():
+            self._exit_stack.enter_context(codec.streaming(self._batch_size))
+        self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None
+        capture_sizes = sorted({int(size) for size in (cudagraph_capture_sizes or []) if int(size) > 0})
+        if capture_sizes and self._device.type == "cuda":
+            self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper(
+                codec,
+                batch_size=self._batch_size,
+                num_quantizers=self._n_vq,
+                reset_streaming_state=lambda: self.reset_slots(list(range(self._batch_size))),
+            )
+            self._cudagraph_wrapper.warmup(self._device, capture_sizes)
+            self.reset_slots(list(range(self._batch_size)))
+            if not self._cudagraph_wrapper.is_ready:
+                self._cudagraph_wrapper = None
+
+    def acquire(self) -> int | None:
+        if not self._free_stream_slots:
+            return None
+        return self._free_stream_slots.pop()
+
+    def release(self, slot: int) -> None:
+        if self._closed:
+            return
+        self.reset_slots([slot])
+        self._free_stream_slots.append(slot)
+
+    def reset_slots(self, slots: list[int]) -> None:
+        if not slots:
+            return
+        reset_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device)
+        reset_mask[slots] = True
+
+        def _reset(module: nn.Module) -> None:
+            state = getattr(module, "_streaming_state", None)
+            if state is not None:
+                state.reset(reset_mask.to(state.device))
+
+        with torch.no_grad():
+            self._codec.apply(_reset)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        with torch.no_grad():
+            self._exit_stack.close()
+        self._closed = True
+
+    @torch.no_grad()
+    def step(self, slot_codes: dict[int, torch.Tensor]) -> dict[int, torch.Tensor]:
+        if not slot_codes:
+            return {}
+        step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()}
+        if len(step_lengths) != 1:
+            raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}")
+        (step_t,) = step_lengths
+        codes_step = torch.zeros(
+            self._n_vq,
+            self._batch_size,
+            step_t,
+            dtype=torch.long,
+            device=self._device,
+        )
+        codes_lengths = torch.zeros(self._batch_size, dtype=torch.long, device=self._device)
+        exec_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device)
+        for slot, codes in slot_codes.items():
+            codes_step[:, slot, :] = codes.to(device=self._device, dtype=torch.long)
+            codes_lengths[slot] = int(codes.shape[1])
+            exec_mask[slot] = True
+
+        graph_output: tuple[torch.Tensor, torch.Tensor] | None = None
+        if self._cudagraph_wrapper is not None:
+            graph_output = self._cudagraph_wrapper.decode(codes_step, exec_mask)
+
+        used_cudagraph = graph_output is not None
+        if used_cudagraph:
+            audio_tensor, lengths_tensor = graph_output
+        else:
+            self._codec._set_streaming_exec_mask(exec_mask)
+            result = self._codec._decode_frame(codes_step, codes_lengths)
+            if result.audio is None:
+                return {}
+            audio_tensor = result.audio
+            lengths_tensor = result.audio_lengths
+
+        audio = audio_tensor.detach().to("cpu", torch.float32)
+        lengths = lengths_tensor.detach().to("cpu") if lengths_tensor is not None else None
+        out: dict[int, torch.Tensor] = {}
+        for slot in slot_codes:
+            wav = audio[slot]
+            if lengths is not None:
+                wav = wav[..., : int(lengths[slot].item())]
+            out[slot] = wav.contiguous()
+        return out
+
+
 class MossTTSCodecDecoder(nn.Module):
     """Stage-1 decoder for all MOSS-TTS variants.
 
@@ -75,6 +200,19 @@ class MossTTSCodecDecoder(nn.Module):
         self._cuda_graph_wrapper: MossTTSCUDAGraphCodecWrapper | None = None
         self._n_channels: int = 1
         self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32)
+        self._stream_session: _MossCodecStreamSession | None = None
+        self._stream_slots: int = self._connector_int("codec_stream_slots", default=0)
+        self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0)
+        self._stream_max_step_frames: int = self._stream_chunk_frames or 100
+        self._stream_req_slots: dict[str, int] = {}
+        self._stream_pending_codes: dict[str, list[torch.Tensor]] = {}
+        self._stream_starved_reqs: set[str] = set()
+        self._codec_streaming: bool = self._connector_bool("codec_streaming", default=False)
+        default_graph_capture_max = self._stream_chunk_frames or min(self._stream_max_step_frames, 32)
+        default_graph_capture_sizes = list(range(1, max(1, default_graph_capture_max) + 1))
+        self._streaming_cudagraph_capture_sizes = self._streaming_cudagraph_capture_sizes_from_compilation_config(
+            default_graph_capture_sizes
+        )
 
     # ------------------------------------------------------------------
     # vLLM-Omni stubs (codec has no AR loop)
@@ -112,9 +250,9 @@ class MossTTSCodecDecoder(nn.Module):
         chunk transfer adapter assigns those to ``request.prompt_token_ids``,
         which arrives here as ``input_ids`` concatenated across all requests.
         Per-request slice boundaries are computed from
-        ``kwargs["num_scheduled_tokens"]`` (a list of token counts, one per
-        request).  ``runtime_additional_information`` carries per-request
-        metadata such as ``left_context_size``.
+        ``kwargs["seq_token_counts"]`` (token counts, one per request).
+        ``runtime_additional_information`` carries per-request metadata such as
+        ``left_context_size``.
 
         Returns
         -------
@@ -123,8 +261,8 @@ class MossTTSCodecDecoder(nn.Module):
           multimodal_outputs["sr"]            — list of scalar int32 tensors
         """
         sr_tensor = self._sr_tensor
-        empty = torch.zeros((0,), dtype=torch.float32)
-        info_list: list[dict[str, Any]] = runtime_additional_information or [{}]
+        empty = self._empty_audio()
+        info_list: list[dict[str, Any]] = list(runtime_additional_information or [{}])
         num_req = max(len(info_list), 1)
 
         if self._codec is None:
@@ -137,27 +275,59 @@ class MossTTSCodecDecoder(nn.Module):
                 },
             )
 
+        if self._codec_streaming and runtime_additional_information is None:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "model_outputs": [empty] * num_req,
+                    "sr": [sr_tensor] * num_req,
+                },
+            )
+
         audios: list[torch.Tensor] = [empty] * num_req
         srs: list[torch.Tensor] = [sr_tensor] * num_req
         device = next(self._codec.parameters()).device
+        streaming_work: list[tuple[int, str, torch.Tensor, bool]] = []
 
         if input_ids is None or input_ids.numel() == 0:
+            for i, wav in self._finish_empty_streaming_requests(info_list).items():
+                audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
             return OmniOutput(
                 text_hidden_states=None,
                 multimodal_outputs={"model_outputs": audios, "sr": srs},
             )
 
-        # ``input_ids`` is concatenated across all requests; split by
-        # query_start_loc-style offsets from ``num_scheduled_tokens`` in
-        # kwargs if available, else assume one request.
+        # ``input_ids`` is concatenated across all requests. vLLM-Omni runners
+        # pass the per-request lengths via the shared code2wav contract
+        # ``seq_token_counts``.
         ids_flat = input_ids.reshape(-1).to(dtype=torch.long)
-        num_scheduled_tokens = kwargs.get("num_scheduled_tokens")
-        if isinstance(num_scheduled_tokens, (list, tuple)) and len(num_scheduled_tokens) == num_req:
-            offsets = [0]
-            for n in num_scheduled_tokens:
-                offsets.append(offsets[-1] + int(n))
-        else:
-            offsets = [0, int(ids_flat.shape[0])]
+        token_counts = self._normalize_seq_token_counts(kwargs.get("seq_token_counts"))
+        if token_counts is None:
+            raise RuntimeError(
+                "MossTTS codec requires seq_token_counts; otherwise concatenated "
+                "codec tokens cannot be split per request."
+            )
+        if sum(token_counts) != int(ids_flat.shape[0]):
+            raise RuntimeError(
+                "MossTTS codec seq_token_counts mismatch: "
+                f"counts={token_counts}, sum={sum(token_counts)}, input_tokens={int(ids_flat.shape[0])}."
+            )
+
+        num_req = len(token_counts)
+        if len(info_list) < num_req:
+            info_list.extend({} for _ in range(num_req - len(info_list)))
+        elif len(info_list) > num_req:
+            info_list = info_list[:num_req]
+        if len(audios) < num_req:
+            audios.extend(empty for _ in range(num_req - len(audios)))
+            srs.extend(sr_tensor for _ in range(num_req - len(srs)))
+        elif len(audios) > num_req:
+            audios = audios[:num_req]
+            srs = srs[:num_req]
+
+        offsets = [0]
+        for n in token_counts:
+            offsets.append(offsets[-1] + int(n))
 
         for i, info in enumerate(info_list):
             if i + 1 >= len(offsets):
@@ -165,6 +335,14 @@ class MossTTSCodecDecoder(nn.Module):
             seg = ids_flat[offsets[i] : offsets[i + 1]]
             if seg.numel() == 0:
                 continue
+            meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
+            finished = bool(meta.get("stream_finished", meta.get("finished", False)))
+            streaming_enabled = bool(meta.get("codec_streaming", self._codec_streaming))
+            code_flat_numel = meta.get("code_flat_numel")
+            if streaming_enabled and finished and code_flat_numel is not None and int(code_flat_numel) == 0:
+                for _, wav in self._finish_empty_streaming_requests([info]).items():
+                    audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
+                continue
             if seg.numel() % self._n_vq != 0:
                 logger.warning(
                     "MossTTS codec input length %d not divisible by n_vq %d; skipping.",
@@ -181,14 +359,19 @@ class MossTTSCodecDecoder(nn.Module):
             codebook_size = self._codec.config.codebook_size
             codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1)
 
-            meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {}
-            left_ctx = self._runtime_value(info, meta, "left_context_size", 0)
+            left_ctx = meta.get("left_context_size", 0)
             if isinstance(left_ctx, (list, tuple)):
                 left_ctx = int(left_ctx[0]) if left_ctx else 0
             elif isinstance(left_ctx, torch.Tensor):
                 left_ctx = int(left_ctx.reshape(-1)[0].item()) if left_ctx.numel() else 0
             left_ctx = int(left_ctx)
 
+            req_key = self._runtime_request_key(info, meta, i)
+
+            if streaming_enabled:
+                streaming_work.append((i, req_key, codes_nq_t, finished))
+                continue
+
             if self._cuda_graph_wrapper is not None:
                 out = self._cuda_graph_wrapper.decode(codes_nq_t)
             else:
@@ -216,21 +399,273 @@ class MossTTSCodecDecoder(nn.Module):
                     )
                 wav = wav[..., trim:]
 
-            audios[i] = wav.reshape(-1) if self._n_channels == 1 else wav
+            audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
+
+        if streaming_work:
+            for i, wav in self._decode_streaming_batch(streaming_work).items():
+                audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav
 
         return OmniOutput(
             text_hidden_states=None,
             multimodal_outputs={"model_outputs": audios, "sr": srs},
         )
 
+    def _finish_empty_streaming_requests(self, info_list: list[dict[str, Any]]) -> dict[int, torch.Tensor]:
+        """Release codec stream state for empty finish sentinels.
+
+        Stage-0 can finish on a step that emits no new audio frame. The stage
+        input processor forwards that as an empty payload with finished=true.
+        If a request never acquired a stream slot, do not offline-decode its
+        buffered codes here: streaming requests must only emit client deltas.
+        """
+        session = self._stream_session
+        outputs: dict[int, torch.Tensor] = {}
+        if session is None:
+            return outputs
+        for i, info in enumerate(info_list):
+            if not isinstance(info, dict):
+                continue
+            meta = (info.get("meta", {}) or {}) if isinstance(info.get("meta", {}), dict) else {}
+            if not bool(meta.get("codec_streaming", self._codec_streaming)):
+                continue
+            finished = bool(meta.get("stream_finished", meta.get("finished", False)))
+            if not finished:
+                continue
+            req_key = self._runtime_request_key(info, meta, i)
+            slot = self._stream_req_slots.get(req_key)
+            pending = req_key in self._stream_pending_codes
+            if slot is not None or pending:
+                if pending and slot is None:
+                    logger.warning(
+                        "MOSS codec stream request %s finished before a codec stream slot became available; "
+                        "dropping buffered codes instead of offline-decoding a non-delta waveform.",
+                        req_key,
+                    )
+                self._finish_stream_request(req_key, session, slot)
+        return outputs
+
     @staticmethod
-    def _runtime_value(info: Any, meta: dict[str, Any], name: str, default: Any = None) -> Any:
-        if name in meta:
-            return meta[name]
-        if isinstance(info, dict) and name in info:
-            return info[name]
+    def _normalize_seq_token_counts(value: Any) -> list[int] | None:
+        if value is None:
+            return None
+        if not isinstance(value, (list, tuple)):
+            raise TypeError(
+                "MossTTS codec expects seq_token_counts to be a list/tuple of per-request token counts, "
+                f"got {type(value).__name__}."
+            )
+        counts = [int(item) for item in value]
+        if not counts:
+            return None
+        for count in counts:
+            if count < 0:
+                raise ValueError(f"MossTTS codec seq_token_counts must be non-negative, got {counts}.")
+        return counts
+
+    def _runtime_request_key(self, info: Any, meta: dict[str, Any], index: int) -> str:
+        for value in (
+            meta.get("req_id"),
+            info.get("request_id") if isinstance(info, dict) else None,
+        ):
+            if isinstance(value, (list, tuple)):
+                value = value[0] if value else None
+            if value is not None:
+                return str(value)
+        return f"moss-codec-stream-{index}"
+
+    def _empty_audio(self) -> torch.Tensor:
+        if self._n_channels > 1:
+            return torch.zeros((self._n_channels, 0), dtype=torch.float32)
+        return torch.zeros((0,), dtype=torch.float32)
+
+    def _ensure_stream_session(self) -> _MossCodecStreamSession | None:
+        if self._codec is None:
+            return None
+        if self._stream_session is not None:
+            return self._stream_session
+        slots = self._stream_slots
+        if slots <= 0:
+            scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None)
+            slots = int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1)
+        self._stream_session = _MossCodecStreamSession(
+            self._codec,
+            stream_slots=max(1, slots),
+            n_vq=self._n_vq,
+            cudagraph_capture_sizes=self._streaming_cudagraph_capture_sizes,
+        )
+        return self._stream_session
+
+    def _decode_streaming_batch(
+        self,
+        items: list[tuple[int, str, torch.Tensor, bool]],
+    ) -> dict[int, torch.Tensor]:
+        session = self._ensure_stream_session()
+        if session is None:
+            return {}
+
+        outputs: dict[int, torch.Tensor] = {}
+        grouped: dict[int, list[tuple[int, str, int, torch.Tensor, bool]]] = {}
+        max_step_frames = max(1, int(self._stream_max_step_frames))
+
+        for output_index, request_id, codes_nq_t, finished in items:
+            pending = self._stream_pending_codes.get(request_id)
+            slot = self._stream_req_slots.get(request_id)
+            if slot is None:
+                slot = session.acquire()
+                if slot is None:
+                    self._append_stream_pending(request_id, codes_nq_t)
+                    if request_id not in self._stream_starved_reqs:
+                        logger.warning(
+                            "MOSS codec streaming slots exhausted; buffering %s until a stream slot is available.",
+                            request_id,
+                        )
+                        self._stream_starved_reqs.add(request_id)
+                    if finished:
+                        logger.warning(
+                            "MOSS codec stream request %s finished before a codec stream slot became available; "
+                            "dropping buffered codes instead of offline-decoding a non-delta waveform.",
+                            request_id,
+                        )
+                        self._finish_stream_request(request_id, session, None)
+                    continue
+                self._stream_req_slots[request_id] = slot
+
+            if pending:
+                self._append_stream_pending(request_id, codes_nq_t)
+                replay_codes = self._pop_stream_pending(request_id)
+                wav = self._decode_stream_slot_sequence(session, slot, replay_codes)
+                if wav is not None:
+                    outputs[output_index] = wav
+                if finished:
+                    self._finish_stream_request(request_id, session, slot)
+                continue
+
+            if int(codes_nq_t.shape[1]) > max_step_frames:
+                wav = self._decode_stream_slot_sequence(session, slot, codes_nq_t)
+                if wav is not None:
+                    outputs[output_index] = wav
+                if finished:
+                    self._finish_stream_request(request_id, session, slot)
+                continue
+
+            grouped.setdefault(int(codes_nq_t.shape[1]), []).append(
+                (output_index, request_id, slot, codes_nq_t, finished)
+            )
+
+        for group in grouped.values():
+            plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group}
+            decoded = session.step(plan)
+            for output_index, request_id, slot, _, finished in group:
+                wav = decoded.get(slot)
+                if wav is not None:
+                    outputs[output_index] = wav
+                if finished:
+                    self._finish_stream_request(request_id, session, slot)
+
+        return outputs
+
+    def _append_stream_pending(self, request_id: str, codes_nq_t: torch.Tensor) -> None:
+        self._stream_pending_codes.setdefault(request_id, []).append(
+            codes_nq_t.detach().to("cpu", torch.long).contiguous()
+        )
+
+    def _pop_stream_pending(self, request_id: str) -> torch.Tensor:
+        pending = self._stream_pending_codes.pop(request_id, [])
+        if not pending:
+            return torch.empty((self._n_vq, 0), dtype=torch.long)
+        return torch.cat(pending, dim=1).contiguous()
+
+    def _decode_stream_slot_sequence(
+        self,
+        session: _MossCodecStreamSession,
+        slot: int,
+        codes_nq_t: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if codes_nq_t.numel() == 0:
+            return None
+        max_step_frames = max(1, int(self._stream_max_step_frames))
+        parts: list[torch.Tensor] = []
+        for start in range(0, int(codes_nq_t.shape[1]), max_step_frames):
+            chunk = codes_nq_t[:, start : start + max_step_frames]
+            decoded = session.step({slot: chunk})
+            wav = decoded.get(slot)
+            if wav is not None:
+                parts.append(wav)
+        if not parts:
+            return None
+        return torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0]
+
+    def _finish_stream_request(
+        self,
+        request_id: str,
+        session: _MossCodecStreamSession,
+        slot: int | None,
+    ) -> None:
+        if slot is not None:
+            session.release(slot)
+        self._stream_req_slots.pop(request_id, None)
+        self._stream_pending_codes.pop(request_id, None)
+        self._stream_starved_reqs.discard(request_id)
+
+    def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None:
+        """Release codec streaming slots when requests finish outside payload flow.
+
+        Normal streaming completion releases slots from ``_decode_streaming_batch``
+        when the Stage-0 payload carries ``finished=True``. Client disconnects
+        and engine-side aborts can finish a request without delivering that
+        terminal payload, so the runner calls this hook from its finished-request
+        path to avoid leaking stream slots and buffered codes.
+        """
+        session = self._stream_session
+        for req_id in finished_req_ids:
+            request_id = str(req_id)
+            slot = self._stream_req_slots.get(request_id)
+            has_state = (
+                slot is not None or request_id in self._stream_pending_codes or request_id in self._stream_starved_reqs
+            )
+            if not has_state:
+                continue
+            if session is not None:
+                self._finish_stream_request(request_id, session, slot)
+            else:
+                self._stream_req_slots.pop(request_id, None)
+                self._stream_pending_codes.pop(request_id, None)
+                self._stream_starved_reqs.discard(request_id)
+
+    def _connector_int(self, name: str, default: int = 0) -> int:
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        connector_cfg = getattr(model_cfg, "stage_connector_config", None)
+        if isinstance(connector_cfg, dict):
+            extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg)
+        else:
+            extra_cfg = getattr(connector_cfg, "extra", None)
+        if isinstance(extra_cfg, dict) and name in extra_cfg:
+            return int(extra_cfg[name])
         return default
 
+    def _connector_bool(self, name: str, default: bool = False) -> bool:
+        model_cfg = getattr(self.vllm_config, "model_config", None)
+        connector_cfg = getattr(model_cfg, "stage_connector_config", None)
+        if isinstance(connector_cfg, dict):
+            extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg)
+        else:
+            extra_cfg = getattr(connector_cfg, "extra", None)
+        if isinstance(extra_cfg, dict) and name in extra_cfg:
+            return bool(extra_cfg[name])
+        return default
+
+    def _streaming_cudagraph_capture_sizes_from_compilation_config(
+        self,
+        default: list[int],
+    ) -> list[int]:
+        if getattr(self.vllm_config.model_config, "enforce_eager", True):
+            return []
+        compilation_config = getattr(self.vllm_config, "compilation_config", None)
+        capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None)
+        if capture_sizes:
+            max_step_frames = max(1, int(self._stream_max_step_frames))
+            return sorted({int(size) for size in capture_sizes if 0 < int(size) <= max_step_frames})
+        return list(default)
+
     # ------------------------------------------------------------------
     # Weight loading
     # ------------------------------------------------------------------
@@ -272,14 +707,12 @@ class MossTTSCodecDecoder(nn.Module):
             (".self_attn.out_projs.0.", ".attn.out_proj."),
             (".linear1.", ".ff1."),
             (".linear2.", ".ff2."),
-            # v2 (MOSS-Audio-Tokenizer-v2) uses different submodule names for
-            # the same attention/FFN sublayers — no trailing "s"/index on
-            # in_proj/out_proj, and an `ffn.{0,2}` Sequential instead of
-            # separate linear1/linear2 attributes.
-            (".self_attn.in_proj.", ".attn.in_proj."),
-            (".self_attn.out_proj.", ".attn.out_proj."),
-            (".ffn.0.", ".ff1."),
-            (".ffn.2.", ".ff2."),
+            # v2 checkpoint names use singular in_proj/out_proj and ffn.{0,2};
+            # the vendored module keeps the original MOSS layer names.
+            (".self_attn.in_proj.", ".self_attn.in_projs.0."),
+            (".self_attn.out_proj.", ".self_attn.out_projs.0."),
+            (".ffn.0.", ".linear1."),
+            (".ffn.2.", ".linear2."),
             (".layer_scale_1.", ".ls1."),
             (".layer_scale_2.", ".ls2."),
             (".input_proj.", ".in_proj."),
@@ -292,21 +725,39 @@ class MossTTSCodecDecoder(nn.Module):
                     return name.replace(src, dst)
             return name
 
-        loaded_total = 0
+        loaded_names: set[str] = set()
         skipped: list[str] = []
+        shape_mismatches: list[tuple[str, str, tuple[int, ...], tuple[int, ...]]] = []
         for name, tensor in codec_weights:
             # Try direct name first (e.g. ``quantizer.input_proj.*`` exists
             # under the same name in both layouts), then the remap (transformer
             # submodules need ``.linear1.``→``.ff1.`` etc.).
             tgt = name if name in params_dict else _remap(name)
             if tgt in params_dict:
+                expected_shape = tuple(params_dict[tgt].shape)
+                actual_shape = tuple(tensor.shape)
+                if expected_shape != actual_shape:
+                    shape_mismatches.append((name, tgt, actual_shape, expected_shape))
+                    continue
                 default_weight_loader(params_dict[tgt], tensor)
-                loaded_total += 1
+                loaded_names.add(tgt)
             else:
                 skipped.append(name)
+
+        missing = sorted(set(params_dict) - loaded_names)
+        if missing or skipped or shape_mismatches:
+            raise RuntimeError(
+                "MOSS Audio Tokenizer weights were not fully loaded: "
+                f"loaded={len(loaded_names)}/{len(params_dict)} "
+                f"missing={len(missing)} skipped={len(skipped)} "
+                f"shape_mismatches={len(shape_mismatches)}; "
+                f"first_missing={missing[:5]} "
+                f"first_skipped={skipped[:5]} "
+                f"first_shape_mismatches={shape_mismatches[:3]}"
+            )
         logger.info(
             "MOSS Audio Tokenizer weights: loaded=%d/%d skipped=%d (first skipped: %s)",
-            loaded_total,
+            len(loaded_names),
             len(params_dict),
             len(skipped),
             skipped[:3] if skipped else "none",
@@ -316,7 +767,15 @@ class MossTTSCodecDecoder(nn.Module):
         codec.to(device=device, dtype=torch.float32)
         codec.eval()
         self._codec = codec
-        self._n_channels = int(getattr(codec_cfg, "number_channels", 1) or 1)
+        inferred_channels = 2 if "v2" in codec_path.lower() else 1
+        self._n_channels = int(
+            getattr(
+                codec_cfg,
+                "number_channels",
+                getattr(codec_cfg, "num_channels", inferred_channels),
+            )
+            or inferred_channels
+        )
         self._sr_tensor = torch.tensor(int(codec_cfg.sampling_rate), dtype=torch.int32)
 
         logger.info(
@@ -326,7 +785,9 @@ class MossTTSCodecDecoder(nn.Module):
             self._n_channels,
         )
 
-        self._maybe_enable_decoder_cudagraph(device)
+        self._configure_decoder_cudagraph(device)
+        if self._codec_streaming and self._streaming_cudagraph_capture_sizes:
+            self._ensure_stream_session()
 
         # vLLM's track_weights_loading() compares the returned set against
         # ``self.named_parameters()``. After ``self._codec = codec`` above,
@@ -336,28 +797,43 @@ class MossTTSCodecDecoder(nn.Module):
 
     def _build_codec(self, codec_path: str) -> tuple[Any, nn.Module]:
         try:
-            from transformers import AutoConfig, AutoModel
-
-            codec_cfg = AutoConfig.from_pretrained(codec_path, trust_remote_code=True)
-            codec = AutoModel.from_config(codec_cfg, trust_remote_code=True)
-            logger.info("Using MOSS Audio Tokenizer remote-code classes from %s", codec_path)
+            codec_cfg = MossAudioTokenizerV2Config.from_pretrained(codec_path)
+            codec = MossAudioTokenizerV2Model(codec_cfg)
+            logger.info("Using vendored MOSS Audio Tokenizer v2 classes from %s", codec_path)
             return codec_cfg, codec
         except Exception:
             logger.exception(
-                "Failed to instantiate official MOSS Audio Tokenizer via HF remote code; "
-                "falling back to vendored codec."
+                "Failed to instantiate vendored MOSS Audio Tokenizer v2; falling back to legacy vendored codec."
             )
 
         codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path)
         codec = MossAudioTokenizerModel(codec_cfg)
         return codec_cfg, codec
 
-    def _maybe_enable_decoder_cudagraph(self, device: torch.device) -> None:
-        """Capture CUDA Graphs for the codec decoder if enforce_eager is False."""
+    def _configure_decoder_cudagraph(self, device: torch.device) -> None:
+        """Select the codec CUDA Graph path.
+
+        ``enforce_eager`` is the single graph on/off switch. If graphing is
+        enabled, ``codec_streaming`` decides whether decode uses the persistent
+        streaming-state wrapper or the offline full-chunk wrapper.
+        """
         if getattr(self.vllm_config.model_config, "enforce_eager", True):
+            self._streaming_cudagraph_capture_sizes = []
             return
         if self._codec is None:
             return
+        if self._codec_streaming:
+            logger.info(
+                "MOSS-TTS codec CUDA Graph selected streaming wrapper: capture_sizes=%s",
+                self._streaming_cudagraph_capture_sizes,
+            )
+            return
+
+        self._enable_non_streaming_decoder_cudagraph(device)
+
+    def _enable_non_streaming_decoder_cudagraph(self, device: torch.device) -> None:
+        if self._codec is None:
+            return
 
         # Read capture sizes from the connector's extra config (same convention
         # as Qwen3-TTS), falling back to a sensible default covering common
diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py
index bf01a83bf..2fd0beefd 100644
--- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py
+++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py
@@ -126,6 +126,7 @@ def _sample_token(
     top_k: int,
     top_p: float,
     do_sample: bool,
+    generator: torch.Generator | None = None,
 ) -> torch.Tensor:
     """Top-k + top-p sampling (matches upstream's ``sample_token`` for the
     inference branch).
@@ -152,7 +153,7 @@ def _sample_token(
 
     probs = F.softmax(logits, dim=-1)
     flat = probs.reshape(-1, probs.shape[-1])
-    sampled = torch.multinomial(flat, num_samples=1).reshape(probs.shape[:-1])
+    sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1])
     return sampled
 
 
diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py
index cb2e95551..e0d871b46 100644
--- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py
+++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py
@@ -160,6 +160,7 @@ class MossTTSLocalDepthTransformer(nn.Module):
         text_top_p: float = 1.0,
         repetition_penalty: float = 1.0,
         history_per_codebook: list[list[int]] | None = None,
+        generator: torch.Generator | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Generate one audio frame for batch B.
 
@@ -183,7 +184,16 @@ class MossTTSLocalDepthTransformer(nn.Module):
         local_hidden = hidden[:, 0, :]
 
         binary_logits = local_text_lm_head(local_hidden).float()
-        binary_choice = _sample_token(binary_logits, text_temperature, text_top_k, text_top_p, do_sample)
+        # This is a binary continue/stop gate. The checkpoint expects sampling
+        # here; greedy argmax is biased toward "continue" and may never stop.
+        binary_choice = _sample_token(
+            binary_logits,
+            text_temperature,
+            text_top_k,
+            text_top_p,
+            do_sample,
+            generator=generator,
+        )
         should_continue = binary_choice.eq(0)
         import os as _os
 
@@ -209,7 +219,14 @@ class MossTTSLocalDepthTransformer(nn.Module):
                     pos = sel > 0
                     sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty)
                     channel_logits.index_copy_(-1, hist_t, sel)
-            channel_token = _sample_token(channel_logits, temperature, top_k, top_p, do_sample)
+            channel_token = _sample_token(
+                channel_logits,
+                temperature,
+                top_k,
+                top_p,
+                do_sample,
+                generator=generator,
+            )
             codes[:, channel_index] = channel_token
 
             if channel_index + 1 < n_vq:
diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
index 9f4037e4b..9758c922e 100644
--- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
+++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py
@@ -1139,26 +1139,19 @@ class MossTTSRealtimeTalkerForGeneration(nn.Module):
 
                 info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
                 state["step"] = int(state.get("step", 0)) + 1
-                per_req_codes[i] = updated_acc
+                per_req_codes[i] = new_codes.unsqueeze(0)
                 have_codes = True
 
         if not have_codes:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
 
-        # Emit one accumulated-codes tensor per request as a list; the splitter
-        # routes ``element[idx]`` to request idx, so each Stage-1 request decodes
-        # its own codes instead of request 0's.
-        # The realtime variant emits raw codes (no delay pattern), so we signal
-        # that to the chunk processor via a 1-D bool tensor in ``meta.finished``
-        # (already a tensor field on the schema). The processor reads its truth
-        # value as "skip de-delay" only for realtime; the delay path doesn't
-        # populate ``meta`` at all.
-        device = hidden.device
+        # Emit one current-frame tensor per request as a list; the splitter
+        # routes ``element[idx]`` to request idx, so each Stage-1 request receives
+        # only its own newly generated raw codec row.
         return OmniOutput(
             text_hidden_states=hidden,
             multimodal_outputs={
                 "codes": {"audio": per_req_codes},
-                "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
             },
         )
 
@@ -1266,7 +1259,7 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
     * Per-step audio generation runs ``MossTTSLocalDepthTransformer``
       (1-layer GPT2-style local transformer, interleaved RoPE, ``n_vq=12``
       codebooks -- structurally different from Realtime's Qwen3-style
-      ``CodePredictorBaseModel``) inside ``make_omni_output``. Stop
+      ``CodePredictorBaseModel``) in the runner-owned ``talker_mtp`` path. Stop
       condition: the local transformer's binary head selects the
       "audio_end" candidate (logits index 1).
     """
@@ -1327,6 +1320,8 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         # Stacked embedding cache built after load_weights() for vectorised
         # additive-fusion gathers (avoids an n_vq-iteration Python loop).
         self._stacked_audio_emb_w: torch.Tensor | None = None
+        self.mtp_hidden_size = hidden_size
+        self.talker_mtp_accepts_req_infos = True
 
         self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
             ("audio_codes", "current"),
@@ -1384,6 +1379,15 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
             num_rows,
         ):
             if state is not None and state.get("is_stopping"):
+                forced_im_end_count = int(state.get("_forced_im_end_count", 0))
+                if forced_im_end_count == 1:
+                    logger.warning(
+                        "MOSS-TTS Local request is still running after "
+                        "im_end_token_id=%d was forced once; verify stage-0 "
+                        "stop_token_ids/eos includes this id.",
+                        self.im_end_token_id,
+                    )
+                state["_forced_im_end_count"] = forced_im_end_count + 1
                 logits[row_start:row_end, self.im_end_token_id] = 0.0
             else:
                 logits[row_start:row_end, self.audio_assistant_slot_token_id] = 0.0
@@ -1455,38 +1459,159 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
                 else torch.full((self.n_vq,), self.audio_pad_token_id, dtype=torch.long, device=device)
             )
 
-            max_new_frames_req = info_dict.get("max_new_frames")
-            if isinstance(max_new_frames_req, (list, tuple)) and max_new_frames_req:
-                max_new_frames_req = max_new_frames_req[0]
-            try:
-                max_new_frames = int(max_new_frames_req) if max_new_frames_req is not None else -1
-            except (TypeError, ValueError):
-                max_new_frames = -1
+            max_new_frames = info_dict.get("max_new_frames", [-1])[0]
 
             info_update: dict[str, Any] = {
-                "audio_state": {"is_stopping": False, "step": 0, "max_new_frames": max_new_frames},
-                "audio_codes": {"current": current_codes},
+                "audio_state": {
+                    "is_stopping": False,
+                    "step": 0,
+                    "max_new_frames": max_new_frames,
+                },
+                "audio_codes": {"current": current_codes, "emit": False},
                 "ref_offset": ref_offset + span_len,
             }
             return input_ids, embeds, info_update
 
-        # Decode step: input_ids is the forced text token from compute_logits.
+        # Decode step: return the plain forced text-token embedding. The local
+        # depth transformer runs in OmniGPUModelRunner.talker_mtp before the
+        # backbone forward; it adds the freshly sampled audio-code embedding
+        # and stores the generated code frame in model_intermediate_buffer.
         text_embed = self.model.embed_tokens(input_ids.reshape(1))  # (1, H)
-        prev_codes = (info_dict.get("audio_codes", {}) or {}).get("current")
-        if isinstance(prev_codes, torch.Tensor) and prev_codes.numel() == self.n_vq:
-            codes = prev_codes.to(device=device, dtype=torch.long).unsqueeze(0)  # (1, n_vq)
-            combined = text_embed + self._audio_embed(codes)
+        hs = info_dict.get("hidden_states", {}) or {}
+        last_hidden = hs.get("last")
+        if isinstance(last_hidden, torch.Tensor) and last_hidden.numel() > 0:
+            mtp_hidden = last_hidden.to(device=device, dtype=text_embed.dtype).reshape(1, -1)
         else:
-            combined = text_embed
-        return input_ids, combined, {}
+            mtp_hidden = torch.zeros((1, self.hidden_size), device=device, dtype=text_embed.dtype)
+        return input_ids, text_embed, {"mtp_inputs": (mtp_hidden, torch.zeros_like(mtp_hidden))}
 
     def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]:
         if hidden_states.numel() == 0:
             return {}
         return {"hidden_states": {"last": hidden_states[-1].detach()}}
 
+    @torch.inference_mode()
+    def talker_mtp(
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        text_step: torch.Tensor,
+        do_sample: bool | None = None,
+        temperature: float | None = None,
+        top_k: int | None = None,
+        top_p: float | None = None,
+        generator: torch.Generator | None = None,
+        req_infos: list[dict[str, Any]] | None = None,
+        **_: Any,
+    ) -> tuple[torch.Tensor, None]:
+        bsz = int(input_embeds.shape[0])
+        dev = input_embeds.device
+        input_embeds_out = input_embeds.reshape(bsz, -1).clone()
+        req_infos = req_infos or [{} for _ in range(bsz)]
+
+        def write_update(info: dict[str, Any], update: dict[str, Any]) -> None:
+            for key, value in update.items():
+                if isinstance(value, dict):
+                    target = info.setdefault(key, {})
+                    target.update(value)
+                else:
+                    info[key] = value
+
+        # Match MOSS-TTS Local v1.5's model-card/SGLang defaults. The generic
+        # stage SamplingParams are for the backbone token loop and should not
+        # silently retune the local depth transformer's audio-channel sampler.
+        audio_temperature = 1.7
+        audio_top_k = 25
+        audio_top_p = 0.8
+        do_sample_val = True if do_sample is None else bool(do_sample)
+
+        for i in range(bsz):
+            info = req_infos[i] if i < len(req_infos) and isinstance(req_infos[i], dict) else {}
+            state = dict(info.get("audio_state", {}) or {})
+            step = int(state.get("step", 0))
+            max_new_frames = int(state.get("max_new_frames", -1))
+            acc = (info.get("audio_codes", {}) or {}).get("accumulated")
+
+            if state.get("is_stopping"):
+                write_update(
+                    info,
+                    {
+                        "audio_state": state,
+                        "audio_codes": {
+                            "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long),
+                            "emit": False,
+                        },
+                    },
+                )
+                continue
+
+            rep_window = 50
+            if isinstance(acc, torch.Tensor) and acc.numel() > 0:
+                tail = acc[-rep_window:].long().cpu().tolist()
+                hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)]
+            else:
+                hist_per_cb = [[] for _ in range(self.n_vq)]
+
+            should_continue_t, new_codes_b = self.local_transformer.generate_frame(
+                last_talker_hidden[i : i + 1],
+                self.audio_lm_heads,
+                self.audio_embeddings,
+                self.local_text_lm_head,
+                n_vq=self.n_vq,
+                do_sample=do_sample_val,
+                temperature=audio_temperature,
+                top_k=audio_top_k,
+                top_p=audio_top_p,
+                repetition_penalty=1.0,
+                history_per_codebook=hist_per_cb,
+                generator=generator,
+            )
+            should_continue = bool(should_continue_t[0].item())
+            new_codes = new_codes_b.squeeze(0).to(device=dev, dtype=torch.long)
+
+            if not should_continue or (0 <= max_new_frames <= step):
+                state["is_stopping"] = True
+                state["step"] = step + 1
+                write_update(
+                    info,
+                    {
+                        "audio_state": state,
+                        "audio_codes": {
+                            "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long),
+                            "accumulated": acc,
+                            "emit": False,
+                        },
+                    },
+                )
+                continue
+
+            if isinstance(acc, torch.Tensor) and acc.numel() > 0:
+                updated_acc = torch.cat([acc.to(dev), new_codes.unsqueeze(0)], dim=0)
+            else:
+                updated_acc = new_codes.unsqueeze(0)
+
+            state["step"] = step + 1
+            current = new_codes.unsqueeze(0)
+            input_embeds_out[i : i + 1] = input_embeds_out[i : i + 1] + self._audio_embed(current).to(
+                dtype=input_embeds_out.dtype
+            )
+            write_update(
+                info,
+                {
+                    "audio_state": state,
+                    "audio_codes": {
+                        "current": current,
+                        "accumulated": updated_acc,
+                        "emit": True,
+                    },
+                },
+            )
+
+        return input_embeds_out, None
+
     # ------------------------------------------------------------------
-    # Per-step audio generation via local transformer
+    # Package runner-generated audio frames
     # ------------------------------------------------------------------
 
     def make_omni_output(
@@ -1511,91 +1636,41 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
         self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts]
         self._batch_state_spans = kwargs.get("request_token_spans")
 
-        spans = kwargs.get("request_token_spans")
-        per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts]
+        per_req_codes: list[torch.Tensor] = []
         have_codes = False
-
-        if hidden.numel() > 0 and info_dicts:
-            num_rows = hidden.shape[0]
-            if spans is None or len(spans) != len(info_dicts):
-                raise RuntimeError(
-                    "MOSS-TTS make_omni_output requires request_token_spans "
-                    f"for {len(info_dicts)} request infos, got {0 if spans is None else len(spans)}"
-                )
-            for i, info in enumerate(info_dicts):
-                if not isinstance(info, dict):
-                    continue
-                row_start, row_end = spans[i]
-                row_end = min(int(row_end), num_rows)
-                if row_start >= row_end:
-                    continue
-
-                state = info.get("audio_state", {}) or {}
-                if state.get("is_stopping"):
-                    continue  # already stopped -- no more audio frames
-
-                last_h = hidden[row_end - 1].unsqueeze(0)  # (1, H)
-
-                # Sampling parameters mirror Realtime's defaults; repetition
-                # penalty applies over a 50-frame rolling window per codebook.
-                rep_window = 50
-                acc_for_hist = (info.get("audio_codes", {}) or {}).get("accumulated")
-                if isinstance(acc_for_hist, torch.Tensor) and acc_for_hist.numel() > 0:
-                    tail = acc_for_hist[-rep_window:].long().cpu().tolist()
-                    hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)]
-                else:
-                    hist_per_cb = [[] for _ in range(self.n_vq)]
-
-                should_continue_t, new_codes = self.local_transformer.generate_frame(
-                    last_h,
-                    self.audio_lm_heads,
-                    self.audio_embeddings,
-                    self.local_text_lm_head,
-                    n_vq=self.n_vq,
-                    do_sample=True,
-                    temperature=0.8,
-                    top_k=30,
-                    top_p=0.6,
-                    repetition_penalty=1.1,
-                    history_per_codebook=hist_per_cb,
-                )
-                should_continue = bool(should_continue_t[0].item())
-                new_codes = new_codes.squeeze(0)  # (n_vq,)
-
-                step = int(state.get("step", 0))
-                max_new_frames = int(state.get("max_new_frames", -1))
-                if not should_continue or (0 <= max_new_frames <= step):
-                    state["is_stopping"] = True
-                    state["step"] = step + 1
-                    info["audio_codes"] = {
-                        "current": new_codes,
-                        "accumulated": (info.get("audio_codes", {}) or {}).get("accumulated"),
-                    }
-                    continue  # don't append the stop frame to accumulated
-
-                acc = (info.get("audio_codes", {}) or {}).get("accumulated")
-                if isinstance(acc, torch.Tensor) and acc.numel() > 0:
-                    updated_acc = torch.cat([acc.to(new_codes.device), new_codes.unsqueeze(0)], dim=0)
-                else:
-                    updated_acc = new_codes.unsqueeze(0)
-
-                info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc}
-                state["step"] = step + 1
-                per_req_codes[i] = updated_acc
+        for info in info_dicts:
+            current = (info.get("audio_codes", {}) or {}).get("current") if isinstance(info, dict) else None
+            emit = bool((info.get("audio_codes", {}) or {}).get("emit")) if isinstance(info, dict) else False
+            if emit and isinstance(current, torch.Tensor) and current.numel() > 0:
+                codes = current.to(device=hidden.device, dtype=torch.long)
+                if codes.dim() == 1:
+                    codes = codes.unsqueeze(0)
+                per_req_codes.append(codes)
                 have_codes = True
+            else:
+                per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
 
         if not have_codes:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
 
-        # Local-v1.5 emits raw codes (no delay pattern), so signal that to the
-        # chunk processor the same way Realtime does, via ``meta.finished``
-        # (read as "skip de-delay" by ``talker2codec_async_chunk``).
-        device = hidden.device
+        # Local-v1.5 emits the current raw code frame (no delay pattern); the
+        # pipeline routes those rows through ``talker2codec_raw_async_chunk``.
+        per_req_ref_codes: list[torch.Tensor] = []
+        for info in info_dicts:
+            ref_codes = (info.get("codes", {}) or {}).get("ref") if isinstance(info, dict) else None
+            if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0:
+                ref_codes = ref_codes.detach()
+                if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0:
+                    ref_codes = ref_codes.view(-1, self.n_vq)
+                if ref_codes.dim() == 2 and ref_codes.shape[1] == self.n_vq:
+                    per_req_ref_codes.append(ref_codes.to(device=hidden.device, dtype=torch.long))
+                    continue
+            per_req_ref_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long))
+
         return OmniOutput(
             text_hidden_states=hidden,
             multimodal_outputs={
-                "codes": {"audio": per_req_codes},
-                "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)},
+                "codes": {"audio": per_req_codes, "ref": per_req_ref_codes},
             },
         )
 
diff --git a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py
index ed03c018e..87f255072 100644
--- a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py
+++ b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py
@@ -1,6 +1,6 @@
 """CUDA Graph acceleration for the MOSS Audio Tokenizer codec decoder.
 
-Captures MossAudioTokenizerModel._decode for a set of fixed frame-count
+Captures MossAudioTokenizerModel._decode_frame for a set of fixed frame-count
 bucket sizes, then replays the captured graph at inference time to eliminate
 kernel-launch overhead.  Inputs that exceed all captured sizes fall back to
 eager execution transparently.
@@ -24,7 +24,7 @@ logger = init_logger(__name__)
 
 
 class MossTTSCUDAGraphCodecWrapper:
-    """CUDA Graph wrapper for MossAudioTokenizerModel._decode.
+    """CUDA Graph wrapper for MossAudioTokenizerModel._decode_frame.
 
     Graphs are keyed by padded_T (int).  On each call the actual T is
     bucket-matched to the smallest pre-captured size >= T.  The static code
@@ -99,7 +99,7 @@ class MossTTSCUDAGraphCodecWrapper:
             dummy_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device)
             dummy_lengths = torch.tensor([size], dtype=torch.long, device=device)
             with torch.no_grad():
-                _ = self.model._decode(dummy_codes, dummy_lengths)
+                _ = self.model._decode_frame(dummy_codes, dummy_lengths)
 
         torch.accelerator.synchronize(device)
 
@@ -132,13 +132,13 @@ class MossTTSCUDAGraphCodecWrapper:
 
         # Extra eager warmup inside capture to ensure all kernels are compiled.
         with torch.no_grad():
-            _ = self.model._decode(static_codes, static_lengths)
+            _ = self.model._decode_frame(static_codes, static_lengths)
         torch.accelerator.synchronize(device)
 
         graph = CUDAGraph()
         with torch.no_grad():
             with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
-                static_out = self.model._decode(static_codes, static_lengths)
+                static_out = self.model._decode_frame(static_codes, static_lengths)
 
         self.graphs[size] = graph
         self.static_codes[size] = static_codes
diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py
index d0cec08e5..c6cf5c86c 100644
--- a/vllm_omni/model_executor/models/moss_tts/pipeline.py
+++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py
@@ -31,7 +31,7 @@ MOSS_TTS_PIPELINE = PipelineConfig(
             input_sources=(),
             owns_tokenizer=True,
             engine_output_type="latent",
-            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_delay_async_chunk"),
             sampling_constraints={
                 "detokenize": False,
             },
@@ -62,7 +62,7 @@ MOSS_TTS_REALTIME_PIPELINE = PipelineConfig(
             input_sources=(),
             owns_tokenizer=True,
             engine_output_type="latent",
-            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"),
             sampling_constraints={"detokenize": False},
         ),
         StagePipelineConfig(
@@ -91,8 +91,11 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig(
             input_sources=(),
             owns_tokenizer=True,
             engine_output_type="latent",
-            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"),
-            sampling_constraints={"detokenize": False},
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"),
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [151645],
+            },
         ),
         StagePipelineConfig(
             stage_id=1,
diff --git a/vllm_omni/model_executor/models/omnivoice/pipeline.py b/vllm_omni/model_executor/models/omnivoice/pipeline.py
new file mode 100644
index 000000000..88d479459
--- /dev/null
+++ b/vllm_omni/model_executor/models/omnivoice/pipeline.py
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""OmniVoice pipeline topology (frozen).
+
+Single-stage diffusion model for text-to-speech.
+Stage 0: DiT — text → iterative unmasking → codebook tokens → DAC decode → 24kHz audio
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+OMNIVOICE_PIPELINE = PipelineConfig(
+    model_type="omnivoice",
+    model_arch="OmniVoicePipeline",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="dit",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="audio",
+            owns_tokenizer=False,
+            requires_multimodal_data=False,
+            engine_output_type="audio",
+        ),
+    ),
+)
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
index 4bd48c043..7b2c2558c 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
@@ -17,7 +17,6 @@ from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
-    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix
 from vllm.model_executor.models.vision import (
@@ -39,6 +38,7 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerDummyInputsBuilder,
     Qwen2_5OmniThinkerMultiModalProcessor,
+    Qwen2_5OmniThinkerProcessingInfo,
 )
 from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges
 from vllm_omni.platforms import current_omni_platform
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
index 5d58a63d4..d0df1c2ea 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
@@ -4,6 +4,7 @@ from collections.abc import Iterable, Iterator, Mapping
 from typing import Any
 
 import numpy as np
+import PIL.Image as PILImage
 import torch
 from torch import nn
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
@@ -26,16 +27,21 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerProcessingInfo,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase,
 )
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniThinkerMultiModalDataParser as _Qwen2_5OmniThinkerMultiModalDataParserBase,
+)
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniThinkerMultiModalProcessor as _Qwen2_5OmniThinkerMultiModalProcessorBase,
 )
+from vllm.model_executor.models.qwen2_5_omni_thinker import (
+    Qwen2_5OmniThinkerProcessingInfo as _Qwen2_5OmniThinkerProcessingInfoBase,
+)
 from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VisionTransformer,
     Qwen2_5_VLImageEmbeddingInputs,
@@ -56,12 +62,13 @@ from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalKwargsItems,
 )
-from vllm.multimodal.parse import MultiModalDataItems
+from vllm.multimodal.parse import MultiModalDataItems, VideoProcessorItems
 from vllm.multimodal.processing.processor import (
     MultiModalPromptUpdates,
     PlaceholderFeaturesInfo,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.utils.collection_utils import is_list_of
 
 from vllm_omni.quantization.component_config import (
     resolve_encoder_quant_config,
@@ -74,11 +81,143 @@ except (ImportError, ModuleNotFoundError):
 logger = init_logger(__name__)
 
 
+def _presampled_videos_hf_kwargs(
+    mm_data: Mapping[str, object],
+    mm_kwargs: Mapping[str, object],
+) -> Mapping[str, object]:
+    """Adjust HF video kwargs for videos pre-sampled by vLLM's video loader.
+
+    When ``video_metadata`` is present (emitted by
+    ``Qwen2_5OmniVideoProcessorItems``), the frames were already sampled
+    according to ``media_io_kwargs``, so the HF processor must not re-sample
+    them, and it needs the sampled fps (instead of its default) to compute
+    ``video_second_per_grid`` correctly. Otherwise the audio/video temporal
+    alignment breaks under ``use_audio_in_video=True``.
+    """
+    video_metadata = mm_data.get("video_metadata")
+    if not video_metadata:
+        return mm_kwargs
+
+    mm_kwargs = dict(mm_kwargs)
+    videos_kwargs = dict(mm_kwargs.get("videos_kwargs") or {})
+    videos_kwargs["do_sample_frames"] = False
+
+    def _compute_sampled_video_fps(metadata) -> float | None:
+        duration = getattr(metadata, "duration", None)
+        indices = getattr(metadata, "frames_indices", None)
+        if not duration or not indices or duration <= 0:
+            return None
+        return len(indices) / float(duration)
+
+    # An explicit user-provided fps takes precedence.
+    if "fps" not in videos_kwargs:
+        fps_values = [_compute_sampled_video_fps(m) for m in video_metadata]
+        # HF accepts a single fps per call and uses it for every video's
+        # video_second_per_grid.
+        known_fps = [fps for fps in fps_values if fps is not None]
+        unique_fps = set(known_fps)
+        if len(unique_fps) == 1:
+            videos_kwargs["fps"] = known_fps[0]
+        elif len(unique_fps) > 1:
+            logger.warning(
+                f"Mixed sampled FPS {sorted(unique_fps)} in one request; HF accepts a single fps, using {known_fps[0]}."
+            )
+            videos_kwargs["fps"] = known_fps[0]
+
+    mm_kwargs["videos_kwargs"] = videos_kwargs
+    return mm_kwargs
+
+
+class Qwen2_5OmniVideoProcessorItems(VideoProcessorItems):
+    """Video items that carry the loader's ``(frames, metadata)`` tuples.
+
+    The tuples are kept in ``data`` so that mm hashing and cache-miss
+    re-parsing retain the metadata; ``get_processor_data`` unpacks them into
+    the ``videos`` + ``video_metadata`` arguments of the HF processor.
+    """
+
+    def get_processor_data(self) -> Mapping[str, object]:
+        from transformers.video_utils import VideoMetadata
+
+        videos = [frames for frames, _ in self.data]
+        video_metadata = [
+            VideoMetadata(**{k: v for k, v in metadata.items() if k != "do_sample_frames"}) for _, metadata in self.data
+        ]
+        return {"videos": videos, "video_metadata": video_metadata}
+
+
+class Qwen2_5OmniThinkerMultiModalDataParser(_Qwen2_5OmniThinkerMultiModalDataParserBase):
+    def _parse_video_data(self, data):
+        if data is None or isinstance(data, dict) or self.is_embeddings(data):
+            return super()._parse_video_data(data)
+
+        # Normalize to a list of per-video items (same as the base parser).
+        if (is_list_of(data, PILImage.Image) and len(data) > 0) or (
+            isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4
+        ):
+            data_items = [data]
+        elif isinstance(data, (np.ndarray, torch.Tensor)):
+            data_items = [elem for elem in data]
+        elif isinstance(data, tuple) and len(data) == 2:
+            data_items = [data]
+        else:
+            data_items = data
+
+        videos_with_metadata = [self._get_video_with_metadata(item) for item in data_items]
+
+        # Metadata is optional: videos fetched by vLLM's media connector
+        # arrive as (frames, metadata) tuples, while plain arrays (e.g.
+        # offline `multi_modal_data` or dummy profiling inputs) carry no
+        # metadata and parse as before.
+        if any(metadata is None for _, metadata in videos_with_metadata):
+            return super()._parse_video_data(data)
+
+        return Qwen2_5OmniVideoProcessorItems(
+            videos_with_metadata,
+            metadata=[metadata for _, metadata in videos_with_metadata],
+        )
+
+
+class Qwen2_5OmniThinkerProcessingInfo(_Qwen2_5OmniThinkerProcessingInfoBase):
+    def get_data_parser(self):
+        feature_extractor = self.get_feature_extractor()
+
+        return Qwen2_5OmniThinkerMultiModalDataParser(
+            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
+            target_sr=feature_extractor.sampling_rate,
+            target_channels=self.get_target_channels(),
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
+
 class Qwen2_5OmniThinkerMultiModalProcessor(
     _Qwen2_5OmniThinkerMultiModalProcessorBase,
 ):
     """Override to fix use_audio_in_video detection when mm cache returns None."""
 
+    def _cached_apply_hf_processor(self, inputs, timing_ctx):
+        # If use_audio_in_video, process video and audio together; otherwise,
+        # skip cache to avoid errors.
+        # Prevents partial cache hits from causing processing failures.
+        if inputs.hf_processor_mm_kwargs.get("use_audio_in_video"):
+            return self._apply_hf_processor(inputs, timing_ctx)
+        return super()._cached_apply_hf_processor(inputs, timing_ctx)
+
+    def _call_hf_processor(
+        self,
+        prompt: str,
+        mm_data: Mapping[str, object],
+        mm_kwargs: Mapping[str, object],
+        tok_kwargs: Mapping[str, object],
+    ):
+        mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs)
+        return super()._call_hf_processor(
+            prompt=prompt,
+            mm_data=mm_data,
+            mm_kwargs=mm_kwargs,
+            tok_kwargs=tok_kwargs,
+        )
+
     def _maybe_apply_prompt_updates(
         self,
         mm_items: MultiModalDataItems,
diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py
index 67806f951..91eafd2a8 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py
@@ -335,19 +335,10 @@ class Qwen2Model(nn.Module):
             ("gate_up_proj", "up_proj", 1),
         ]
         params_dict = dict(self.named_parameters(remove_duplicate=False))
-        get_cache_scale = getattr(self.quant_config, "get_cache_scale", None)
         loaded_params: set[str] = set()
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
-            if get_cache_scale is not None and (scale_name := get_cache_scale(name)):
-                # Loading kv cache quantization scales
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                loaded_params.add(scale_name)
-                continue
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py
index 07debc724..ca3a09f89 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py
@@ -9,6 +9,7 @@ Stage 2: Code2Wav — RVQ codes → audio waveform
 
 from transformers import Qwen3OmniMoeConfig
 
+from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
@@ -21,6 +22,12 @@ _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_omni"
 QWEN3_OMNI_PIPELINE = PipelineConfig(
     model_type="qwen3_omni_moe",
     model_arch="Qwen3OmniMoeForConditionalGeneration",
+    endpoint_restrictions=(
+        EndpointRestriction(
+            OmniServingCapability.COMPLETIONS,
+            "Qwen3-Omni requires chat template structure for thinker-talker handoff. Use /v1/chat/completions instead.",
+        ),
+    ),
     stages=(
         StagePipelineConfig(
             stage_id=0,
@@ -44,7 +51,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             input_sources=(0,),
             hf_config_name="talker_config",
             engine_output_type="latent",
-            custom_process_input_func=f"{_PROC}.thinker2talker",
             sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
             custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
@@ -62,7 +68,6 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             final_output_type="audio",
             hf_config_name="thinker_config",
             engine_output_type="audio",
-            custom_process_input_func=f"{_PROC}.talker2code2wav",
             sampling_constraints={"detokenize": True},
         ),
     ),
diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
index 0d7f69062..011496f86 100644
--- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
+++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py
@@ -74,7 +74,6 @@ from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.qwen2_5_omni_thinker import (
     Qwen2_5OmniAudioFeatureInputs,
     Qwen2_5OmniThinkerDummyInputsBuilder,
-    Qwen2_5OmniThinkerMultiModalProcessor,
     check_interleaved_audio_video,
     merge_interleaved_embeddings,
 )
@@ -120,6 +119,9 @@ from vllm.transformers_utils.processor import cached_processor_from_config
 
 from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
     Qwen2_5OmniConditionalGenerationMixin,
+    Qwen2_5OmniThinkerMultiModalDataParser,
+    Qwen2_5OmniThinkerMultiModalProcessor,
+    _presampled_videos_hf_kwargs,
 )
 from vllm_omni.quantization.component_config import (
     PRE_QUANTIZED_METHODS,
@@ -621,6 +623,19 @@ class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProc
         assert isinstance(feature_extractor, WhisperFeatureExtractor)
         return feature_extractor
 
+    def get_data_parser(self):
+        feature_extractor = self.get_feature_extractor()
+
+        # Install the Omni parser, which keeps video metadata from vLLM's
+        # video loader (see Qwen2_5OmniVideoProcessorItems); the parser
+        # inherited from Qwen2AudioProcessingInfo would drop it.
+        return Qwen2_5OmniThinkerMultiModalDataParser(
+            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
+            target_sr=feature_extractor.sampling_rate,
+            target_channels=1,
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
     def get_supported_mm_limits(self) -> Mapping[str, int | None]:
         return {"audio": None, "image": None, "video": None}
 
@@ -671,8 +686,11 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
         tok_kwargs: Mapping[str, object],
     ) -> BatchFeature:
         mm_data = dict(mm_data)
+        mm_kwargs = dict(mm_kwargs)
         audios = mm_data.pop("audios", [])
 
+        mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs))
+
         def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray:
             length = x.shape[-1]
             if length % hop_length != 0:
@@ -927,9 +945,13 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
 
             audio_in_video_item_idx += 1
 
-            second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
-            if second_per_grid_ts:
-                video_second_per_grid_t = second_per_grid_ts[item_idx]
+            # Prefer the HF-computed value (temporal_patch_size / sampled fps)
+            # over request kwargs;
+            second_per_grid_ts = out_mm_data.get("second_per_grid_ts")
+            if second_per_grid_ts is None:
+                second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)
+            if second_per_grid_ts is not None:
+                video_second_per_grid_t = float(second_per_grid_ts[item_idx])
             else:
                 video_second_per_grid_t = 2.0
 
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py
index 8d2f0686a..9d00cc2ba 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py
@@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
     CodePredictorWrapper,
     CodePredictorWrapperConfig,
 )
+from vllm_omni.platforms import current_omni_platform
 
 from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
 
@@ -21,7 +22,7 @@ Qwen3TTSTalkerCodePredictorModelVLLM = CodePredictorBaseModel
 
 
 class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapper):
-    """Qwen3-TTS code predictor (CUDA graphs, per-call sampling, projection)."""
+    """Qwen3-TTS code predictor (per-call sampling, projection)."""
 
     def __init__(
         self,
@@ -35,7 +36,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
             vllm_config=vllm_config,
             cp_config=config,
             wrapper_config=CodePredictorWrapperConfig(
-                use_cuda_graphs=True,
+                use_cuda_graphs=current_omni_platform.is_npu(),
                 use_parallel_embedding=False,
                 use_projection=(config.hidden_size != talker_config.hidden_size),
                 return_proj_buf=False,
diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
index 94c8b2b7b..5420cb422 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable, Mapping, Sequence
 from typing import Any
 
 import numpy as np
@@ -314,6 +314,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # tensor read; postprocess receives the tail-only slice instead, which
         # avoids ~18 ms merge + ~6 ms write per step (Sy0307 profile, #3665).
         self.requires_full_prefix_cached_hidden_states = False
+        # Stage 1 (code2wav) consumes runtime audio codes, not hidden states,
+        # so the inter-stage pooler payload never needs a CPU hidden-states
+        # view. Opting out lets the runner skip the per-step blocking
+        # ``hidden_states[:n].to("cpu")`` (measured at 24% of stage-0 on-CPU
+        # samples at c=64 on H20).
+        self.omni_pooler_payload_include_hidden = False
         # ``codes.audio`` is only needed for future prefix-hit reconstruction
         # after a request has produced codec rows. Keep per-step rows on GPU and
         # materialize the CPU OmniTensorPrefixCache entry once at completion.
@@ -323,7 +329,33 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # OmniGPUModelRunner will store talker_mtp output under this key in
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
-
+        self.talker_mtp_graph_safe = True
+        # talker_mtp samples with per-row generators, so explicitly-seeded
+        # requests stay batched instead of one scalar forward per row (#4883).
+        # This is only valid while talker_mtp receives the *unpadded* active
+        # batch, i.e. when it is NOT graph-wrapped. The runner wraps talker_mtp
+        # (padded batches + a single captured RNG stream) whenever full
+        # cudagraphs are enabled, so disable per-row generators in that case.
+        # Consequence: under full cudagraphs, per-request ``tts_local_seed`` is
+        # not reproducible (matching Qwen3-Omni, which has no per-row seeding);
+        # eager mode keeps the per-row generator fast-path.
+        # TODO(#4923): the model should not read ``cudagraph_mode`` — it is an
+        # upstream (runner) concern, and the runner already re-derives the same
+        # wrap decision in ``_init_talker_mtp``. Once all TTS models declare
+        # ``talker_mtp_graph_safe`` we should move this gating into the model
+        # runner (models only expose the static capability flag, and users can
+        # then configure ``cudagraph_mode`` freely). This also means mtp seeding
+        # is silently dropped when the decode batch > 1 under full cudagraphs;
+        # the follow-up should refactor the per-row generator in the runner to
+        # be cudagraph-safe so seeded batches stay reproducible.
+        cudagraph_mode = getattr(vllm_config.compilation_config, "cudagraph_mode", None)
+        talker_mtp_graph_wrapped = (
+            self.talker_mtp_graph_safe and cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs()
+        )
+        self.talker_mtp_accepts_per_row_generators = not talker_mtp_graph_wrapped
+        self.use_async_omni_output = True
+        self.eager_omni_postprocess_before_async_output = True
+        self.omni_pooler_payload_include_hidden = False
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
 
         if get_pp_group().is_last_rank:
@@ -1060,6 +1092,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         top_k: int | None = None,
         top_p: float | None = None,
         generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
         **kwargs: Any,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1).
@@ -1099,6 +1132,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             top_k=top_k,
             top_p=top_p,
             generator=generator,
+            generators=generators,
         )  # [B, Q]
 
         # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes.
diff --git a/vllm_omni/model_executor/models/voxcpm2/pipeline.py b/vllm_omni/model_executor/models/voxcpm2/pipeline.py
index 347fce177..0579ebc8c 100644
--- a/vllm_omni/model_executor/models/voxcpm2/pipeline.py
+++ b/vllm_omni/model_executor/models/voxcpm2/pipeline.py
@@ -26,6 +26,7 @@ VOXCPM2_PIPELINE = PipelineConfig(
             final_output_type="audio",
             owns_tokenizer=True,
             engine_output_type="audio",
+            scheduler_cls="vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler",
             sampling_constraints={
                 "detokenize": False,
                 "stop_token_ids": [1],
diff --git a/vllm_omni/model_executor/models/voxcpm2/runtime_config.py b/vllm_omni/model_executor/models/voxcpm2/runtime_config.py
new file mode 100644
index 000000000..1446729e7
--- /dev/null
+++ b/vllm_omni/model_executor/models/voxcpm2/runtime_config.py
@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclasses.dataclass(frozen=True)
+class _VoxCPM2RuntimeConfig:
+    enable_profiling: bool = False
+    enable_nvtx_profile: bool = False
+    enable_loc_dit_layer_nvtx: bool = False
+    enable_loc_dit_fused_qkv: bool = True
+    enable_loc_dit_fused_mlp: bool = True
+    enable_loc_dit_zero_dt_cache: bool = True
+    enable_loc_dit_fast_rope: bool = False
+    enable_loc_dit_skip_qkv_contig: bool = True
+    enable_loc_dit_reduce_overhead_no_cg: bool = False
+    enable_loc_dit_fullgraph_no_cg: bool = False
+    cfg_cutoff_ratio: float = 1.0
+    decode_graph_capture_policy: str = "all"
+    enable_vae_cuda_graph: bool = False
+    enable_cfm_cuda_graph: bool = False
+    enable_cfm_prealloc_output: bool = False
+    enable_batched_cfm: bool = True
+    deterministic_cfm_noise: bool = False
+    deterministic_cfm_seed: int = 20260601
+    audio_emit_every: int = 1
+    vae_decode_every: int = 3
+    enable_delayed_audio_copy: bool = False
+    delayed_audio_copy_use_events: bool = False
+    coalesce_audio_d2h: bool = True
+    enable_batched_vae_decode: bool = True
+    enable_batched_fsq_fusion: bool = False
+    batched_fsq_fusion_max_batch: int = 32
+    enable_batched_prefill_tail: bool = False
+    enable_unified_decode_graph: bool = True
+    unified_decode_graph_max_batch_size: int = 64
+
+    @classmethod
+    def from_vllm_config(cls, vllm_config: Any) -> _VoxCPM2RuntimeConfig:
+        model_config = vllm_config.model_config
+        raw = getattr(model_config, "voxcpm2_runtime_config", None)
+        if raw is None:
+            raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None)
+        if raw is None:
+            return cls()
+        if hasattr(raw, "to_dict"):
+            raw = raw.to_dict()
+        elif not isinstance(raw, dict) and hasattr(raw, "__dict__"):
+            raw = vars(raw)
+        if not isinstance(raw, dict):
+            logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw)
+            return cls()
+
+        fields = {field.name: field for field in dataclasses.fields(cls)}
+        values: dict[str, Any] = {}
+        for key, value in raw.items():
+            if key not in fields:
+                logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key)
+                continue
+            default = getattr(cls(), key)
+            values[key] = cls._coerce_value(key, value, default)
+        return cls(**values)._normalized()
+
+    def _normalized(self) -> _VoxCPM2RuntimeConfig:
+        cfg = self
+        if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph:
+            logger.warning(
+                "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; "
+                "disabling CFM CUDA Graph and keeping batched CFM."
+            )
+            cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False)
+        if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph:
+            logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.")
+            cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False)
+        if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy:
+            logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.")
+            cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False)
+        if cfg.enable_batched_vae_decode and (
+            cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph
+        ):
+            logger.warning(
+                "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, "
+                "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode."
+            )
+            cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False)
+        if cfg.enable_unified_decode_graph:
+            if cfg.enable_cfm_cuda_graph:
+                logger.info("VoxCPM2 unified decode graph supersedes independent CFM CUDA Graph; disabling it.")
+                cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False, enable_cfm_prealloc_output=False)
+            if not cfg.enable_batched_cfm:
+                logger.info("VoxCPM2 unified decode graph requires batched CFM; enabling it.")
+                cfg = dataclasses.replace(cfg, enable_batched_cfm=True)
+        return cfg
+
+    def unified_decode_graph_available(self, *, use_cuda_graph: bool) -> bool:
+        return bool(use_cuda_graph and self.enable_unified_decode_graph and not self.deterministic_cfm_noise)
+
+    @staticmethod
+    def _coerce_value(key: str, value: Any, default: Any) -> Any:
+        if isinstance(default, bool):
+            if isinstance(value, str):
+                return value.strip().lower() in ("1", "true", "yes", "on")
+            return bool(value)
+        if isinstance(default, int) and not isinstance(default, bool):
+            value = int(value)
+            if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}:
+                return max(1, value)
+            return value
+        if isinstance(default, float):
+            if key == "cfg_cutoff_ratio":
+                return min(1.0, max(0.0, float(value)))
+            return float(value)
+        if isinstance(default, str):
+            return str(value)
+        return value
diff --git a/vllm_omni/model_executor/models/voxcpm2/scheduler.py b/vllm_omni/model_executor/models/voxcpm2/scheduler.py
new file mode 100644
index 000000000..a780c6955
--- /dev/null
+++ b/vllm_omni/model_executor/models/voxcpm2/scheduler.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler
+from vllm_omni.platforms import current_omni_platform
+
+from .runtime_config import _VoxCPM2RuntimeConfig
+
+
+class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler):
+    """VoxCPM2 scheduler variant for full unified decode graph serving.
+
+    VoxCPM2's full unified decode graph only applies to pure decode batches.
+    When a decode-ready request is already running, this scheduler defers new
+    waiting admissions for the current tick so the decode batch can stay on the
+    VoxCPM2 graph path. This is a model-local serving policy, not a generic AR
+    scheduler rule.
+    """
+
+    def _unified_decode_graph_enabled(self) -> bool:
+        runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)
+        return runtime_config.unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())
+
+    def _should_defer_waiting_for_unified_decode_graph(self) -> bool:
+        if not self._unified_decode_graph_enabled():
+            return False
+        if not self.waiting or not self.running:
+            return False
+
+        for request in self.running:
+            if getattr(request, "status", None) != RequestStatus.RUNNING or request.is_finished():
+                continue
+            if self._get_confirmed_num_computed_tokens(request) >= request.num_prompt_tokens:
+                return True
+        return False
+
+    def _should_defer_waiting_admission(self) -> bool:
+        return self._should_defer_waiting_for_unified_decode_graph()
diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
index 9505bfbee..94ac454d8 100644
--- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
+++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py
@@ -16,12 +16,13 @@ import hashlib
 import logging
 import math
 import time
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable, Sequence
 from types import MethodType
-from typing import Any
+from typing import Any, NamedTuple, Protocol, TypedDict
 
 import torch
 import torch.nn as nn
+from typing_extensions import Unpack
 from vllm.config import VllmConfig
 from vllm.forward_context import get_forward_context, override_forward_context
 from vllm.inputs import tokens_input
@@ -42,8 +43,10 @@ from vllm_omni.utils.speaker_cache import (
     iter_custom_voice_profiles,
     load_validated_profile_tensors,
 )
+from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
 
 from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM
+from .runtime_config import _VoxCPM2RuntimeConfig
 from .voxcpm2_import_utils import import_voxcpm2_core
 
 logger = init_logger(__name__)
@@ -56,103 +59,50 @@ _ENABLE_NVTX_PROFILE = False
 _ACTIVE_STATE_LEAK_WARN_MIN = 512
 
 
-@dataclasses.dataclass(frozen=True)
-class _VoxCPM2RuntimeConfig:
-    enable_profiling: bool = False
-    enable_nvtx_profile: bool = False
-    enable_loc_dit_layer_nvtx: bool = False
-    enable_loc_dit_fused_qkv: bool = True
-    enable_loc_dit_fused_mlp: bool = True
-    enable_loc_dit_zero_dt_cache: bool = True
-    enable_loc_dit_fast_rope: bool = False
-    enable_loc_dit_skip_qkv_contig: bool = True
-    enable_loc_dit_reduce_overhead_no_cg: bool = False
-    enable_loc_dit_fullgraph_no_cg: bool = False
-    cfg_cutoff_ratio: float = 1.0
-    decode_graph_capture_policy: str = "all"
-    enable_vae_cuda_graph: bool = False
-    enable_cfm_cuda_graph: bool = False
-    enable_cfm_prealloc_output: bool = False
-    enable_batched_cfm: bool = True
-    deterministic_cfm_noise: bool = False
-    deterministic_cfm_seed: int = 20260601
-    audio_emit_every: int = 1
-    vae_decode_every: int = 3
-    enable_delayed_audio_copy: bool = False
-    delayed_audio_copy_use_events: bool = False
-    coalesce_audio_d2h: bool = True
-    enable_batched_vae_decode: bool = True
-    enable_batched_fsq_fusion: bool = False
-    batched_fsq_fusion_max_batch: int = 32
+class _VoxCPM2PromptConfigLike(Protocol):
+    audio_vae_config: dict[str, Any]
+    patch_size: int
 
-    @classmethod
-    def from_vllm_config(cls, vllm_config: VllmConfig) -> _VoxCPM2RuntimeConfig:
-        model_config = vllm_config.model_config
-        raw = getattr(model_config, "voxcpm2_runtime_config", None)
-        if raw is None:
-            raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None)
-        if raw is None:
-            return cls()
-        if hasattr(raw, "to_dict"):
-            raw = raw.to_dict()
-        elif not isinstance(raw, dict) and hasattr(raw, "__dict__"):
-            raw = vars(raw)
-        if not isinstance(raw, dict):
-            logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw)
-            return cls()
 
-        fields = {field.name: field for field in dataclasses.fields(cls)}
-        values: dict[str, Any] = {}
-        for key, value in raw.items():
-            if key not in fields:
-                logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key)
-                continue
-            default = getattr(cls(), key)
-            values[key] = cls._coerce_value(key, value, default)
-        return cls(**values)._normalized()
+class _AttentionMetadataLike(Protocol):
+    scheduler_metadata: Any
 
-    def _normalized(self) -> _VoxCPM2RuntimeConfig:
-        cfg = self
-        if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph:
-            logger.warning(
-                "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; "
-                "disabling CFM CUDA Graph and keeping batched CFM."
-            )
-            cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False)
-        if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph:
-            logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.")
-            cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False)
-        if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy:
-            logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.")
-            cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False)
-        if cfg.enable_batched_vae_decode and (
-            cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph
-        ):
-            logger.warning(
-                "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, "
-                "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode."
-            )
-            cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False)
-        return cfg
 
-    @staticmethod
-    def _coerce_value(key: str, value: Any, default: Any) -> Any:
-        if isinstance(default, bool):
-            if isinstance(value, str):
-                return value.strip().lower() in ("1", "true", "yes", "on")
-            return bool(value)
-        if isinstance(default, int) and not isinstance(default, bool):
-            value = int(value)
-            if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}:
-                return max(1, value)
-            return value
-        if isinstance(default, float):
-            if key == "cfg_cutoff_ratio":
-                return min(1.0, max(0.0, float(value)))
-            return float(value)
-        if isinstance(default, str):
-            return str(value)
-        return value
+class _ForwardContextLike(Protocol):
+    attn_metadata: dict[str, _AttentionMetadataLike] | Any
+
+
+class VoxCPM2PreprocessInput(TypedDict, total=False):
+    additional_information: dict[str, Any]
+    request_id: str
+    text_token_ids: list[list[int]]
+    reference_audio: object
+    ref_audio: object
+    prompt_audio: object
+    prompt_text: str | list[str] | None
+    voice_profile: dict[str, Any] | list[dict[str, Any]] | None
+    voice_name: str | list[str] | None
+    voice_created_at: int | str | None
+
+
+class VoxCPM2PostprocessInput(TypedDict, total=False):
+    request_id: str
+
+
+class _PrefillInputs(NamedTuple):
+    text_token: torch.Tensor
+    audio_feat: torch.Tensor
+    text_mask: torch.Tensor
+    audio_mask: torch.Tensor
+
+
+class _PrefillResidualMeta(NamedTuple):
+    lm_hidden: torch.Tensor
+    prefix_feat_cond: torch.Tensor
+
+
+class _DecodeResidualMeta(NamedTuple):
+    new_lm_hidden: torch.Tensor
 
 
 def is_cjk_char(c: str) -> bool:
@@ -195,11 +145,11 @@ def split_multichar_chinese(token_ids: list[int], split_map: dict[int, list[int]
 
 
 def build_voxcpm2_prompt(
-    hf_config: Any,
+    hf_config: _VoxCPM2PromptConfigLike,
     tokenizer: Any,
     split_map: dict[int, list[int]],
     text: str,
-    ref_audio: Any | None = None,
+    ref_audio: Sequence[float] | torch.Tensor | None = None,
     ref_sr: int | None = None,
     ref_text: str | None = None,
     voice_profile: dict[str, Any] | None = None,
@@ -234,6 +184,8 @@ def build_voxcpm2_prompt(
                 prefill_len += len(prompt_ids)
             prefill_len += audio_feat_len
     elif ref_audio is not None:
+        if ref_sr is None:
+            raise ValueError("VoxCPM2 ref_sr is required when ref_audio is provided.")
         vae = hf_config.audio_vae_config
         patch_samples = hf_config.patch_size * math.prod(vae["encoder_rates"])
         ref_len = math.ceil(math.ceil(len(ref_audio) * vae["sample_rate"] / ref_sr) / patch_samples)
@@ -343,6 +295,46 @@ class _CapturedCFMGraph:
     buffers: _CFMBufferManager
 
 
+@dataclasses.dataclass
+class _CapturedUnifiedDecodeGraph:
+    graph: torch.cuda.CUDAGraph
+    batch_size: int
+    # Static input buffers (copied before replay)
+    input_embeds: torch.Tensor
+    positions: torch.Tensor
+    prev_feat_embed: torch.Tensor
+    prefix_feat_cond: torch.Tensor
+    # Static output buffers (read after replay)
+    next_feat_embed: torch.Tensor
+    cfm_output: torch.Tensor
+    lm_hidden: torch.Tensor
+    # Internal buffers for CFM solver
+    cfm_buffers: _CFMBufferManager
+    cfm_noise: torch.Tensor
+
+
+@dataclasses.dataclass
+class _UnifiedDecodeGraphStats:
+    captures: int = 0
+    replays: int = 0
+    skips: dict[str, int] = dataclasses.field(default_factory=dict)
+    real_batch_sizes: dict[int, int] = dataclasses.field(default_factory=dict)
+    graph_bucket_sizes: dict[int, int] = dataclasses.field(default_factory=dict)
+    logged_replays: int = 0
+    logged_skips: int = 0
+
+    def record_skip(self, reason: str) -> None:
+        self.skips[reason] = self.skips.get(reason, 0) + 1
+
+    def record_replay_batch(self, *, real_batch_size: int, graph_bucket_size: int) -> None:
+        self.real_batch_sizes[real_batch_size] = self.real_batch_sizes.get(real_batch_size, 0) + 1
+        self.graph_bucket_sizes[graph_bucket_size] = self.graph_bucket_sizes.get(graph_bucket_size, 0) + 1
+
+    @property
+    def total_skips(self) -> int:
+        return sum(self.skips.values())
+
+
 @dataclasses.dataclass
 class _PendingAudioCopy:
     host: torch.Tensor
@@ -927,6 +919,17 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         self._max_cached_graphs = self._max_batch_size
         self._cuda_graph_warmup_steps = 0
         self._cuda_graph_warmup_threshold = 3
+        self._unified_graphs: dict[int, _CapturedUnifiedDecodeGraph] = {}
+        self._enable_unified_decode_graph = self._runtime_config.unified_decode_graph_available(
+            use_cuda_graph=use_cuda_graph
+        )
+        self._unified_decode_graph_max_batch_size = self._runtime_config.unified_decode_graph_max_batch_size
+        self._unified_graph_bucket_sizes: frozenset[int] = self._build_unified_graph_bucket_sizes(
+            min(self._max_batch_size, self._unified_decode_graph_max_batch_size)
+        )
+        self._unified_graph_stats = _UnifiedDecodeGraphStats()
+        self._runner_assisted_unified_decode_graph_active = False
+        self._runner_assisted_unified_decode_graph_batch_size = 0
 
         self._multichar_zh_split: dict[int, list[int]] | None = None
 
@@ -1035,10 +1038,10 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
 
     def _build_prompt_cache(
         self,
-        ref_audio: Any = None,
-        prompt_audio: Any = None,
+        ref_audio: object = None,
+        prompt_audio: object = None,
         prompt_text: str | None = None,
-    ) -> dict | None:
+    ) -> dict[str, Any] | None:
         """Build prompt cache, handling both file paths and raw audio data.
 
         The OpenAI speech API sends decoded audio as [samples_list, sr]
@@ -1106,6 +1109,84 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             max_batch_size=self._max_batch_size,
         )
 
+    @staticmethod
+    def _voxcpm2_unwrap_torch_compile(module: nn.Module | Callable) -> nn.Module | Callable:
+        """Return the eager module behind torch.compile wrappers."""
+        seen: set[int] = set()
+        while hasattr(module, "_orig_mod") and id(module) not in seen:
+            seen.add(id(module))
+            module = getattr(module, "_orig_mod")
+        return module
+
+    @staticmethod
+    def _voxcpm2_compile_without_inductor_cudagraphs(
+        module: nn.Module | Callable,
+        *,
+        mode: str | None,
+        fullgraph: bool,
+    ) -> Callable:
+        """Compile VoxCPM2 LocDiT for capture inside VoxCPM2-owned CUDA graphs."""
+        kwargs: dict[str, Any] = {
+            "fullgraph": fullgraph,
+            "options": {
+                "triton.cudagraphs": False,
+                "triton.cudagraph_trees": False,
+            },
+        }
+        if mode is not None:
+            kwargs["mode"] = mode
+        return torch.compile(module, **kwargs)
+
+    @staticmethod
+    def _raise_if_cuda_runtime_error(error: BaseException) -> None:
+        if isinstance(error, torch.cuda.OutOfMemoryError):
+            raise error
+        if isinstance(error, RuntimeError):
+            message = str(error).lower()
+            cuda_markers = (
+                "cuda error",
+                "device-side assert",
+                "illegal memory access",
+                "out of memory",
+                "cublas",
+                "cudnn",
+                "cufft",
+            )
+            if any(marker in message for marker in cuda_markers):
+                raise error
+
+    def _voxcpm2_compile_unified_capture_estimator(self, module: nn.Module | Callable) -> Callable:
+        eager_module = self._voxcpm2_unwrap_torch_compile(module)
+        if getattr(self, "_estimator_for_unified_capture_source", None) is eager_module and hasattr(
+            self, "_estimator_for_unified_capture"
+        ):
+            return self._estimator_for_unified_capture
+        compiled = self._voxcpm2_compile_without_inductor_cudagraphs(
+            eager_module,
+            mode=None,
+            fullgraph=False,
+        )
+        compiled._compiled = True
+        self._estimator_for_unified_capture_source = eager_module
+        self._estimator_for_unified_capture = compiled
+        return compiled
+
+    def _voxcpm2_compile_unified_capture_feat_encoder(self, module: nn.Module | Callable) -> Callable:
+        eager_module = self._voxcpm2_unwrap_torch_compile(module)
+        if getattr(self, "_feat_encoder_for_unified_capture_source", None) is eager_module and hasattr(
+            self, "_feat_encoder_for_unified_capture"
+        ):
+            return self._feat_encoder_for_unified_capture
+        compiled = self._voxcpm2_compile_without_inductor_cudagraphs(
+            eager_module,
+            mode=None,
+            fullgraph=False,
+        )
+        compiled._compiled = True
+        self._feat_encoder_for_unified_capture_source = eager_module
+        self._feat_encoder_for_unified_capture = compiled
+        return compiled
+
     def _setup_torch_compile(self) -> None:
         if not self._enable_torch_compile:
             return
@@ -1133,30 +1214,42 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             if _install_locdit_zero_dt_cache(estimator):
                 targets.append("LocDiT zero-dt embedding cache")
 
+        external_cfm_capture = self._enable_cfm_cuda_graph or self._enable_unified_decode_graph
         if cfg.enable_loc_dit_layer_nvtx:
             patched = _install_locdit_layer_nvtx(estimator)
             if patched:
                 targets.append(f"LocDiT layer NVTX ({patched} layers, compile skipped)")
             estimator._compiled = True
-        elif self._enable_cfm_cuda_graph:
+        elif self._enable_unified_decode_graph:
+            try:
+                compiled_ro = torch.compile(estimator, mode="reduce-overhead", fullgraph=False)
+                compiled_ro._compiled = True
+                tts.feat_decoder.estimator = compiled_ro
+                targets.append("LocDiT (reduce-overhead serving + lazy no-cg unified capture)")
+            except Exception as e:
+                self._raise_if_cuda_runtime_error(e)
+                logger.warning("torch.compile LocDiT dual-mode failed: %s", e)
+                try:
+                    tts.feat_decoder.estimator = self._voxcpm2_compile_without_inductor_cudagraphs(
+                        estimator,
+                        mode="reduce-overhead",
+                        fullgraph=False,
+                    )
+                    tts.feat_decoder.estimator._compiled = True
+                    self._estimator_for_unified_capture = tts.feat_decoder.estimator
+                    targets.append("LocDiT (no-cg only, unified capture)")
+                except Exception as inner_e:
+                    self._raise_if_cuda_runtime_error(inner_e)
+                    logger.warning("torch.compile LocDiT failed completely: %s", inner_e)
+        elif external_cfm_capture:
             try:
                 if cfg.enable_loc_dit_reduce_overhead_no_cg:
-                    import torch._inductor.config as inductor_config
-
-                    old_cudagraphs = inductor_config.triton.cudagraphs
-                    old_cudagraph_trees = inductor_config.triton.cudagraph_trees
-                    inductor_config.triton.cudagraphs = False
-                    inductor_config.triton.cudagraph_trees = False
-                    try:
-                        tts.feat_decoder.estimator = torch.compile(
-                            estimator,
-                            mode="reduce-overhead",
-                            fullgraph=False,
-                        )
-                    finally:
-                        inductor_config.triton.cudagraphs = old_cudagraphs
-                        inductor_config.triton.cudagraph_trees = old_cudagraph_trees
-                    targets.append("LocDiT (reduce-overhead no-inductor-cudagraph + CUDA Graph)")
+                    tts.feat_decoder.estimator = self._voxcpm2_compile_without_inductor_cudagraphs(
+                        estimator,
+                        mode="reduce-overhead",
+                        fullgraph=False,
+                    )
+                    targets.append("LocDiT (reduce-overhead no-inductor-cudagraph + external CUDA Graph)")
                 else:
                     tts.feat_decoder.estimator = torch.compile(
                         estimator,
@@ -1166,25 +1259,11 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                             "triton.cudagraph_trees": False,
                         },
                     )
-                    graph_mode = "fullgraph" if cfg.enable_loc_dit_fullgraph_no_cg else "partitioned"
-                    targets.append(f"LocDiT ({graph_mode} compile no-inductor-cudagraph + CUDA Graph)")
+                    targets.append("LocDiT (no-inductor-cudagraph + external CUDA Graph)")
                 tts.feat_decoder.estimator._compiled = True
             except Exception as e:
-                logger.warning("torch.compile LocDiT reduce-overhead for CUDA Graph failed: %s", e)
-                try:
-                    tts.feat_decoder.estimator = torch.compile(
-                        estimator,
-                        fullgraph=cfg.enable_loc_dit_fullgraph_no_cg,
-                        options={
-                            "triton.cudagraphs": False,
-                            "triton.cudagraph_trees": False,
-                        },
-                    )
-                    tts.feat_decoder.estimator._compiled = True
-                    graph_mode = "fullgraph" if cfg.enable_loc_dit_fullgraph_no_cg else "partitioned"
-                    targets.append(f"LocDiT ({graph_mode} compile no-inductor-cudagraph + CUDA Graph)")
-                except Exception as inner_e:
-                    logger.warning("torch.compile LocDiT for CUDA Graph failed: %s", inner_e)
+                self._raise_if_cuda_runtime_error(e)
+                logger.warning("torch.compile LocDiT for external CUDA Graph failed: %s", e)
         else:
             try:
                 tts.feat_decoder.estimator = torch.compile(
@@ -1195,14 +1274,34 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                 tts.feat_decoder.estimator._compiled = True
                 targets.append("LocDiT")
             except Exception as e:
+                self._raise_if_cuda_runtime_error(e)
                 logger.warning("torch.compile LocDiT failed: %s", e)
 
         try:
             if not hasattr(tts.feat_encoder, "_compiled"):
-                tts.feat_encoder = torch.compile(tts.feat_encoder, mode="reduce-overhead", fullgraph=False)
-                tts.feat_encoder._compiled = True
-                targets.append("feat_encoder")
+                if self._enable_unified_decode_graph:
+                    feat_encoder = tts.feat_encoder
+                    compiled_ro = torch.compile(feat_encoder, mode="reduce-overhead", fullgraph=False)
+                    compiled_ro._compiled = True
+                    tts.feat_encoder = compiled_ro
+                    targets.append("feat_encoder (reduce-overhead serving + lazy no-cg unified capture)")
+                elif external_cfm_capture:
+                    tts.feat_encoder = torch.compile(
+                        tts.feat_encoder,
+                        fullgraph=False,
+                        options={
+                            "triton.cudagraphs": False,
+                            "triton.cudagraph_trees": False,
+                        },
+                    )
+                    tts.feat_encoder._compiled = True
+                    targets.append("feat_encoder (no-inductor-cudagraph)")
+                else:
+                    tts.feat_encoder = torch.compile(tts.feat_encoder, mode="reduce-overhead", fullgraph=False)
+                    tts.feat_encoder._compiled = True
+                    targets.append("feat_encoder")
         except Exception as e:
+            self._raise_if_cuda_runtime_error(e)
             logger.warning("torch.compile feat_encoder failed: %s", e)
 
         if self._compile_vae and not self._enable_vae_cuda_graph:
@@ -1212,6 +1311,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                     tts.audio_vae._compiled = True
                     targets.append("AudioVAE")
             except Exception as e:
+                self._raise_if_cuda_runtime_error(e)
                 logger.warning("torch.compile AudioVAE failed: %s", e)
 
         if not self._enable_cuda_graph:
@@ -1220,6 +1320,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                     targets.extend(f"scaffold.{t}" for t in self.model.compile_selective())
                     self.model._selective_compiled = True
                 except Exception as e:
+                    self._raise_if_cuda_runtime_error(e)
                     logger.warning("scaffold compile failed: %s", e)
 
             if not getattr(self.residual_model, "_selective_compiled", False):
@@ -1227,6 +1328,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                     targets.extend(f"residual.{t}" for t in self.residual_model.compile_selective())
                     self.residual_model._selective_compiled = True
                 except Exception as e:
+                    self._raise_if_cuda_runtime_error(e)
                     logger.warning("residual compile failed: %s", e)
         else:
             self.model.precompute_fused_qkv()
@@ -1240,6 +1342,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                 self._projections_compiled = True
                 targets.append("projections")
             except Exception as e:
+                self._raise_if_cuda_runtime_error(e)
                 self._compiled_dit_proj = self._compiled_stop_fn = None
                 logger.warning("projections compile failed: %s", e)
 
@@ -1255,7 +1358,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         return tts.stop_head(tts.stop_actn(tts.stop_proj(lm_h)))
 
     @staticmethod
-    def _nullify_volatile_metadata(ctx: Any) -> Any:
+    def _nullify_volatile_metadata(ctx: _ForwardContextLike) -> _ForwardContextLike:
         """Set ``scheduler_metadata`` to None on all attention layers.
 
         This is the only tensor FA3 reallocates each step (variable shape).
@@ -1351,6 +1454,31 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             )
         return True
 
+    @staticmethod
+    def _build_unified_graph_bucket_sizes(max_size: int) -> frozenset[int]:
+        if max_size <= 0:
+            return frozenset()
+        sizes = {1}
+        size = 2
+        while size <= min(max_size, 16):
+            sizes.add(size)
+            size <<= 1
+        bucket = 24
+        while bucket <= max_size:
+            sizes.add(bucket)
+            bucket += 8
+        sizes.add(max_size)
+        return frozenset(size for size in sizes if 0 < size <= max_size)
+
+    def _select_unified_graph_bucket_size(self, batch_size: int) -> int | None:
+        if (
+            batch_size <= 0
+            or batch_size > self._unified_decode_graph_max_batch_size
+            or batch_size > self._max_cached_graphs
+        ):
+            return None
+        return min((size for size in self._unified_graph_bucket_sizes if size >= batch_size), default=None)
+
     def _capture_vae_graph(self, feat: torch.Tensor) -> _CapturedVAEGraph:
         batch_size = feat.shape[0]
         num_frames = feat.shape[-1]
@@ -1535,8 +1663,297 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             state.cfm_output_gpu.copy_(graph.output)
         return state.cfm_output_gpu.transpose(1, 2)
 
+    # -------------------- unified decode graph --------------------
+
+    def _capture_unified_decode_graph(self, batch_size: int) -> _CapturedUnifiedDecodeGraph:
+        H = self.config.hidden_size
+        D = self._feat_dim
+        P = self._patch_size
+        dev = torch.device(self._device)
+        dtype = self._side_dtype
+
+        self._setup_cfm_buffers()
+        if self._enable_torch_compile:
+            self._setup_torch_compile()
+        self.model.precompute_fused_qkv()
+        self.residual_model.precompute_fused_qkv()
+
+        tts = self.tts
+        dit_hidden = tts.lm_to_dit_proj.out_features + tts.res_to_dit_proj.out_features
+        cfm_buffers = _CFMBufferManager(
+            device=dev,
+            dtype=dtype,
+            feat_dim=D,
+            patch_size=P,
+            dit_hidden_size=dit_hidden,
+            max_batch_size=batch_size,
+        )
+
+        g = _CapturedUnifiedDecodeGraph(
+            graph=torch.cuda.CUDAGraph(),
+            batch_size=batch_size,
+            input_embeds=torch.zeros(batch_size, H, device=dev, dtype=dtype),
+            positions=torch.zeros(batch_size, device=dev, dtype=torch.long),
+            prev_feat_embed=torch.zeros(batch_size, H, device=dev, dtype=dtype),
+            prefix_feat_cond=torch.zeros(batch_size, P, D, device=dev, dtype=dtype),
+            next_feat_embed=torch.zeros(batch_size, H, device=dev, dtype=dtype),
+            cfm_output=torch.zeros(batch_size, D, P, device=dev, dtype=dtype),
+            lm_hidden=torch.zeros(batch_size, H, device=dev, dtype=dtype),
+            cfm_buffers=cfm_buffers,
+            cfm_noise=torch.zeros(batch_size, D, P, device=dev, dtype=dtype),
+        )
+
+        dit_proj = self._dit_proj_fn
+
+        def unified_fwd():
+            scaffold_h = self.model(input_ids=None, positions=g.positions, inputs_embeds=g.input_embeds)
+            if isinstance(scaffold_h, tuple):
+                scaffold_h = scaffold_h[0]
+            lm_h = tts.fsq_layer(scaffold_h)
+            res_input = tts.fusion_concat_proj(torch.cat([lm_h, g.prev_feat_embed], dim=-1))
+            res_h = self.residual_model(positions=g.positions, inputs_embeds=res_input)
+            dit_h = dit_proj(lm_h, res_h)
+            cond = g.prefix_feat_cond.transpose(1, 2).contiguous()
+            cfm_out = _optimized_solve_euler_with_noise(
+                tts.feat_decoder,
+                dit_h,
+                P,
+                cond,
+                g.cfm_noise,
+                self._inference_timesteps,
+                self._cfg_value,
+                cfm_buffers,
+                cfg_cutoff_ratio=self._cfg_cutoff_ratio,
+            )
+            feat_enc = tts.feat_encoder(cfm_out.transpose(1, 2).unsqueeze(1)).squeeze(1)
+            next_embed = tts.enc_to_lm_proj(feat_enc)
+            return next_embed, cfm_out, lm_h
+
+        ctx = get_forward_context()
+        capture_context = override_forward_context(self._nullify_volatile_metadata(ctx))
+
+        original_estimator = tts.feat_decoder.estimator
+        capture_estimator = self._voxcpm2_compile_unified_capture_estimator(original_estimator)
+        tts.feat_decoder.estimator = capture_estimator
+
+        original_feat_encoder = tts.feat_encoder
+        capture_feat_encoder = self._voxcpm2_compile_unified_capture_feat_encoder(original_feat_encoder)
+        tts.feat_encoder = capture_feat_encoder
+
+        try:
+            with capture_context:
+                with torch.no_grad():
+                    for _ in range(3):
+                        g.cfm_noise.normal_()
+                        unified_fwd()
+
+                    g.cfm_noise.normal_()
+                    with torch.cuda.graph(g.graph, pool=current_platform.get_global_graph_pool()):
+                        g.next_feat_embed, g.cfm_output, g.lm_hidden = unified_fwd()
+        finally:
+            tts.feat_decoder.estimator = original_estimator
+            tts.feat_encoder = original_feat_encoder
+        self._unified_graph_stats.captures += 1
+        logger.info(
+            "CUDA Graph captured for unified decode (batch_size=%d, captures=%d)",
+            batch_size,
+            self._unified_graph_stats.captures,
+        )
+        return g
+
+    def _unified_decode_graph_skip_reason(
+        self,
+        *,
+        can_use_graph: bool,
+        is_all_decode: bool,
+        num_reqs: int,
+    ) -> str | None:
+        if not self._enable_unified_decode_graph:
+            return "disabled"
+        if not can_use_graph:
+            return "graph_not_ready"
+        if not is_all_decode:
+            return "not_all_decode"
+        if num_reqs > self._unified_decode_graph_max_batch_size:
+            return "batch_exceeds_unified_max"
+        if num_reqs > self._max_cached_graphs:
+            return "batch_too_large"
+        graph_size = self._select_unified_graph_bucket_size(num_reqs)
+        if graph_size is None:
+            return "capture_policy"
+        if num_reqs > 1 and not self._runner_assisted_unified_decode_graph_active:
+            return "runner_full_metadata_missing"
+        for req_id, _, _, _ in self._pending_requests:
+            state = self._active_states[req_id]
+            if not state.prefill_completed:
+                return "prefill_incomplete"
+            if state.prev_feat_embed is None or state.curr_prefix_feat_cond is None:
+                return "state_not_ready"
+        return None
+
+    def _maybe_log_unified_graph_stats(self) -> None:
+        if not self._enable_profiling:
+            return
+        stats = self._unified_graph_stats
+        total_skips = stats.total_skips
+        should_log = (stats.replays - stats.logged_replays >= 100) or (total_skips - stats.logged_skips >= 50)
+        if not should_log:
+            return
+        stats.logged_replays = stats.replays
+        stats.logged_skips = total_skips
+        logger.info(
+            "VoxCPM2 unified graph stats: captures=%d replays=%d skips=%s real_batch_sizes=%s graph_bucket_sizes=%s",
+            stats.captures,
+            stats.replays,
+            stats.skips,
+            dict(sorted(stats.real_batch_sizes.items())),
+            dict(sorted(stats.graph_bucket_sizes.items())),
+        )
+
+    def _forward_unified_decode(
+        self,
+        inputs_embeds: torch.Tensor,
+        positions: torch.Tensor,
+        num_reqs: int,
+    ) -> torch.Tensor:
+        graph_size = self._select_unified_graph_bucket_size(num_reqs)
+        if graph_size is None:
+            raise RuntimeError(f"No unified decode graph bucket for batch size {num_reqs}")
+        g = self._unified_graphs.get(graph_size)
+        if g is None:
+            g = self._capture_unified_decode_graph(graph_size)
+            self._unified_graphs[graph_size] = g
+
+        self._unified_graph_stats.replays += 1
+        if self._enable_profiling:
+            self._unified_graph_stats.record_replay_batch(real_batch_size=num_reqs, graph_bucket_size=graph_size)
+        self._maybe_log_unified_graph_stats()
+        states: list[_RequestState] = []
+        commit_mask: list[bool] = []
+        for req_id, _is_prefill, _embeds, _n in self._pending_requests:
+            state = self._active_states[req_id]
+            states.append(state)
+            already_stopping = state.is_stopping
+            commit_mask.append(not already_stopping)
+            if not already_stopping:
+                state.decode_step_count += 1
+                if state.decode_step_count >= self._max_decode_steps:
+                    state.is_stopping = True
+
+        self._perf.start("unified.copy_inputs")
+        for i, state in enumerate(states):
+            pfe = state.prev_feat_embed
+            g.prev_feat_embed[i].copy_(pfe.squeeze(0) if pfe.ndim > 1 else pfe)
+            pfc = state.curr_prefix_feat_cond
+            if pfc.ndim == 2:
+                g.prefix_feat_cond[i].copy_(pfc)
+            else:
+                g.prefix_feat_cond[i].copy_(pfc.squeeze(0))
+
+        g.input_embeds[:num_reqs].copy_(inputs_embeds[:num_reqs])
+        g.positions[:num_reqs].copy_(positions[:num_reqs])
+        if graph_size > num_reqs:
+            # Runner-assisted FULL metadata pads attention to graph_size. Keep
+            # padded rows as valid duplicates instead of position-0 zero rows;
+            # some attention backends may still touch padded slots before
+            # honoring seq_lens=0.
+            last = num_reqs - 1
+            g.input_embeds[num_reqs:graph_size].copy_(g.input_embeds[last : last + 1].expand(graph_size - num_reqs, -1))
+            g.positions[num_reqs:graph_size].copy_(g.positions[last : last + 1].expand(graph_size - num_reqs))
+            g.prev_feat_embed[num_reqs:graph_size].copy_(
+                g.prev_feat_embed[last : last + 1].expand(graph_size - num_reqs, -1)
+            )
+            g.prefix_feat_cond[num_reqs:graph_size].copy_(
+                g.prefix_feat_cond[last : last + 1].expand(graph_size - num_reqs, -1, -1)
+            )
+        g.cfm_noise.normal_()
+        self._perf.stop("unified.copy_inputs")
+
+        self._perf.start("unified.replay")
+        g.graph.replay()
+        self._perf.stop("unified.replay")
+
+        self._perf.start("unified.commit")
+        with torch.no_grad():
+            all_stop_logits = self._stop_fn(g.lm_hidden[:num_reqs])
+        for i, state in enumerate(states):
+            stop_logits_i = all_stop_logits[i : i + 1]
+            if not commit_mask[i]:
+                state.precomputed_stop_logits = stop_logits_i
+                continue
+            next_embed_i = g.next_feat_embed[i : i + 1].clone()
+            cfm_out_i = g.cfm_output[i : i + 1].transpose(1, 2)
+            self._commit_decode_state(state, stop_logits_i, next_embed_i, cfm_out_i)
+        self._perf.stop("unified.commit")
+
+        self._perf.start("unified.audio")
+        self._precompute_stop_flags_for_audio_collect(states)
+        ready_audio = self._drain_ready_audio_copies_for_states(states)
+        audio_by_req = self._collect_audio_batch(states, initial_delayed_chunks_by_req=ready_audio)
+        for state in states:
+            self._results_queue.append((state.request_id, state.precomputed_stop_logits))
+            self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id)))
+
+        self._pending_requests.clear()
+        self._flush_deferred_cleanup()
+        self._perf.stop("unified.audio")
+        return g.next_feat_embed[:num_reqs]
+
     # -------------------- vllm hooks --------------------
 
+    def get_runner_assisted_full_attention_metadata_request(
+        self,
+        *,
+        req_ids: Sequence[str],
+        num_reqs: int,
+        num_scheduled_tokens: Sequence[int],
+        num_computed_tokens: Sequence[int],
+        max_num_scheduled_tokens: int,
+    ) -> RunnerAssistedFullAttentionMetadataRequest | None:
+        if not self._enable_unified_decode_graph or not self._enable_cuda_graph:
+            return None
+        if num_reqs <= 1 or num_reqs > self._unified_decode_graph_max_batch_size:
+            return None
+        if num_reqs > self._max_cached_graphs or self._select_unified_graph_bucket_size(num_reqs) is None:
+            return None
+        if max_num_scheduled_tokens != 1 or len(num_scheduled_tokens) != num_reqs:
+            return None
+        if any(int(n) != 1 for n in num_scheduled_tokens):
+            return None
+        if len(req_ids) != num_reqs or len(num_computed_tokens) != num_reqs:
+            return None
+        if any(int(n) <= 0 for n in num_computed_tokens):
+            return None
+        tts_compiled = getattr(self.tts.feat_decoder.estimator, "_compiled", False) if self._tts is not None else False
+        if not tts_compiled or self._cuda_graph_warmup_steps < self._cuda_graph_warmup_threshold:
+            return None
+        for req_id in req_ids:
+            state = self._active_states.get(req_id)
+            if (
+                state is None
+                or not state.prefill_completed
+                or state.prev_feat_embed is None
+                or state.curr_prefix_feat_cond is None
+            ):
+                return None
+        bucket_size = self._select_unified_graph_bucket_size(num_reqs)
+        if bucket_size is None:
+            return None
+        return RunnerAssistedFullAttentionMetadataRequest(
+            num_reqs_padded=bucket_size,
+            for_cudagraph_capture=bucket_size not in self._unified_graphs,
+        )
+
+    def set_runner_assisted_full_attention_metadata_context(
+        self,
+        *,
+        enabled: bool,
+        num_reqs: int = 0,
+    ) -> None:
+        enabled = bool(enabled)
+        self._runner_assisted_unified_decode_graph_active = enabled
+        self._runner_assisted_unified_decode_graph_batch_size = int(num_reqs) if enabled else 0
+
     def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor:
         return self.model.embed_input_ids(input_ids)
 
@@ -1565,12 +1982,30 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             self._enable_cuda_graph and graph_ready and intermediate_tensors is None and inputs_embeds is not None
         )
 
-        if can_use_graph and is_all_decode and num_reqs <= self._max_cached_graphs:
-            use_scaffold_graph = self._should_use_decode_graph(num_reqs)
-        else:
-            use_scaffold_graph = False
+        unified_skip_reason = self._unified_decode_graph_skip_reason(
+            can_use_graph=can_use_graph,
+            is_all_decode=is_all_decode,
+            num_reqs=num_reqs,
+        )
+        if unified_skip_reason is None:
+            self._perf.start("unified_decode")
+            result = self._forward_unified_decode(inputs_embeds, positions, num_reqs)
+            self._perf.stop("unified_decode")
+            self._perf.stop("forward_total")
+            return result
+        if self._enable_unified_decode_graph and num_reqs > 0 and unified_skip_reason is not None:
+            self._unified_graph_stats.record_skip(unified_skip_reason)
+            self._maybe_log_unified_graph_stats()
 
-        if use_scaffold_graph:
+        use_segmented_decode_graph = (
+            can_use_graph
+            and is_all_decode
+            and num_reqs <= self._max_cached_graphs
+            and (not self._enable_unified_decode_graph or unified_skip_reason is not None)
+            and self._should_use_decode_graph(num_reqs)
+        )
+
+        if use_segmented_decode_graph:
             self._perf.start("scaffold_fwd")
             if num_reqs not in self._scaffold_graphs:
                 self._scaffold_graphs[num_reqs] = self._capture_graph(self.model, num_reqs, "scaffold")
@@ -1590,7 +2025,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         token_offset = 0
         residual_inputs: list[torch.Tensor] = []
         residual_positions: list[torch.Tensor] = []
-        req_metas: list[tuple[_RequestState, bool, dict]] = []
+        req_metas: list[tuple[_RequestState, bool, _PrefillResidualMeta | _DecodeResidualMeta]] = []
         pending_decode_fsq: list[tuple[_RequestState, torch.Tensor, torch.Tensor]] = []
 
         def flush_decode_fsq_batch() -> None:
@@ -1656,6 +2091,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             residual_batch_size = batch_in.shape[0]
             use_residual_graph = (
                 self._enable_cuda_graph
+                and use_segmented_decode_graph
                 and is_all_decode
                 and graph_ready
                 and residual_batch_size == num_reqs  # 1 token per request
@@ -1679,12 +2115,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                     batch_out = self.residual_model(batch_pos, batch_in)
             self._perf.stop("residual_fwd")
 
-            can_batch_decode_tail = (
+            can_finish_decode_batch = (
                 bool(req_metas)
                 and all(not is_prefill for _, is_prefill, _ in req_metas)
                 and all(x.shape[0] == 1 for x in residual_inputs)
             )
-            if can_batch_decode_tail:
+            if can_finish_decode_batch:
                 self._finish_decode_batch(req_metas, batch_out)
                 self._precompute_stop_flags_for_audio_collect([state for state, _, _ in req_metas])
                 ready_audio_by_req = self._drain_ready_audio_copies_for_states([state for state, _, _ in req_metas])
@@ -1698,17 +2134,24 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             else:
                 offset = 0
                 decoded_states: list[_RequestState] = []
+                prefill_batch: list[tuple[_RequestState, _PrefillResidualMeta, torch.Tensor]] = []
                 for idx, (state, is_prefill, meta) in enumerate(req_metas):
                     n = residual_inputs[idx].shape[0]
                     res_out = batch_out[offset : offset + n]
                     offset += n
 
                     if is_prefill:
-                        self._finish_prefill(state, meta, res_out, dev)
+                        prefill_batch.append((state, meta, res_out))
                     else:
-                        self._finish_decode(state, meta, res_out, dev)
+                        self._finish_decode(state, meta, res_out)
                         decoded_states.append(state)
 
+                if prefill_batch and self._runtime_config.enable_batched_prefill_tail:
+                    self._finish_prefill_batch(prefill_batch)
+                else:
+                    for state, meta, res_out in prefill_batch:
+                        self._finish_prefill(state, meta, res_out, dev)
+
                 collect_states = [state for state, _, _ in req_metas]
                 self._precompute_stop_flags_for_audio_collect(collect_states)
                 ready_audio_by_req = self._drain_ready_audio_copies_for_states(collect_states)
@@ -1731,7 +2174,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
 
     # -------------------- prefill / decode helpers --------------------
 
-    def _prepare_residual_prefill(self, state: _RequestState, base_lm_out: torch.Tensor, dev: Any):
+    def _prepare_residual_prefill(
+        self,
+        state: _RequestState,
+        base_lm_out: torch.Tensor,
+        dev: torch.device,
+    ) -> tuple[torch.Tensor, _PrefillResidualMeta]:
         tts = self.tts
         text_mask, feat_mask, feat, feat_embed = state.prefill_masks
         state.prefill_masks = None
@@ -1754,10 +2202,15 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         lm_hidden = enc_outputs[:, -1, :]
 
         residual_input = tts.fusion_concat_proj(torch.cat([enc_outputs, feat_mask.unsqueeze(-1) * feat_embed], dim=-1))
-        meta = {"lm_hidden": lm_hidden, "prefix_feat_cond": prefix_feat_cond}
+        meta = _PrefillResidualMeta(lm_hidden=lm_hidden, prefix_feat_cond=prefix_feat_cond)
         return residual_input.squeeze(0), meta
 
-    def _prepare_residual_decode(self, state: _RequestState, base_lm_out: torch.Tensor, dev: Any):
+    def _prepare_residual_decode(
+        self,
+        state: _RequestState,
+        base_lm_out: torch.Tensor,
+        dev: torch.device,
+    ) -> tuple[torch.Tensor, _DecodeResidualMeta]:
         tts = self.tts
         state.decode_step_count += 1
 
@@ -1774,14 +2227,14 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         if prev.ndim == 1:
             prev = prev.unsqueeze(0)
         res_input = tts.fusion_concat_proj(torch.cat([lm_h, prev], dim=-1))
-        return res_input, {"new_lm_hidden": lm_h}
+        return res_input, _DecodeResidualMeta(new_lm_hidden=lm_h)
 
     def _prepare_residual_decode_batch(
         self,
         states: list[_RequestState],
         base_lm_outs: list[torch.Tensor],
-        dev: Any,
-    ) -> list[tuple[torch.Tensor, dict]]:
+        dev: torch.device,
+    ) -> list[tuple[torch.Tensor, _DecodeResidualMeta]]:
         tts = self.tts
 
         for state in states:
@@ -1811,7 +2264,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         return [
             (
                 res_batch[i : i + 1],
-                {"new_lm_hidden": lm_h_batch[i : i + 1]},
+                _DecodeResidualMeta(new_lm_hidden=lm_h_batch[i : i + 1]),
             )
             for i in range(len(states))
         ]
@@ -1906,14 +2359,96 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         out.normal_(generator=gen)
         state.cfm_noise_step += 1
 
-    def _finish_prefill(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any):
+    def _finish_prefill_batch(self, batch: list[tuple[_RequestState, _PrefillResidualMeta, torch.Tensor]]) -> None:
+        if len(batch) == 1:
+            state, meta, res_out = batch[0]
+            self._finish_prefill(state, meta, res_out, None)
+            return
+
+        self._perf.start("prefill_tail_batch")
         tts = self.tts
-        lm_hidden = meta["lm_hidden"]
-        prefix_feat_cond = meta["prefix_feat_cond"]
+        self._setup_cfm_buffers()
+        if self._enable_torch_compile:
+            self._setup_torch_compile()
+
+        dit_hs: list[torch.Tensor] = []
+        conds: list[torch.Tensor] = []
+        stop_logits_list: list[torch.Tensor] = []
+
+        for state, meta, res_out in batch:
+            lm_hidden = meta.lm_hidden
+            prefix_feat_cond = meta.prefix_feat_cond
+            residual_hidden = res_out[-1:, :]
+            stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach()
+            dit_h = torch.cat([tts.lm_to_dit_proj(lm_hidden), tts.res_to_dit_proj(residual_hidden)], dim=-1)
+            cond = prefix_feat_cond.transpose(1, 2).contiguous()
+            dit_hs.append(dit_h)
+            conds.append(cond)
+            stop_logits_list.append(stop_logits)
+
+        batched_dit_h = torch.cat(dit_hs, dim=0)
+        batched_cond = torch.cat(conds, dim=0)
+        b = batched_dit_h.shape[0]
+
+        if self._deterministic_cfm_noise and self._cfm_buffers is not None:
+            noise = self._cfm_buffers.noise[:b]
+            for i, (state, _, _) in enumerate(batch):
+                self._fill_deterministic_cfm_noise(state, noise[i : i + 1])
+        else:
+            noise = torch.randn(
+                b,
+                self._feat_dim,
+                self._patch_size,
+                device=batched_dit_h.device,
+                dtype=batched_dit_h.dtype,
+            )
+        pred_feats = _optimized_solve_euler_with_noise(
+            tts.feat_decoder,
+            batched_dit_h,
+            self._patch_size,
+            batched_cond,
+            noise,
+            self._inference_timesteps,
+            self._cfg_value,
+            self._cfm_buffers,
+            cfg_cutoff_ratio=self._cfg_cutoff_ratio,
+        ).transpose(1, 2)
+
+        with torch.no_grad():
+            curr_embeds = tts.enc_to_lm_proj(tts.feat_encoder(pred_feats.unsqueeze(1))).squeeze(1)
+
+        for i, (state, meta, res_out) in enumerate(batch):
+            pred_feat_i = pred_feats[i : i + 1]
+            curr_embed_i = curr_embeds[i : i + 1]
+            self._commit_decode_state(state, stop_logits_list[i], curr_embed_i, pred_feat_i)
+            state.decode_step_count = 0
+            state.request_start_time = time.perf_counter()
+            state.prefill_completed = True
+
+        self._perf.stop("prefill_tail_batch")
+        if self._enable_profiling:
+            logger.info("PREFILL_BATCH[%d requests] tail breakdown:\n%s", len(batch), self._perf.breakdown())
+            self._perf.reset()
+
+    def _finish_prefill(
+        self,
+        state: _RequestState,
+        meta: _PrefillResidualMeta,
+        res_out: torch.Tensor,
+        dev: torch.device | None,
+    ) -> None:
+        self._perf.start("prefill_tail")
+        tts = self.tts
+        lm_hidden = meta.lm_hidden
+        prefix_feat_cond = meta.prefix_feat_cond
         residual_hidden = res_out[-1:, :]
 
+        self._perf.start("prefill.stop_fn")
         stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach()
+        self._perf.stop("prefill.stop_fn")
+        self._perf.start("prefill.dit_proj")
         dit_h = torch.cat([tts.lm_to_dit_proj(lm_hidden), tts.res_to_dit_proj(residual_hidden)], dim=-1)
+        self._perf.stop("prefill.dit_proj")
 
         self._setup_cfm_buffers()
         if self._enable_torch_compile:
@@ -1921,24 +2456,29 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
 
         pred_feat = self._run_cfm_for_state(state, dit_h, prefix_feat_cond.transpose(1, 2).contiguous())
 
+        self._perf.start("prefill.feat_encoder")
         with torch.no_grad():
             curr_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1)
+        self._perf.stop("prefill.feat_encoder")
 
         self._commit_decode_state(state, stop_logits, curr_embed, pred_feat)
         state.decode_step_count = 0
         state.request_start_time = time.perf_counter()
         state.prefill_completed = True
+        self._perf.stop("prefill_tail")
 
         if logger.isEnabledFor(logging.DEBUG):
             # Only compute the norm (which forces a GPU->CPU sync) if we will log it.
             logger.debug("PREFILL[%s]: patch norm=%.4f", state.request_id, pred_feat.norm().item())
+        if self._enable_profiling:
+            logger.info("PREFILL[%s] tail breakdown:\n%s", state.request_id, self._perf.breakdown())
         self._perf.reset()
 
-    def _finish_decode(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any):
+    def _finish_decode(self, state: _RequestState, meta: _DecodeResidualMeta, res_out: torch.Tensor) -> None:
         self._perf.start("decode_step")
         tts = self.tts
 
-        lm_h = meta["new_lm_hidden"]
+        lm_h = meta.new_lm_hidden
         res_h = res_out.unsqueeze(0) if res_out.ndim == 1 else res_out
 
         dit_proj = getattr(self, "_compiled_dit_proj", None) or self._dit_proj_fn
@@ -1959,7 +2499,11 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         if self._enable_profiling and state.decode_step_count % 20 == 0:
             logger.info("Step %d[%s]:\n%s", state.decode_step_count, state.request_id, self._perf.breakdown())
 
-    def _finish_decode_batch(self, req_metas: list[tuple], batch_out: torch.Tensor):
+    def _finish_decode_batch(
+        self,
+        req_metas: list[tuple[_RequestState, bool, _DecodeResidualMeta]],
+        batch_out: torch.Tensor,
+    ) -> None:
         self._perf.start("decode_step")
         tts = self.tts
 
@@ -1968,7 +2512,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
 
         if self._enable_batched_cfm and not self._enable_cfm_cuda_graph:
             states = [state for state, _, _ in req_metas]
-            lm_h = torch.cat([meta["new_lm_hidden"] for _, _, meta in req_metas], dim=0)
+            lm_h = torch.cat([meta.new_lm_hidden for _, _, meta in req_metas], dim=0)
             pfc = torch.cat(
                 [
                     state.curr_prefix_feat_cond.to(self._side_dtype).unsqueeze(0)
@@ -2016,7 +2560,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             return
 
         for i, (state, _, meta) in enumerate(req_metas):
-            lm_h = meta["new_lm_hidden"]
+            lm_h = meta.new_lm_hidden
             res_h = batch_out[i : i + 1]
             pfc = state.curr_prefix_feat_cond.to(self._side_dtype)
             if pfc.ndim == 2:
@@ -2474,7 +3018,10 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
     # -------------------- preprocess / postprocess --------------------
 
     def preprocess(
-        self, input_ids: torch.Tensor, input_embeds: torch.Tensor | None, **info_dict: Any
+        self,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor | None,
+        **info_dict: Unpack[VoxCPM2PreprocessInput],
     ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]:
         additional = info_dict.get("additional_information")
         if isinstance(additional, dict):
@@ -2570,34 +3117,28 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
                 )
 
             if state.prompt_cache is None and (ref_audio or (prompt_audio and prompt_text)):
-                try:
-                    state.prompt_cache = self._build_prompt_cache(
-                        ref_audio=ref_audio,
-                        prompt_audio=prompt_audio,
-                        prompt_text=prompt_text,
-                    )
-                    if (
-                        voice_name
-                        and state.prompt_cache is not None
-                        and state.prompt_cache.get("mode") == "reference"
-                        and "ref_audio_feat" in state.prompt_cache
-                    ):
-                        _key = self._speaker_cache.make_cache_key(
-                            voice_name, model_type="voxcpm2", created_at=_created_at
-                        )
-                        self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()})
-                        logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
-                except Exception as e:
-                    logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e)
-                    state.prompt_cache = None
+                state.prompt_cache = self._build_prompt_cache(
+                    ref_audio=ref_audio,
+                    prompt_audio=prompt_audio,
+                    prompt_text=prompt_text,
+                )
+                if (
+                    voice_name
+                    and state.prompt_cache is not None
+                    and state.prompt_cache.get("mode") == "reference"
+                    and "ref_audio_feat" in state.prompt_cache
+                ):
+                    _key = self._speaker_cache.make_cache_key(voice_name, model_type="voxcpm2", created_at=_created_at)
+                    self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()})
+                    logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name)
 
             inputs = self._build_prefill_inputs(token_ids, dev, req_id)
             tts = self.tts
-            feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs["audio_feat"]))
-            text_embed = self.model.embed_input_ids(inputs["text_token"].to(dev))
-            text_mask, feat_mask = inputs["text_mask"], inputs["audio_mask"]
+            feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs.audio_feat))
+            text_embed = self.model.embed_input_ids(inputs.text_token.to(dev))
+            text_mask, feat_mask = inputs.text_mask, inputs.audio_mask
             embeds = (text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed).squeeze(0)
-            state.prefill_masks = (text_mask, feat_mask, inputs["audio_feat"], feat_embed)
+            state.prefill_masks = (text_mask, feat_mask, inputs.audio_feat, feat_embed)
         else:
             state = self._active_states.get(req_id)
             curr = state.curr_embed_for_next if state else None
@@ -2609,23 +3150,31 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
         self._pending_requests.append((req_id, is_prefill, embeds, span_len))
         return input_ids, embeds, {}
 
-    def postprocess(self, hidden_states: torch.Tensor, **info: Any) -> dict[str, Any]:
+    def postprocess(self, hidden_states: torch.Tensor, **info: Unpack[VoxCPM2PostprocessInput]) -> dict[str, Any]:
         req_id = info.get("request_id", self._current_request_id or "default")
         if self._enable_profiling:
             state = self._active_states.get(req_id)
             if state and state.decode_step_count > 0:
                 logger.info(
-                    "REQUEST DONE[%s]: %d steps, %.2fs\n%s",
+                    "REQUEST DONE[%s]: %d steps, %.2fs\n%s\nUnified graph: captures=%d replays=%d skips=%s",
                     req_id,
                     state.decode_step_count,
                     time.perf_counter() - state.request_start_time,
                     self._perf.breakdown(),
+                    self._unified_graph_stats.captures,
+                    self._unified_graph_stats.replays,
+                    self._unified_graph_stats.skips,
                 )
         return {}
 
     # -------------------- build prefill inputs --------------------
 
-    def _build_prefill_inputs(self, token_ids: list[int], dev: Any, req_id: str = "default") -> dict:
+    def _build_prefill_inputs(
+        self,
+        token_ids: list[int],
+        dev: torch.device,
+        req_id: str = "default",
+    ) -> _PrefillInputs:
         tts = self.tts
         dtype = self._side_dtype
         state = self._active_states.get(req_id)
@@ -2673,12 +3222,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module):
             text_mask = torch.cat([rtm.cpu(), ones_t, zeros_p])
             audio_mask = torch.cat([ram.cpu(), zeros_t, ones_p])
 
-        return {
-            "text_token": text_token.unsqueeze(0).to(dev),
-            "audio_feat": audio_feat.unsqueeze(0).to(dev).to(dtype),
-            "text_mask": text_mask.unsqueeze(0).to(dev),
-            "audio_mask": audio_mask.unsqueeze(0).to(dev),
-        }
+        return _PrefillInputs(
+            text_token=text_token.unsqueeze(0).to(dev),
+            audio_feat=audio_feat.unsqueeze(0).to(dev).to(dtype),
+            text_mask=text_mask.unsqueeze(0).to(dev),
+            audio_mask=audio_mask.unsqueeze(0).to(dev),
+        )
 
     # -------------------- weight loading --------------------
 
diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
index f262ede61..40d4ec620 100644
--- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
+++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py
@@ -209,7 +209,12 @@ class VoxtralTTSForConditionalGeneration(
 
     def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict | None):
         self.post_process_idx = 0
-        audio_tokens = info_dict.pop("audio", None)
+        codes = info_dict.get("codes")
+        audio_tokens = codes.get("audio") if isinstance(codes, Mapping) else None
+        if audio_tokens is None:
+            # Backward compatible with request state written before #4527 moved
+            # Voxtral TTS feedback frames under ``codes.audio``.
+            audio_tokens = info_dict.pop("audio", None)
         if audio_tokens is not None:
             kwargs = {"audio_tokens": audio_tokens.to(input_ids.device)}
             multimodal_embeddings = self.model.embed_multimodal(**kwargs)
diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml
deleted file mode 100644
index f5e449253..000000000
--- a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml
+++ /dev/null
@@ -1,51 +0,0 @@
-# Stage config for running MammothModa2 with full AR→DiT pipeline (text-to-image).
-# Stage 0: AR  (multimodal understanding + latent generation)
-# Stage 1: DiT (latent → image)
-#
-# For text/image understanding tasks (text output only), use mammoth_moda2_ar.yaml.
-
-stage_args:
-  - stage_id: 0
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: ar
-      max_num_seqs: 100
-      model_arch: MammothModa2ForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 8192
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      enable_prefix_caching: false
-    final_output: false
-
-  - stage_id: 1
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: dit
-      max_num_seqs: 1
-      model_arch: MammothModa2ForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.3
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: image
-      enable_prefix_caching: false
-    # Declares extra_args so has_sampling_extra_args is enabled, routing the t2i knobs
-    # (text_guidance_scale, cfg_range, num_inference_steps) declared in
-    # vllm_omni/model_extras/mammothmodal2_preview.py from extra_body into the DiT
-    # forward. Values here are the defaults applied when a request omits a knob.
-    default_sampling_params:
-      extra_args:
-        text_guidance_scale: 9.0
-        cfg_range: [0.0, 1.0]
-        num_inference_steps: 50
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit
-    final_output: true
-    final_output_type: image
diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml
deleted file mode 100644
index d455c7973..000000000
--- a/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml
+++ /dev/null
@@ -1,24 +0,0 @@
-# Stage config for running MammothModa2 with single AR stage (understanding / summarization).
-# Stage 0: AR (multimodal understanding + text generation)
-#
-# Use this config for image/video understanding tasks where only text output is needed.
-# For text-to-image generation (AR → DiT), use mammoth_moda2.yaml instead.
-
-stage_args:
-  - stage_id: 0
-    runtime:
-      devices: "0"
-    engine_args:
-      model_stage: ar
-      max_num_seqs: 16
-      model_arch: MammothModa2ForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 8192
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: text
-      enable_prefix_caching: false
-    final_output: true
-    final_output_type: text
diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py
index 0477ea4e1..2956be94f 100644
--- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py
+++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py
@@ -119,30 +119,6 @@ def _decode_additional_information(raw_info: Any) -> dict[str, Any]:
     return decoded
 
 
-def text2flow(
-    source_outputs: list[Any],
-    prompt: OmniTokensPrompt | TextPrompt = None,
-    _requires_multimodal_data: bool = True,
-):
-    """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs."""
-    engine_inputs: list[OmniTokensPrompt] = []
-    for source_output in source_outputs:
-        output = source_output.outputs[0]
-        multi_modal_data = output.multimodal_output
-        if multi_modal_data is None:
-            raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
-
-        prefix_ids = _ensure_list(source_output.prompt_token_ids)
-        raw_output_ids = _ensure_list(output.cumulative_token_ids)
-        prompt_speech_ids = _prompt_speech_token_ids(multi_modal_data)
-        output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids)
-        output_ids = _strip_prompt_prefix(output_ids, prompt_speech_ids)
-        additional_info = dict(multi_modal_data)
-        additional_info.setdefault("ids", {})["prompt"] = prefix_ids
-        engine_inputs.append(OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=additional_info))
-    return engine_inputs
-
-
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,
diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py
index 668ff2ed4..f428c34a4 100644
--- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py
+++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py
@@ -114,37 +114,6 @@ def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | Non
     raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}")
 
 
-def slow_ar_to_dac_decoder(
-    source_outputs: list[Any],
-    _prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder."""
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    slow_ar_outputs = source_outputs
-    dac_inputs: list[OmniTokensPrompt] = []
-
-    for output in slow_ar_outputs:
-        out = output.outputs[0]
-        # audio_codes shape: [num_frames, num_codebooks]
-        audio_codes = out.multimodal_output["audio_codes"].to(torch.long)
-        # Filter zero-padded frames.
-        valid_mask = audio_codes.any(dim=1)
-        audio_codes = audio_codes[valid_mask]
-        # Codebook-major flat: [num_codebooks * num_frames]
-        codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
-        dac_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=None,
-            )
-        )
-    return dac_inputs
-
-
 def slow_ar_to_dac_decoder_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,
diff --git a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py
index 9a81a8cc3..963adb9c2 100644
--- a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py
+++ b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py
@@ -95,128 +95,130 @@ def ar2diffusion(
     source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | list | None = None,
     requires_multimodal_data: bool = False,
-) -> list[dict[str, Any]]:
+) -> dict[str, Any] | None:
     """Process AR stage outputs to create Diffusion stage inputs.
 
+    HunyuanImage3 produces one downstream diffusion request per parent AR
+    request. ``source_outputs`` may include companion outputs, but the parent
+    output is expected first: the orchestrator bundles diffusion inputs as
+    ``[parent_output, *cfg_companion_outputs]`` before invoking this bridge.
+
     Args:
         prompt: Original user prompt (may contain multimodal data).
         requires_multimodal_data: Whether to forward multimodal data.
 
     Returns:
-        List of dicts, each consumable by the HunyuanImage3 diffusion pipeline.
+        One prompt dict consumable by the HunyuanImage3 diffusion pipeline, or
+        ``None`` when no parent output is available.
     """
-    ar_outputs = source_outputs
-    diffusion_inputs = []
+    if not source_outputs:
+        return None
 
-    # Normalize prompt to list
-    if not isinstance(prompt, list):
-        prompt = [prompt] if prompt is not None else [{}]
+    # The orchestrator constructs this list as [parent, *cfg_companions].
+    ar_output = source_outputs[0]
+    output = ar_output.outputs[0]
+    generated_token_ids = output.cumulative_token_ids
+    # Prefer cumulative_text, fallback to text if aggregation dropped it
+    generated_text = getattr(output, "cumulative_text", None) or getattr(output, "text", "") or ""
 
-    for i, ar_output in enumerate(ar_outputs):
-        output = ar_output.outputs[0]
-        generated_token_ids = output.cumulative_token_ids
-        # Prefer cumulative_text, fallback to text if aggregation dropped it
-        generated_text = getattr(output, "cumulative_text", None) or getattr(output, "text", "") or ""
+    if isinstance(prompt, list):
+        original_prompt = prompt[0] if prompt else {}
+    elif prompt is not None:
+        original_prompt = prompt
+    else:
+        original_prompt = {}
+    if isinstance(original_prompt, dict):
+        pass
+    elif hasattr(original_prompt, "_asdict"):
+        original_prompt = original_prompt._asdict()
+    elif hasattr(original_prompt, "__dict__"):
+        original_prompt = vars(original_prompt)
+    else:
+        original_prompt = {}
 
-        # Get original prompt info
-        original_prompt = prompt[i] if i < len(prompt) else {}
-        if isinstance(original_prompt, dict):
-            pass
-        elif hasattr(original_prompt, "_asdict"):
-            original_prompt = original_prompt._asdict()
-        elif hasattr(original_prompt, "__dict__"):
-            original_prompt = vars(original_prompt)
-        else:
-            original_prompt = {}
+    height = original_prompt.get("height", 1024)
+    width = original_prompt.get("width", 1024)
+    text_prompt = original_prompt.get("prompt", "")
+    use_system_prompt = original_prompt.get("use_system_prompt")
+    custom_system_prompt = original_prompt.get("system_prompt")
 
-        height = original_prompt.get("height", 1024)
-        width = original_prompt.get("width", 1024)
-        text_prompt = original_prompt.get("prompt", "")
-        use_system_prompt = original_prompt.get("use_system_prompt")
-        custom_system_prompt = original_prompt.get("system_prompt")
+    # Prefer the AR's predicted output aspect (`<img_size_*><img_ratio_*>`
+    # tail emitted by `HunyuanImage3ForCausalMM.sample` under the
+    # ratio-restriction logits processor) over the carried-through
+    # height/width, which the serving layer fills with the first
+    # reference image's bucket and so collapses non-square targets to
+    # square in the multi-image / mismatched-aspect case. Mirrors the
+    # official upstream where `reso_group[ratio_index]` is the
+    # canonical source of the diffusion target shape.
+    ratio_idx = _extract_ratio_index(generated_token_ids)
+    ar_predicted = False
+    if ratio_idx is not None:
+        base_size = int(original_prompt.get("image_base_size", 1024))
+        reso_group: ResolutionGroup = get_cached_resolution_group(base_size=base_size)
+        try:
+            reso: Resolution = reso_group[ratio_idx]
+            height = reso.height
+            width = reso.width
+            ar_predicted = True
+        except IndexError:
+            logger.warning(
+                "[ar2diffusion] Request 0: ratio_index=%d out of range [0,%d), keeping prompt size %dx%d",
+                ratio_idx,
+                len(reso_group),
+                height,
+                width,
+            )
 
-        # Prefer the AR's predicted output aspect (`<img_size_*><img_ratio_*>`
-        # tail emitted by `HunyuanImage3ForCausalMM.sample` under the
-        # ratio-restriction logits processor) over the carried-through
-        # height/width, which the serving layer fills with the first
-        # reference image's bucket and so collapses non-square targets to
-        # square in the multi-image / mismatched-aspect case. Mirrors the
-        # official upstream where `reso_group[ratio_index]` is the
-        # canonical source of the diffusion target shape.
-        ratio_idx = _extract_ratio_index(generated_token_ids)
-        ar_predicted = False
-        if ratio_idx is not None:
-            base_size = int(original_prompt.get("image_base_size", 1024))
-            reso_group: ResolutionGroup = get_cached_resolution_group(base_size=base_size)
-            try:
-                reso: Resolution = reso_group[ratio_idx]
-                height = reso.height
-                width = reso.width
-                ar_predicted = True
-            except IndexError:
-                logger.warning(
-                    "[ar2diffusion] Request %d: ratio_index=%d out of range [0,%d), keeping prompt size %dx%d",
-                    i,
-                    ratio_idx,
-                    len(reso_group),
-                    height,
-                    width,
-                )
+    cot_text_for_dit = _truncate_at_cot_end(generated_text)
 
-        cot_text_for_dit = _truncate_at_cot_end(generated_text)
+    logger.info(
+        "[ar2diffusion] Request 0: AR generated %d tokens, text length=%d, cot_text length=%d, target size=%dx%d (%s)",
+        len(generated_token_ids),
+        len(generated_text),
+        len(cot_text_for_dit),
+        height,
+        width,
+        f"AR ratio_idx={ratio_idx}" if ar_predicted else "from prompt (no AR ratio token)",
+    )
 
-        logger.info(
-            "[ar2diffusion] Request %d: AR generated %d tokens, text length=%d, "
-            "cot_text length=%d, target size=%dx%d (%s)",
-            i,
-            len(generated_token_ids),
-            len(generated_text),
-            len(cot_text_for_dit),
-            height,
-            width,
-            f"AR ratio_idx={ratio_idx}" if ar_predicted else "from prompt (no AR ratio token)",
-        )
+    diffusion_input: dict[str, Any] = {
+        "prompt": text_prompt,
+        "height": height,
+        "width": width,
+        "extra": {
+            "ar_generated_text": cot_text_for_dit,
+        },
+    }
 
-        diffusion_input: dict[str, Any] = {
-            "prompt": text_prompt,
-            "height": height,
-            "width": width,
-            "extra": {
-                "ar_generated_text": cot_text_for_dit,
-            },
-        }
+    # Forward use_system_prompt so the DiT can build the same system prefix.
+    # Also forward the custom system prompt body when sys_type=custom so
+    # DiT's `get_system_prompt(use, "image", body)` doesn't fall back to
+    # an empty prefix and silently diverge from AR.
+    if use_system_prompt is not None:
+        diffusion_input["use_system_prompt"] = use_system_prompt
+    if custom_system_prompt is not None:
+        diffusion_input["system_prompt"] = custom_system_prompt
 
-        # Forward use_system_prompt so the DiT can build the same system prefix.
-        # Also forward the custom system prompt body when sys_type=custom so
-        # DiT's `get_system_prompt(use, "image", body)` doesn't fall back to
-        # an empty prefix and silently diverge from AR.
-        if use_system_prompt is not None:
-            diffusion_input["use_system_prompt"] = use_system_prompt
-        if custom_system_prompt is not None:
-            diffusion_input["system_prompt"] = custom_system_prompt
+    # Forward multimodal data (original image for IT2I conditioning).
+    # The diffusion pre_process_func reads multi_modal_data["image"], which
+    # matches vLLM's standard prompt schema, so we only need to pass it once.
+    mm_data = original_prompt.get("multi_modal_data")
+    if mm_data:
+        prompt_images = mm_data.get("image")
+        if prompt_images is None:
+            prompt_images = mm_data.get("images")
+        if prompt_images is not None:
+            diffusion_input["multi_modal_data"] = {"image": prompt_images}
 
-        # Forward multimodal data (original image for IT2I conditioning).
-        # The diffusion pre_process_func reads multi_modal_data["image"], which
-        # matches vLLM's standard prompt schema, so we only need to pass it once.
-        mm_data = original_prompt.get("multi_modal_data")
-        if mm_data:
-            prompt_images = mm_data.get("image")
-            if prompt_images is None:
-                prompt_images = mm_data.get("images")
-            if prompt_images is not None:
-                diffusion_input["multi_modal_data"] = {"image": prompt_images}
+    # Forward multimodal output from AR (if any)
+    if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
+        mm_output = ar_output.multimodal_output
+        if isinstance(mm_output, Mapping):
+            diffusion_input["extra"]["ar_multimodal_output"] = mm_output
 
-        # Forward multimodal output from AR (if any)
-        if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output:
-            mm_output = ar_output.multimodal_output
-            if isinstance(mm_output, Mapping):
-                diffusion_input["extra"]["ar_multimodal_output"] = mm_output
+    # Forward sampling params
+    for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]:
+        if key in original_prompt:
+            diffusion_input[key] = original_prompt[key]
 
-        # Forward sampling params
-        for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]:
-            if key in original_prompt:
-                diffusion_input[key] = original_prompt[key]
-
-        diffusion_inputs.append(diffusion_input)
-
-    return diffusion_inputs
+    return diffusion_input
diff --git a/vllm_omni/model_executor/stage_input_processors/indextts2.py b/vllm_omni/model_executor/stage_input_processors/indextts2.py
index 9d38c8f2a..39d9d011c 100644
--- a/vllm_omni/model_executor/stage_input_processors/indextts2.py
+++ b/vllm_omni/model_executor/stage_input_processors/indextts2.py
@@ -143,65 +143,6 @@ def _request_id(request: Any) -> str:
     return str(getattr(request, "external_req_id", None) or getattr(request, "request_id", "?"))
 
 
-def talker2s2mel(
-    source_outputs: list[Any],
-    prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Legacy orchestrator path: collect all Stage 0 output, format Stage 1 input."""
-    _old_nt = torch.get_num_threads()
-    torch.set_num_threads(1)
-    try:
-        return _talker2s2mel_impl(source_outputs)
-    finally:
-        torch.set_num_threads(_old_nt)
-
-
-def _talker2s2mel_impl(source_outputs: list[Any]) -> list[Any]:
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    s2mel_inputs: list[OmniTokensPrompt] = []
-
-    for i, talker_output in enumerate(source_outputs):
-        if not talker_output.finished:
-            continue
-
-        output = talker_output.outputs[0]
-        mm = output.multimodal_output
-        if not isinstance(mm, dict):
-            logger.warning("Talker output %d has no multimodal_output dict", i)
-            continue
-
-        codes_dict = mm.get("codes", {})
-        mel_codes = codes_dict.get("mel")
-        if not isinstance(mel_codes, torch.Tensor) or mel_codes.numel() == 0:
-            logger.warning("Talker output %d has empty mel_codes", i)
-            continue
-        mel_codes = mel_codes.to(torch.long)
-
-        hs_dict = mm.get("hidden_states", {})
-        latent = hs_dict.get("latent")
-        if not isinstance(latent, torch.Tensor) or latent.numel() == 0:
-            logger.warning("Talker output %d has empty latent", i)
-            continue
-
-        meta = mm.get("meta", {})
-        if not isinstance(meta, dict):
-            meta = {}
-
-        additional_information = _build_s2mel_additional_information(mel_codes, latent, meta, context="talker2s2mel")
-        s2mel_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=[0],
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=additional_information,
-            )
-        )
-
-    return s2mel_inputs
-
-
 def talker2s2mel_token_only(
     source_outputs: list[Any],
     prompt: Any = None,
diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py
index 7c0db4131..3a1a4441f 100644
--- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py
+++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py
@@ -1,7 +1,6 @@
 from typing import Any
 
 import torch
-from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
@@ -10,7 +9,6 @@ from vllm_omni.data_entry_keys import (
     OmniPayload,
     OmniPayloadStruct,
 )
-from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID
 
 logger = init_logger(__name__)
@@ -230,88 +228,6 @@ def llm2code2wav_async_chunk(
     )
 
 
-def llm2code2wav(
-    source_outputs: list[Any],
-    prompt: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,
-) -> list[OmniTokensPrompt]:
-    """
-    Process talker outputs to create code2wav inputs.
-
-    Workflow:
-    1. Extract talker's codec code outputs (8-layer RVQ codes)
-    2. Flatten codes for code2wav input
-    3. Package for code2wav stage
-
-    Args:
-        prompt: Original prompt data
-        requires_multimodal_data: Whether multimodal data is required
-
-    Returns:
-        List of OmniTokensPrompt for code2wav stage
-    """
-    talker_outputs = source_outputs
-    code2wav_inputs = []
-
-    # Process each talker output
-    for i, talker_output in enumerate(talker_outputs):
-        output = talker_output.outputs[0]
-
-        # Extract codec codes from talker output
-        # Expected shape: [8, seq_len] (8-layer RVQ codes)
-        mm = output.multimodal_output
-        mm_codes = mm.get("codes", {})
-        mm_hs = mm.get("hidden_states", {})
-        if "audio" in mm_codes:
-            codec_codes = mm_codes["audio"].to(torch.long)  # [seq_batch_size, 1, 8, 4]
-            is_all_zero = (codec_codes == 0).all(dim=(1, 2, 3))
-            non_zero_indices = (~is_all_zero).nonzero(as_tuple=True)[0]
-            if len(non_zero_indices) == 0:
-                # All codec codes are zero - skip this request with a warning
-                request_id = getattr(talker_output, "request_id", f"unknown_{i}")
-                logger.warning(
-                    "Skipping request %s: all codec codes are zero (empty output from Stage-0). "
-                    "This may indicate the model failed to generate valid audio codes.",
-                    request_id,
-                )
-            else:
-                if len(non_zero_indices) < codec_codes.shape[0]:
-                    codec_codes = codec_codes[non_zero_indices]
-        elif "output" in mm_hs and "audio" not in mm_codes:
-            codec_codes = torch.zeros(1, 1, 8, 4, dtype=torch.long)
-        else:
-            raise ValueError(f"Invalid multimodal_output: {output.multimodal_output}")
-
-        pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4)
-
-        code_final = prepend_and_flatten_colmajor(codec_codes, pad_vec)
-        code_final = code_final.tolist()
-
-        # Guard against flattened sequences longer than code2wav's max_model_len.
-        # Without this, add_request raises ``could not broadcast input array
-        # from shape (N,) into shape (max_model_len,)`` and kills the engine
-        # core (see issue #2683). Mirrors the offline end2end.py safeguard.
-        if len(code_final) > MAX_CODE2WAV_TOKENS:
-            request_id = getattr(talker_output, "request_id", f"unknown_{i}")
-            logger.warning(
-                "Request %s: code_final len=%d > MAX_CODE2WAV_TOKENS=%d, truncating.",
-                request_id,
-                len(code_final),
-                MAX_CODE2WAV_TOKENS,
-            )
-            code_final = code_final[:MAX_CODE2WAV_TOKENS]
-
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=code_final,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return code2wav_inputs
-
-
 # ============================================================================
 # Worker-connector data plane (non-async-chunk path).
 # AR runner's `flatten_payload` converts the model emit
@@ -329,8 +245,9 @@ _FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset()
 def _filter_zero_codec_rows(codec_codes: torch.Tensor) -> torch.Tensor:
     """Drop zero-padded codec rows from a 4-D `[N, 1, 8, 4]` tensor.
 
-    Mirrors the zero-row filter in the orchestrator-path `llm2code2wav`
-    body (see this file's ``llm2code2wav`` around line 224).
+    Shared by the sync placeholder builder (``llm2code2wav_token_only``) and
+    the full-payload producer (``llm2code2wav_full_payload``) so both size the
+    downstream codec sequence off the same non-zero frames.
     """
     if codec_codes.ndim != 4 or codec_codes.numel() == 0:
         return codec_codes
diff --git a/vllm_omni/model_executor/stage_input_processors/ming_tts.py b/vllm_omni/model_executor/stage_input_processors/ming_tts.py
index 0f2effdcb..058cda788 100644
--- a/vllm_omni/model_executor/stage_input_processors/ming_tts.py
+++ b/vllm_omni/model_executor/stage_input_processors/ming_tts.py
@@ -12,6 +12,7 @@ from vllm.logger import init_logger
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
 from vllm_omni.inputs.data import OmniTokensPrompt
 from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
+    INITIAL_LATENT_CHUNK_SIZE,
     KEY_CHUNK_ID,
     KEY_REQUEST_ID,
     LATENT_CHUNK_SIZE,
@@ -116,7 +117,7 @@ def _decode_stop_reason(value: Any) -> str | None:
     return MING_STOP_REASON_BY_CODE.get(int(value))
 
 
-def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]:
+def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int, int]:
     connector = getattr(transfer_manager, "connector", None)
     raw_cfg = getattr(connector, "config", {}) or {}
     cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, Mapping) else {}
@@ -131,9 +132,19 @@ def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]:
             LATENT_LEFT_CONTEXT,
         )
     chunk_size = int(cfg.get("latent_chunk_size", LATENT_CHUNK_SIZE))
+    initial_chunk_size = int(cfg.get("initial_latent_chunk_size", INITIAL_LATENT_CHUNK_SIZE))
     left_context = int(cfg.get("latent_left_context", LATENT_LEFT_CONTEXT))
     if chunk_size <= 0:
         raise ValueError(f"Invalid Ming latent_chunk_size={chunk_size}")
+    if initial_chunk_size < 0:
+        raise ValueError(f"Invalid Ming initial_latent_chunk_size={initial_chunk_size}")
+    if initial_chunk_size > chunk_size:
+        logger.warning(
+            "Ming initial_latent_chunk_size=%d > latent_chunk_size=%d, clamping to latent_chunk_size.",
+            initial_chunk_size,
+            chunk_size,
+        )
+        initial_chunk_size = chunk_size
     # Stage-2 VAE caches past_key_values and stream_state by request_id.
     # Replaying left-context latents would double-feed cached decoder state.
     if left_context != 0:
@@ -142,7 +153,7 @@ def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]:
             "Ming boundary continuity is handled by per-request decoder state cache, not "
             f"latent replay. Got latent_left_context={left_context}."
         )
-    return chunk_size, left_context
+    return chunk_size, initial_chunk_size, left_context
 
 
 def _build_chunk_observability(
@@ -196,7 +207,7 @@ def llm2audio_vae_async_chunk(
     if patch is not None:
         transfer_manager.code_prompt_token_ids[request_id].append(patch)
 
-    chunk_size, _ = _get_async_chunk_config(transfer_manager)
+    chunk_size, initial_chunk_size, _ = _get_async_chunk_config(transfer_manager)
 
     patches = transfer_manager.code_prompt_token_ids[request_id]
     seen_patch_len = int(state.get("seen_patch_len", 0))
@@ -225,11 +236,17 @@ def llm2audio_vae_async_chunk(
             )
         return None
 
-    chunk_length = length % chunk_size
-    if chunk_length != 0 and not finished:
-        return None
-
-    emit_count = chunk_length if chunk_length != 0 else chunk_size
+    use_first_chunk = 0 < initial_chunk_size < chunk_size and seen_patch_len == 0
+    if finished:
+        emit_count = length
+    elif use_first_chunk:
+        if length < initial_chunk_size:
+            return None
+        emit_count = initial_chunk_size
+    else:
+        if length < chunk_size:
+            return None
+        emit_count = chunk_size
     emit_patches = list(new_patches[:emit_count])
     state["seen_patch_len"] = seen_patch_len + len(emit_patches)
     latent_patches = torch.stack(emit_patches, dim=0)
diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py
index 48c0f98cb..3dfe2a603 100644
--- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py
+++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py
@@ -5,6 +5,8 @@
 
 from __future__ import annotations
 
+from collections import defaultdict
+from collections.abc import Mapping
 from typing import Any
 
 import torch
@@ -15,6 +17,8 @@ from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct
 
 logger = init_logger(__name__)
 
+_MOSS_AUDIO_PAD_CODE = 1024
+
 
 # ---------------------------------------------------------------------------
 # Helpers
@@ -90,13 +94,13 @@ def talker2codec(
 # ---------------------------------------------------------------------------
 
 
-def talker2codec_async_chunk(
+def talker2codec_delay_async_chunk(
     transfer_manager: Any,
     multimodal_output: dict[str, Any] | None,
     request: Any,
     is_finished: bool = False,
 ) -> OmniPayloadStruct | None:
-    """Emit accumulated audio codes to Stage 1 as they arrive from Stage 0.
+    """Emit delay-patterned accumulated audio codes to Stage 1.
 
     State is maintained in ``transfer_manager`` keyed by request ID.
     A chunk is forwarded to Stage 1 when either:
@@ -138,14 +142,6 @@ def talker2codec_async_chunk(
                     req_state["accumulated"] = new_rows
                 else:
                     req_state["accumulated"] = torch.cat([req_state["accumulated"], new_rows], dim=0)
-        # Realtime variant emits raw (un-delay-patterned) codes; record that
-        # so the emit step below skips the apply_de_delay_pattern transform.
-        # Realtime sets ``meta.finished`` to a 1-D bool tensor; the delay
-        # variant leaves ``meta`` empty.
-        meta_in = pooling_output.get("meta", {}) or {}
-        flag = meta_in.get("finished")
-        if isinstance(flag, torch.Tensor) and flag.numel() >= 1 and bool(flag.reshape(-1)[0].item()):
-            req_state["skip_dedelay"] = True
 
     acc = req_state["accumulated"]
     if acc is None or acc.numel() == 0:
@@ -181,15 +177,12 @@ def talker2codec_async_chunk(
     #   1. de-delay   ``(T+nq-1, nq)`` → ``(T, nq)``
     #   2. drop rows that are entirely pad (separators between audio segments,
     #      and the leading text-mode rows that precede the first audio_start).
-    # The realtime talker emits raw codes (no delay), so steps 1 is skipped.
     chunk_codes_long = chunk_codes.to(torch.long).cpu().contiguous()  # (T_chunk, NQ)
     nq = int(chunk_codes_long.shape[1])
     t_chunk = int(chunk_codes_long.shape[0])
     audio_pad_code = 1024  # MOSS-TTS audio_pad_code; same value across variants.
 
-    if req_state.get("skip_dedelay"):
-        de_delayed = chunk_codes_long
-    elif t_chunk > nq:
+    if t_chunk > nq:
         de_delayed = chunk_codes_long.new_zeros((t_chunk - nq + 1, nq))
         for i in range(nq):
             de_delayed[:, i] = chunk_codes_long[i : i + de_delayed.shape[0], i]
@@ -224,9 +217,133 @@ def talker2codec_async_chunk(
         codes=CodesStruct(audio=codec_flat),
         meta=MetaStruct(
             left_context_size=left_context,
-            finished=torch.tensor(is_finished, dtype=torch.bool),
+            finished=torch.tensor(bool(is_finished), dtype=torch.bool),
         ),
     )
 
 
-__all__ = ["talker2codec", "talker2codec_async_chunk"]
+def talker2codec_raw_async_chunk(
+    transfer_manager: Any,
+    multimodal_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> OmniPayloadStruct | None:
+    """Async processor for MOSS-TTS Local/Realtime raw codec rows.
+
+    Stage 0 emits newly generated raw codec rows shaped ``[T, n_vq]`` (normally
+    ``[1, n_vq]`` per decode step). This processor buffers those new rows until
+    a codec chunk is ready, then forwards the chunk to Stage 1. No delay-pattern
+    de-delay is applied on this path.
+    """
+    external_req_id = getattr(request, "external_req_id", None)
+    req_id = str(external_req_id if external_req_id is not None else getattr(request, "request_id", id(request)))
+
+    if not hasattr(transfer_manager, "code_prompt_token_ids"):
+        transfer_manager.code_prompt_token_ids = defaultdict(list)
+    if not hasattr(transfer_manager, "request_payload"):
+        transfer_manager.request_payload = {}
+    if not hasattr(transfer_manager, "put_req_chunk"):
+        transfer_manager.put_req_chunk = defaultdict(int)
+
+    pending_frames = transfer_manager.code_prompt_token_ids[req_id]
+
+    if isinstance(multimodal_output, Mapping):
+        codes_dict = multimodal_output.get("codes", {}) or {}
+        new_frames = codes_dict.get("audio")
+        if isinstance(new_frames, torch.Tensor) and new_frames.numel() > 0:
+            frames_cpu = new_frames.detach().to("cpu", torch.long).contiguous()
+            if frames_cpu.ndim == 1:
+                frames_cpu = frames_cpu.reshape(1, -1)
+            if frames_cpu.ndim != 2:
+                raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_cpu.shape)}")
+            valid_rows = frames_cpu.ne(_MOSS_AUDIO_PAD_CODE).any(dim=1)
+            for frame in frames_cpu[valid_rows]:
+                pending_frames.append(frame.clone())
+        # Raw/local streaming should mirror the non-streaming path: the codec
+        # decodes only generated audio rows. Reference audio conditions the
+        # talker, but feeding its codes into the codec streaming state adds a
+        # long first-packet prime step and changes the decoder state relative
+        # to non-streaming output.
+
+    connector = getattr(transfer_manager, "connector", None)
+    raw_cfg = getattr(connector, "config", {}) or {}
+    cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+    cfg = cfg if isinstance(cfg, dict) else {}
+    chunk_frames = int(cfg.get("codec_chunk_frames", 15) or 15)
+    initial_chunk_frames = int(cfg.get("initial_codec_chunk_frames") or 0)
+    if chunk_frames <= 0:
+        raise ValueError(f"codec_chunk_frames must be positive for MOSS raw streaming, got {chunk_frames}")
+    if initial_chunk_frames < 0:
+        raise ValueError(f"initial_codec_chunk_frames must be non-negative, got {initial_chunk_frames}")
+    if initial_chunk_frames > chunk_frames:
+        logger.warning(
+            "initial_codec_chunk_frames=%d > codec_chunk_frames=%d, clamping.",
+            initial_chunk_frames,
+            chunk_frames,
+        )
+        initial_chunk_frames = chunk_frames
+
+    pending = len(pending_frames)
+    emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0
+    threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames
+    if pending <= 0:
+        if is_finished:
+            transfer_manager.code_prompt_token_ids.pop(req_id, None)
+            transfer_manager.request_payload.pop(req_id, None)
+            return OmniPayloadStruct(
+                # A non-empty sentinel is required so Stage-1 is scheduled.
+                # ``code_flat_numel=0`` tells the codec this is a control-only
+                # finish packet, not an audio code.
+                codes=CodesStruct(audio=torch.tensor([0], dtype=torch.long)),
+                meta=MetaStruct(
+                    req_id=[req_id],
+                    left_context_size=0,
+                    codec_streaming=True,
+                    codec_chunk_frames=0,
+                    codec_left_context_frames=0,
+                    code_flat_numel=0,
+                    stream_finished=torch.tensor(True, dtype=torch.bool),
+                    finished=torch.tensor(True, dtype=torch.bool),
+                ),
+                request_id=req_id,
+            )
+        return None
+    if not is_finished and pending < threshold:
+        return None
+
+    emit_frames = pending if is_finished else threshold
+    chunk_rows = pending_frames[:emit_frames]
+    del pending_frames[:emit_frames]
+    chunk_codes = torch.stack(
+        [row.to(torch.long).cpu() for row in chunk_rows],
+        dim=0,
+    ).contiguous()
+    finished = bool(is_finished and len(pending_frames) == 0)
+
+    codec_flat = chunk_codes.transpose(0, 1).contiguous().reshape(-1).to(torch.long)
+
+    if finished:
+        transfer_manager.code_prompt_token_ids.pop(req_id, None)
+        transfer_manager.request_payload.pop(req_id, None)
+
+    return OmniPayloadStruct(
+        codes=CodesStruct(audio=codec_flat),
+        meta=MetaStruct(
+            req_id=[req_id],
+            left_context_size=0,
+            codec_streaming=True,
+            codec_chunk_frames=int(chunk_codes.shape[0]),
+            codec_left_context_frames=0,
+            code_flat_numel=int(codec_flat.numel()),
+            stream_finished=torch.tensor(finished, dtype=torch.bool),
+            finished=torch.tensor(finished, dtype=torch.bool),
+        ),
+        request_id=req_id,
+    )
+
+
+__all__ = [
+    "talker2codec",
+    "talker2codec_delay_async_chunk",
+    "talker2codec_raw_async_chunk",
+]
diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
index d624b4a49..edde6972d 100644
--- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
+++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
@@ -10,7 +10,6 @@ from typing import Any
 
 import torch
 from vllm.inputs import TextPrompt
-from vllm.platforms import current_platform
 
 from vllm_omni.data_entry_keys import (
     CodesStruct,
@@ -569,25 +568,31 @@ def thinker2talker_full_payload(
         output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
         all_token_ids = list(prompt_token_ids) + list(output_token_ids)
 
-    # Trim the trailing stop-token row from the accumulated thinker output.
-    # The accumulator captures one hidden-state row per executed thinker
-    # forward (prefill + every decode step including the one that emitted
-    # the stop_token), so for a finished request thinker_emb has exactly one
-    # row more than the rows the talker should consume. async_chunk's
-    # chunk-0 path naturally captures only the prefill / non-stop portion,
-    # which is why the [async_chunk] parametrization passes while [default]
-    # over-generates one codec frame on short outputs (e.g.
-    # test_one_word_prompt_001[default]: audio extends "London" with
-    # spurious phonemes).
-    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 0:
+    # Drop the terminal stop-token row only when more than one row was
+    # accumulated; trimming a single row would ship 0 conditioning tensors
+    # while ids still has tokens and break talker prefill alignment.
+    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 1:
         thinker_emb_prefill = thinker_emb[:-1]
     else:
         thinker_emb_prefill = thinker_emb
-    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 0:
+    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 1:
         thinker_hid_prefill = thinker_hid[:-1]
     else:
         thinker_hid_prefill = thinker_hid
 
+    emb_rows = int(thinker_emb_prefill.shape[0]) if isinstance(thinker_emb_prefill, torch.Tensor) else 0
+    hid_rows = int(thinker_hid_prefill.shape[0]) if isinstance(thinker_hid_prefill, torch.Tensor) else 0
+    if len(all_token_ids) > 0 and (emb_rows == 0 or hid_rows == 0):
+        logger.warning(
+            "thinker2talker_full_payload: empty thinker conditioning for req=%s "
+            "(ids_len=%s embed_rows=%s hidden_rows=%s); withholding payload.",
+            rid,
+            len(all_token_ids),
+            emb_rows,
+            hid_rows,
+        )
+        return None
+
     payload: OmniPayload = {
         "embed": {
             "prefill": thinker_emb_prefill.detach().cpu(),
@@ -599,7 +604,6 @@ def thinker2talker_full_payload(
         "ids": {"all": list(all_token_ids), "prompt": list(prompt_token_ids)},
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
     }
-    payload["next_stage_prompt_len"] = _compute_talker_prompt_ids_length(payload, device="cpu")
     speaker = extract_speaker_from_request(request)
     if speaker is not None:
         payload["speaker"] = speaker
@@ -609,154 +613,30 @@ def thinker2talker_full_payload(
     return payload
 
 
-def thinker2talker(
-    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.
-
-    Workflow:
-    1. Extract thinker's text generation outputs (token IDs + hidden states)
-    2. Split hidden states into: prompt embeddings + generated embeddings
-    3. Package for talker with additional information
-
-    In PD disaggregation mode, merges prefill-stage prompt embeddings with
-    decode-stage generated embeddings before handing off to the talker.
-
-    Args:
-        prompt: Original prompt data
-        requires_multimodal_data: Whether multimodal data is required
-
-    Returns:
-        List of OmniTokensPrompt for talker stage
-    """
-    thinker_outputs = source_outputs
-    talker_inputs: list[OmniTokensPrompt] = []
-
-    device = torch.device(current_platform.device_type)
-
-    # Process each thinker output
-    for i, thinker_output in enumerate(thinker_outputs):
-        output = thinker_output.outputs[0]
-        req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
-        prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
-        output_ids = _ensure_list(output.cumulative_token_ids)
-        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
-        if is_streaming_session:
-            prompt_token_ids, output_ids = _get_streaming_talker_tokens(
-                req_id,
-                prompt_token_ids,
-                output_ids,
-                getattr(streaming_context, "new_prompt_len_snapshot", None),
-                streaming_context,
-                clear_state=bool(getattr(thinker_output, "finished", False)),
-            )
-        thinker_sequences = prompt_token_ids + output_ids
-        thinker_input_ids = prompt_token_ids
-        new_seq_length = len(prompt_token_ids + output_ids) - 1
-        thinker_mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(thinker_mm_raw, Mapping):
-            logger.debug("thinker2talker: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        thinker_mm: OmniPayload = thinker_mm_raw
-        mm_hs = thinker_mm.get("hidden_states", {})
-        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
-        emb_layer = _layer_tensor(mm_layers, _EMBED_LAYER_KEY)
-        hid_layer = _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY)
-        if emb_layer is None or hid_layer is None:
-            logger.debug("thinker2talker: skip req=%s due to missing hidden-state layers", req_id)
-            continue
-        thinker_emb = emb_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
-        thinker_hid = hid_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
-
-        prefill_mm: dict[str, Any] | None = None
-        prefill_mm = _get_prefill_multimodal_output(req_id, streaming_context)
-
-        if prefill_mm is not None:
-            expected_total = len(prompt_token_ids) + len(output_ids)
-            try:
-                thinker_emb, thinker_hid = _merge_pd_embeddings(
-                    thinker_emb, thinker_hid, prefill_mm, device, expected_total=expected_total
-                )
-            except Exception as exc:
-                logger.warning("[PD] Could not merge prefill embeddings: %s", exc)
-
-        payload = OmniPayloadStruct(
-            embed=EmbeddingsStruct(
-                prefill=thinker_emb,
-                tts_bos=_resolve_tts_token_embedding(
-                    "tts_bos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-                tts_eos=_resolve_tts_token_embedding(
-                    "tts_eos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-                tts_pad=_resolve_tts_token_embedding(
-                    "tts_pad", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device
-                ),
-            ),
-            hidden_states=HiddenStatesStruct(output=thinker_hid),
-            ids=IdsStruct(all=thinker_sequences, prompt=thinker_input_ids),
-            speaker=extract_speaker_from_prompt(prompt, index=i),
-            language=extract_language_from_prompt(prompt, index=i),
-        )
-        info = to_dict(payload)
-        prompt_len = _compute_talker_prompt_ids_length(info, device=device)
-
-        talker_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=[0] * prompt_len,
-                additional_information=info,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return talker_inputs
-
-
 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]:
-    """Non-async-chunk Stage-1 input builder for the connector data plane.
+    """Orchestrator-side placeholder builder for Stage-1 (Talker) when
+    ``async_chunk=False``.
 
-    The worker connector (Stage-0 ``thinker2talker_full_payload`` →
-    ``_sync_local_stage_payloads``) supplies the bulk talker conditioning
-    tensors (embed / hidden_states / ids) via ``model_intermediate_buffer``.
-    The orchestrator only needs to ship a placeholder prefill prompt of the
-    correct length so the scheduler can allocate KV-cache slots.
+    After the communication-layer refactor, this function only allocates a
+    placeholder ``prompt_token_ids`` of the correct length so the scheduler can
+    reserve KV-cache slots. It does **not** forward bulk tensors.
 
-    Small per-request voice metadata (``speaker`` / ``language``) is forwarded
-    here from the user prompt so the worker's line-408 buffer seed picks it
-    up. The connector-side ``extract_speaker_from_request`` reads the
-    strongly-typed ``request.additional_information.entries["speaker"]`` which
-    currently does not always round-trip the user-supplied voice; until that
-    plumbing is normalized, providing the small fields directly preserves
-    voice selection (regression discovered on Buildkite 9668:
-    ``test_speaker_002[default]`` lost the preset voice).
+    Bulk talker conditioning is sent through the connector. Speaker and
+    language are also copied from the original prompt so they survive when
+    Stage-0 request metadata is unavailable to the connector payload.
+
+    ``prompt`` / ``requires_multimodal_data`` are kept for call-site signature
+    compatibility with other orchestrator input processors; they are unused.
     """
     talker_inputs: list[OmniTokensPrompt] = []
     for i, thinker_output in enumerate(source_outputs):
         output = thinker_output.outputs[0]
         req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
-        # Skip-on-missing parity with thinker2talker_full_payload: if the
-        # connector builder would drop this request (no MM dict or missing
-        # hidden-state layers), do the same here so the worker buffer
-        # presence agrees with the orchestrator's scheduling decision.
-        thinker_mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(thinker_mm_raw, Mapping):
-            logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        mm_hs = thinker_mm_raw.get("hidden_states", {})
-        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {}
-        if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None:
-            logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id)
-            continue
         prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
         output_ids = _ensure_list(output.cumulative_token_ids)
         is_streaming_session = bool(getattr(streaming_context, "enabled", False))
@@ -773,21 +653,17 @@ def thinker2talker_token_only(
         thinker_input_ids = prompt_token_ids
         info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
         prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
-
-        # Forward only small voice metadata; bulk tensors come from the
-        # connector path via _sync_local_stage_payloads.
-        small_info: dict[str, Any] = {}
-        speaker = extract_speaker_from_prompt(prompt, index=i)
-        if speaker is not None:
-            small_info["speaker"] = speaker
-        language = extract_language_from_prompt(prompt, index=i)
-        if language is not None:
-            small_info["language"] = language
-
+        # Keep this fallback until the connector reliably preserves voice metadata.
+        additional_information = to_dict(
+            OmniPayloadStruct(
+                speaker=extract_speaker_from_prompt(prompt, index=i),
+                language=extract_language_from_prompt(prompt, index=i),
+            )
+        )
         talker_inputs.append(
             OmniTokensPrompt(
                 prompt_token_ids=[0] * prompt_len,
-                additional_information=(small_info if small_info else None),
+                additional_information=additional_information or None,
                 multi_modal_data=None,
                 mm_processor_kwargs=None,
             )
@@ -949,59 +825,5 @@ def talker2code2wav_full_payload(
     )
     return {
         "codes": {"audio": codec_codes},
-        "code_predictor_codes": codec_codes,
         "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
     }
-
-
-def talker2code2wav(
-    source_outputs: list[Any],
-    _prompt: OmniTokensPrompt | TextPrompt | None = None,
-    _requires_multimodal_data: bool = False,
-    streaming_context: Any | None = None,
-) -> list[OmniTokensPrompt]:
-    """
-    Process talker outputs to create code2wav inputs.
-
-    Workflow:
-    1. Extract talker's codec code outputs (8-layer RVQ codes)
-    2. Flatten codes for code2wav input
-    3. Package for code2wav stage
-
-    Args:
-    Returns:
-        List of OmniTokensPrompt for code2wav stage
-    """
-    talker_outputs = source_outputs
-    code2wav_inputs: list[OmniTokensPrompt] = []
-    # Process each talker output
-    for i, talker_output in enumerate(talker_outputs):
-        output = talker_output.outputs[0]
-        req_id = str(getattr(talker_output, "request_id", f"idx-{i}"))
-        cur_seq_len = len(output.cumulative_token_ids) - 1
-        seq_len = cur_seq_len
-        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
-        if is_streaming_session:
-            seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context)
-        mm_raw = getattr(output, "multimodal_output", None)
-        if not isinstance(mm_raw, Mapping):
-            logger.debug("talker2code2wav: skip req=%s due to empty multimodal_output", req_id)
-            continue
-        mm: OmniPayload = mm_raw
-        if "codes" not in mm or not isinstance(mm.get("codes"), dict) or "audio" not in mm["codes"]:
-            logger.debug("talker2code2wav: skip req=%s due to missing codes.audio", req_id)
-            continue
-        # Extract codec codes from talker output
-        # Expected shape: [8, seq_len] (8-layer RVQ codes)
-        codec_codes = (
-            mm["codes"]["audio"][-seq_len:].to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist()
-        )  # 16, seq_len
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return code2wav_inputs
diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
index 6cebc4c28..597fc4db5 100644
--- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
+++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
@@ -38,103 +38,6 @@ def _qwen3_tts_empty_finished_payload():
     }
 
 
-def talker2code2wav(
-    source_outputs: list[Any],
-    prompt: Any = None,
-    _requires_multimodal_data: bool = False,
-) -> list[Any]:
-    """Non-async: collect all talker codes, then pass to code2wav at once."""
-    from vllm_omni.inputs.data import OmniTokensPrompt
-
-    talker_outputs = source_outputs
-    code2wav_inputs: list[OmniTokensPrompt] = []
-    for i, talker_output in enumerate(talker_outputs):
-        if not talker_output.finished:
-            # Non-async decode should only run once, after talker has
-            # accumulated the final code sequence.
-            continue
-        output = talker_output.outputs[0]
-        mm = output.multimodal_output
-        mm_codes = mm.get("codes", {})
-
-        # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16)
-        audio_codes = mm_codes["audio"].to(torch.long)
-        token_ids = output.cumulative_token_ids
-
-        # token_ids provides an upper bound on the newly generated codec span.
-        # audio_codes may still contain zero-padded / invalid rows, so trim only
-        # after filtering valid frames instead of trying to align EOS indices.
-        seq_len = max(len(token_ids) - 1, 0)
-        # Filter invalid frames: zero-padded (EOS), out-of-range values (e.g.
-        # stop_token_id=2150 exceeds codebook_size=2048), and negative
-        # sentinels (e.g. -1 padding).
-        _CODEBOOK_SIZE = 2048
-        valid_mask = (
-            (audio_codes >= 0).all(dim=1) & audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE)
-        )
-        audio_codes = audio_codes[valid_mask]
-        if seq_len > 0 and audio_codes.ndim == 2 and int(audio_codes.shape[0]) > seq_len:
-            audio_codes = audio_codes[-seq_len:]
-        ref_code = mm_codes.get("ref")
-        ref_code_len = mm.get("meta", {}).get("ref_code_len")
-        if isinstance(ref_code_len, torch.Tensor):
-            ref_code_len = int(ref_code_len.reshape(-1)[-1].item()) if ref_code_len.numel() > 0 else 0
-        elif ref_code_len is None:
-            ref_code_len = 0
-        else:
-            ref_code_len = int(ref_code_len)
-        if isinstance(ref_code, list):
-            ref_code = ref_code[0] if ref_code else None
-        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
-            ref_code = ref_code.to(torch.long).cpu().contiguous()
-            if ref_code.ndim == 1:
-                num_quantizers = int(audio_codes.shape[1]) if audio_codes.ndim == 2 and audio_codes.shape[1] > 0 else 16
-                if ref_code.numel() % num_quantizers != 0:
-                    logger.warning(
-                        "Ignoring malformed ref_code with %d elements not divisible by num_quantizers=%d",
-                        ref_code.numel(),
-                        num_quantizers,
-                    )
-                    ref_code = None
-                else:
-                    ref_code = ref_code.reshape(-1, num_quantizers)
-            elif ref_code.ndim != 2:
-                logger.warning("Ignoring malformed ref_code shape %s", tuple(ref_code.shape))
-                ref_code = None
-            if isinstance(ref_code, torch.Tensor) and ref_code_len > 0 and int(ref_code.shape[0]) > ref_code_len:
-                logger.warning(
-                    "Trimming ref_code from %d frames to ref_code_len=%d before Code2Wav.",
-                    int(ref_code.shape[0]),
-                    ref_code_len,
-                )
-                ref_code = ref_code[:ref_code_len]
-            if not isinstance(ref_code, torch.Tensor):
-                ref_code_len = 0
-            else:
-                ref_code_len = int(ref_code.shape[0])
-                audio_codes = torch.cat([ref_code.to(audio_codes.device), audio_codes], dim=0)
-        else:
-            ref_code_len = 0
-        # Code2Wav expects codebook-major flat: [Q*num_frames]
-        codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist()
-        additional_information = to_dict(
-            OmniPayloadStruct(
-                meta=MetaStruct(left_context_size=ref_code_len) if ref_code_len > 0 else None,
-                speaker=extract_speaker_from_prompt(prompt, index=i),
-                language=extract_language_from_prompt(prompt, index=i),
-            )
-        )
-        code2wav_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=codec_codes,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-                additional_information=additional_information if additional_information else None,
-            )
-        )
-    return code2wav_inputs
-
-
 def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
     audio_codes = multimodal_output.get("codes", {}).get("audio")
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
@@ -351,10 +254,7 @@ _NUM_QUANTIZERS_DEFAULT = 16
 
 
 def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
-    """Filter zero-padded, out-of-range, and negative-padded codec frames.
-
-    Mirrors the orchestrator-path body in `talker2code2wav` above.
-    """
+    """Filter zero-padded, out-of-range, and negative-padded codec frames."""
     if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0:
         return audio_codes
     if audio_codes.ndim != 2:
@@ -367,10 +267,8 @@ def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor:
 
 def _coerce_ref_code_len(raw) -> int:
     """Coerce mm["meta"]["ref_code_len"] / pooling_output["meta.ref_code_len"]
-    raw value (Tensor | int | None) into a non-negative int.  Mirrors the
-    extraction inlined in the legacy ``talker2code2wav`` path; clamps any
-    negative input to 0 since downstream code treats this as a non-negative
-    frame count."""
+    raw value (Tensor | int | None) into a non-negative int; clamps any
+    negative input to 0 since downstream code treats this as a frame count."""
     if isinstance(raw, torch.Tensor):
         value = int(raw.reshape(-1)[-1].item()) if raw.numel() > 0 else 0
     elif raw is None:
@@ -381,7 +279,7 @@ def _coerce_ref_code_len(raw) -> int:
 
 
 def _normalize_ref_code(ref_code, num_quantizers: int, ref_code_len: int):
-    """Coerce ref_code into a [ref_len, Q] tensor or None.  Mirrors orchestrator path."""
+    """Coerce ref_code into a [ref_len, Q] tensor or None."""
     if isinstance(ref_code, list):
         ref_code = ref_code[0] if ref_code else None
     if not isinstance(ref_code, torch.Tensor) or ref_code.numel() == 0:
@@ -407,8 +305,7 @@ def talker2code2wav_token_only(
 
     Sized to the expected codec token count (codebook-major flat:
     Q * (ref_frames + audio_frames)).  Speaker / language metadata are
-    extracted from `prompt` and threaded via `additional_information`
-    (orchestrator-style; same as the legacy `talker2code2wav` builder).
+    extracted from `prompt` and threaded via `additional_information`.
     Actual codec ids are delivered via the worker connector payload built
     by `talker2code2wav_full_payload`.
     """
@@ -473,8 +370,8 @@ def talker2code2wav_full_payload(
     Reads accumulated codec from `pooling_output["codes.audio"]` (CONCAT
     across steps via flatten_payload), latest `pooling_output["codes.ref"]`
     (prefill-emitted), and latest `pooling_output["meta.ref_code_len"]`.
-    Replicates the orchestrator-path body of `talker2code2wav` (filter,
-    crop to seq_len, prepend ref, codebook-major flatten).
+    Filters invalid frames, crops to seq_len, prepends ref, and flattens
+    codebook-major for code2wav consumption.
     """
     del transfer_manager
     rid = getattr(request, "request_id", "?")
diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
index d3c1219b7..af2cb0a25 100644
--- a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
+++ b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
@@ -2,7 +2,6 @@ from collections.abc import Mapping
 from typing import Any
 
 import torch
-from vllm.inputs import TextPrompt
 from vllm.logger import init_logger
 
 from vllm_omni.data_entry_keys import (
@@ -11,7 +10,6 @@ from vllm_omni.data_entry_keys import (
     OmniPayload,
     OmniPayloadStruct,
 )
-from vllm_omni.inputs.data import OmniTokensPrompt
 
 logger = init_logger(__name__)
 
@@ -31,29 +29,6 @@ def _codec_audio_tensors(multimodal_output: OmniPayload | dict[str, Any]) -> lis
     return None
 
 
-def generator2tokenizer(
-    source_outputs: list[Any],
-    _prompt: OmniTokensPrompt | TextPrompt = None,
-    _requires_multimodal_data: bool = False,
-):
-    tokenizer_inputs = []
-    for generator_output in source_outputs:
-        output = generator_output.outputs[0]
-        audio_tensors = _codec_audio_tensors(output.multimodal_output)
-        if not audio_tensors:
-            raise ValueError("Missing codes.audio in generator multimodal_output")
-        audio_tokens = torch.cat(audio_tensors, dim=-1).flatten().detach().cpu().tolist()
-        tokenizer_inputs.append(
-            OmniTokensPrompt(
-                prompt_token_ids=audio_tokens,
-                multi_modal_data=None,
-                mm_processor_kwargs=None,
-            )
-        )
-
-    return tokenizer_inputs
-
-
 def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None:
     audio_tensors = _codec_audio_tensors(multimodal_output)
     if not audio_tensors:
diff --git a/vllm_omni/model_extras/audiox.py b/vllm_omni/model_extras/audiox.py
index 2ab7bc4d8..220972148 100644
--- a/vllm_omni/model_extras/audiox.py
+++ b/vllm_omni/model_extras/audiox.py
@@ -19,7 +19,7 @@ AUDIOX_EXTRA_BODY_PARAMS = frozenset(
         "audio_path",  # optional reference-audio conditioning
     }
 )
-# Echoed back in ``DiffusionOutput.custom_output`` (pipeline_audiox.py).
+# Echoed back in diffusion output metadata (pipeline_audiox.py).
 AUDIOX_EXTRA_OUTPUT_PARAMS = frozenset(
     {
         "audiox_task",
diff --git a/vllm_omni/model_extras/ming_flash_omni.py b/vllm_omni/model_extras/ming_flash_omni.py
new file mode 100644
index 000000000..8fece255d
--- /dev/null
+++ b/vllm_omni/model_extras/ming_flash_omni.py
@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from typing import Any
+
+from PIL import Image
+
+
+def build_text_to_image_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    height: int | None = None,
+    width: int | None = None,
+) -> dict[str, Any]:
+    text_prompt: dict[str, Any] = {
+        "prompt": prompt,
+        "modalities": ["image"],
+        "mm_processor_kwargs": {"modalities": ["image"]},
+    }
+    if height is not None:
+        text_prompt["mm_processor_kwargs"]["target_h"] = height
+    if width is not None:
+        text_prompt["mm_processor_kwargs"]["target_w"] = width
+    if negative_prompt is not None:
+        text_prompt["negative_prompt"] = negative_prompt
+    return text_prompt
+
+
+def build_image_to_image_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    input_image: Image.Image | list[Image.Image],
+    height: int | None = None,
+    width: int | None = None,
+) -> dict[str, Any]:
+    img_prompt: dict[str, Any] = {
+        "prompt": prompt,
+        "modalities": ["img2img"],
+        "multi_modal_data": {"img2img": input_image},
+        "mm_processor_kwargs": {"modalities": ["img2img"]},
+    }
+    if height is not None:
+        img_prompt["mm_processor_kwargs"]["target_h"] = height
+    if width is not None:
+        img_prompt["mm_processor_kwargs"]["target_w"] = width
+    if negative_prompt is not None:
+        img_prompt["negative_prompt"] = negative_prompt
+    return img_prompt
+
+
+# For Image-generation (diffusion-stage)
+MING_FLASH_OMNI_EXTRA_BODY_PARAMS = frozenset(
+    {
+        "height",
+        "width",
+        "steps",
+        "cfg",
+        "seed",
+        "byte5_text",
+        "negative_prompt",
+    }
+)
+MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS = frozenset()
+MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True
diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py
index cec304213..79f8b4360 100644
--- a/vllm_omni/model_extras/registry.py
+++ b/vllm_omni/model_extras/registry.py
@@ -46,6 +46,17 @@ from vllm_omni.model_extras.mammothmodal2_preview import (
 from vllm_omni.model_extras.mammothmodal2_preview import (
     build_text_to_image_prompt as build_mammothmoda2_text_to_image_prompt,
 )
+from vllm_omni.model_extras.ming_flash_omni import (
+    MING_FLASH_OMNI_EXTRA_BODY_PARAMS,
+    MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS,
+    MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
+)
+from vllm_omni.model_extras.ming_flash_omni import (
+    build_image_to_image_prompt as build_ming_flash_omni_image_to_image_prompt,
+)
+from vllm_omni.model_extras.ming_flash_omni import (
+    build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt,
+)
 from vllm_omni.model_extras.sensenova_u1 import (
     SENSENOVA_U1_EXTRA_BODY_PARAMS,
     SENSENOVA_U1_EXTRA_OUTPUT_PARAMS,
@@ -142,6 +153,11 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
         "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt,
     },
+    "Cosmos3OmniPipeline": {
+        "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
+        "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
+        "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt,
+    },
     "MagiHumanPipeline": {
         "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
         "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
@@ -165,6 +181,13 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "init_extra_args_for_non_diffusion_stages": MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
         "text_to_image_prompt_builder": build_mammothmoda2_text_to_image_prompt,
     },
+    "MingImagePipeline": {
+        "extra_body_params": MING_FLASH_OMNI_EXTRA_BODY_PARAMS,
+        "extra_output_params": MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS,
+        "init_extra_args_for_non_diffusion_stages": MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
+        "text_to_image_prompt_builder": build_ming_flash_omni_text_to_image_prompt,
+        "image_to_image_prompt_builder": build_ming_flash_omni_image_to_image_prompt,
+    },
 }
 
 
diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs/__init__.py
similarity index 100%
rename from vllm_omni/outputs.py
rename to vllm_omni/outputs/__init__.py
diff --git a/vllm_omni/outputs/mm_outputs.py b/vllm_omni/outputs/mm_outputs.py
new file mode 100644
index 000000000..8c5fe6442
--- /dev/null
+++ b/vllm_omni/outputs/mm_outputs.py
@@ -0,0 +1,261 @@
+"""Multimodal output data structures for vLLM-Omni.
+
+This module defines structured types for multimodal outputs.
+
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+from vllm.outputs import CompletionOutput
+
+from vllm_omni.outputs.output_modality import TensorAccumulationStrategy
+from vllm_omni.outputs.utils import _is_tensor_list, _to_cpu
+
+logger = init_logger(__name__)
+
+# Keys whose values are metadata scalars (e.g. audio sample rate) but may
+# arrive as 0-d torch.Tensors — from_dict routes all tensors into .tensors,
+# so we relocate them to .metadata before consolidation to avoid a bogus
+# torch.cat attempt and its warn-and-keep-last fallback.
+_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"})
+
+
+def _cat_tensors(
+    tensors: list[torch.Tensor],
+    strategy: TensorAccumulationStrategy,
+) -> torch.Tensor:
+    """Concatenate a list of tensors according to *strategy*."""
+    if strategy == TensorAccumulationStrategy.CONCAT_LAST:
+        return torch.cat(tensors, dim=-1)
+    if strategy == TensorAccumulationStrategy.REPLACE:
+        return tensors[-1]
+    # CONCAT_DIM0 / APPEND_LIST / default
+    return torch.cat(tensors, dim=0)
+
+
+def _consolidate_tensor_list(
+    key: str,
+    tensor_list: list[torch.Tensor],
+    strategy: TensorAccumulationStrategy,
+) -> torch.Tensor:
+    """Concatenate a deferred tensor list, with fallbacks on shape mismatch."""
+    try:
+        return _cat_tensors(tensor_list, strategy)
+    except RuntimeError:
+        # [TODO: this part is for async chunk, not for audio only. can we come up with a design
+        # that identify this by whether using async chunk instead of modality?]
+        if key != "audio":
+            logger.warning("Error concatenating tensor for key %s; keeping last tensor", key)
+            return tensor_list[-1]
+        # Audio chunks may have mismatched shapes; retry along the last dim,
+        # then fall back to flattening each chunk.
+        try:
+            return torch.cat(tensor_list, dim=-1)
+        except RuntimeError:
+            return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0)
+
+
+def _append_entries(store: dict[str, Any], incoming: dict[str, Any]) -> None:
+    """Append *incoming* entries into *store*.
+
+    Tensor values are collected into lists (deferred concatenation, avoiding
+    O(n²) repeated torch.cat); non-tensor values replace the existing entry.
+    """
+    for key, new_value in incoming.items():
+        existing = store.get(key)
+        if isinstance(existing, list) and isinstance(new_value, torch.Tensor):
+            existing.append(new_value)
+        elif isinstance(existing, torch.Tensor) and isinstance(new_value, torch.Tensor):
+            store[key] = [existing, new_value]
+        else:
+            store[key] = new_value
+
+
+@dataclass(eq=False)
+class MultimodalPayload(Mapping):
+    """Structured multimodal output payload.
+
+    Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)``
+    style checks in downstream code can be replaced with duck-typing, and
+    ``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)``
+    all work seamlessly for both tensors and metadata.
+
+    Attributes:
+        tensors: Dictionary mapping modality/key names to their tensors.
+        metadata: Optional dictionary for non-tensor metadata
+            (e.g., sample rate for audio, image dimensions).
+    """
+
+    tensors: dict[str, torch.Tensor] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def primary_tensor(self) -> torch.Tensor | None:
+        """Return the first tensor in the payload, or None if empty."""
+        if self.tensors:
+            return next(iter(self.tensors.values()))
+        return None
+
+    @property
+    def is_empty(self) -> bool:
+        """Return True if the payload has no tensors and no metadata."""
+        return len(self.tensors) == 0 and len(self.metadata) == 0
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get a value by key, searching tensors first then metadata."""
+        if key in self.tensors:
+            return self.tensors[key]
+        return self.metadata.get(key, default)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.tensors or key in self.metadata
+
+    def __getitem__(self, key: str) -> Any:
+        """Dict-like indexing: search tensors first, then metadata."""
+        if key in self.tensors:
+            return self.tensors[key]
+        if key in self.metadata:
+            return self.metadata[key]
+        raise KeyError(key)
+
+    def __len__(self) -> int:
+        return len(self.tensors) + len(self.metadata)
+
+    def __iter__(self) -> Iterator[str]:
+        yield from self.tensors
+        yield from self.metadata
+
+    def __bool__(self) -> bool:
+        return bool(self.tensors) or bool(self.metadata)
+
+    def __eq__(self, other: object) -> bool:
+        """Support equality with plain dicts and other Mappings."""
+        if isinstance(other, MultimodalPayload):
+            return self.tensors == other.tensors and self.metadata == other.metadata
+        if isinstance(other, Mapping):
+            return dict(self) == dict(other)
+        return NotImplemented
+
+    def merged_with(self, incoming: MultimodalPayload) -> MultimodalPayload:
+        """Merge *incoming* onto this payload and return the result.
+
+        Tensor values accumulate into lists for deferred concatenation;
+        non-tensor values are replaced with the latest. When this payload
+        is empty, *incoming* is returned as-is, so callers should use the
+        return value: ``accumulated = accumulated.merged_with(incoming)``.
+        """
+        if self.is_empty:
+            return incoming
+        _append_entries(self.tensors, incoming.tensors)
+        _append_entries(self.metadata, incoming.metadata)
+        return self
+
+    def consolidate_tensors(self, strategy: TensorAccumulationStrategy) -> None:
+        """Concatenate deferred tensor lists into single tensors.
+
+        Tensors are generated content accumulated as chunks, so lists are
+        concatenated according to *strategy* (e.g. audio chunks along the
+        time dimension, latent frames along the batch dimension).
+        """
+        # Relocate scalar metadata tensors (e.g. sample rate) that from_dict
+        # routed into .tensors, so they take the REPLACE path via .metadata
+        # instead of failing torch.cat as 0-d tensors.
+        for key in _METADATA_TENSOR_KEYS:
+            value = self.tensors.pop(key, None)
+            if value is None:
+                continue
+            self.metadata[key] = value[-1] if isinstance(value, list) else value
+
+        for key, value in list(self.tensors.items()):
+            if _is_tensor_list(value):
+                self.tensors[key] = _consolidate_tensor_list(key, value, strategy)
+
+    def consolidate_metadata(self) -> None:
+        """Resolve deferred tensor lists in metadata by keeping the latest value.
+
+        Metadata values are per-step snapshots (e.g. sample rate), not
+        content deltas, so the latest value supersedes earlier ones.
+        Nested dicts (unflattened payloads) are resolved one level down.
+        """
+        for key, value in list(self.metadata.items()):
+            if _is_tensor_list(value):
+                self.metadata[key] = value[-1]
+            elif isinstance(value, dict):
+                for sub_key, sub_value in list(value.items()):
+                    if _is_tensor_list(sub_value):
+                        value[sub_key] = sub_value[-1]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert back to a plain dict (tensors + metadata merged)."""
+        result: dict[str, Any] = dict(self.tensors)
+        result.update(self.metadata)
+        return result
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None:
+        """Create a MultimodalPayload from a raw dictionary.
+
+        Separates torch.Tensor values into tensors and everything
+        else into metadata.
+        """
+        if not data:
+            return None
+        tensors: dict[str, torch.Tensor] = {}
+        metadata: dict[str, Any] = {}
+        for k, v in data.items():
+            if isinstance(v, torch.Tensor):
+                tensors[k] = v
+            else:
+                metadata[k] = v
+        if not tensors and not metadata:
+            return None
+        return cls(tensors=tensors, metadata=metadata)
+
+    @classmethod
+    def from_raw(cls, payload: Any, modality_key: str) -> MultimodalPayload | None:
+        """Create a MultimodalPayload from a raw producer payload.
+
+        Accepts a MultimodalPayload (returned as-is), a dict, or a bare
+        tensor (stored under *modality_key*). Tensors are moved to CPU.
+        Producer-specific dict keys are remapped to the semantic modality
+        key (e.g. "audio", "latent"): AR runners produce {"hidden": ...}
+        and generation runners produce {"model_outputs": ...}.
+        """
+        if isinstance(payload, MultimodalPayload):
+            return payload
+
+        if not isinstance(payload, dict):
+            return cls.from_dict({modality_key: _to_cpu(payload)})
+
+        remapped: dict[str, Any] = {}
+        for key, value in payload.items():
+            is_producer_key = key == "model_outputs" or (key == "hidden" and modality_key != "hidden")
+            remapped[modality_key if is_producer_key else key] = _to_cpu(value)
+        return cls.from_dict(remapped)
+
+
+@dataclass
+class MultimodalCompletionOutput(CompletionOutput):
+    """CompletionOutput with multimodal support.
+
+    Inherits all CompletionOutput fields and adds multimodal_output.
+    As a CompletionOutput subclass, compatible with all existing vLLM consumers.
+    """
+
+    def __init__(
+        self,
+        multimodal_output: MultimodalPayload | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.multimodal_output = multimodal_output
+
+    def __repr__(self) -> str:
+        base = super().__repr__()
+        return f"{base[:-1]}, multimodal_output={self.multimodal_output!r})"
diff --git a/vllm_omni/outputs/output_metadata.py b/vllm_omni/outputs/output_metadata.py
new file mode 100644
index 000000000..f3e364a34
--- /dev/null
+++ b/vllm_omni/outputs/output_metadata.py
@@ -0,0 +1,238 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any, Literal, TypeAlias, TypedDict
+
+import numpy as np
+import PIL.Image
+import torch
+
+# Final user-facing output modality returned by pipeline stages.
+FinalOutputModalityType: TypeAlias = Literal["text", "image", "audio", "video"]
+
+# Scalar values that may appear in either payload or metadata.
+# Examples: text_output="done", fps=16, action_only_output=True.
+DiffusionScalar: TypeAlias = str | int | float | bool | None
+
+# Shape-like metadata. Examples: [1, 16, 3, 480, 832], (16, 480, 832).
+DiffusionShape: TypeAlias = list[int] | tuple[int, ...]
+
+# Named tensor controls used by transfer-style metadata.
+# Example: {"pose": pose_tensor, "depth": depth_tensor}.
+DiffusionTensorMap: TypeAlias = Mapping[str, torch.Tensor]
+
+# Values allowed under metadata groups such as video/audio/actions/transfer.
+# Examples: {"fps": 16}, {"shape": [1, 16, 3, 480, 832]},
+# {"controls": {"pose": pose_tensor}}, {"hints": ["pose"]}.
+DiffusionMetadataValue: TypeAlias = (
+    DiffusionScalar
+    | DiffusionShape
+    | DiffusionTensorMap
+    | torch.Tensor
+    | list["DiffusionMetadataValue"]
+    | tuple["DiffusionMetadataValue", ...]
+    | Mapping[str, "DiffusionMetadataValue"]
+)
+
+# Mutable metadata envelope side. Example:
+# {"video": {"fps": 16}, "actions": {"raw_action_dim": 14}}.
+DiffusionMetadata: TypeAlias = dict[str, DiffusionMetadataValue]
+
+# Read-only metadata view for validators and serving readers.
+DiffusionMetadataMapping: TypeAlias = Mapping[str, DiffusionMetadataValue]
+
+
+class DiffusionTrajectoryPayload(TypedDict, total=False):
+    # Optional denoising trajectory payload for RL/debug consumers.
+    latents: torch.Tensor | dict[str, Any] | None
+    timesteps: torch.Tensor | dict[str, Any] | None
+    log_probs: torch.Tensor | dict[str, Any] | None
+    decoded: list[PIL.Image.Image] | None
+
+
+# Concrete payload values produced by post_process.
+# Examples: PIL image, video tensor, audio bytes/tensor, action tensor.
+DiffusionPayloadValue: TypeAlias = (
+    DiffusionScalar
+    | bytes
+    | torch.Tensor
+    | np.ndarray
+    | PIL.Image.Image
+    | DiffusionTrajectoryPayload
+    | list["DiffusionPayloadValue"]
+    | tuple["DiffusionPayloadValue", ...]
+    | Mapping[str, "DiffusionPayloadValue"]
+)
+
+# Mutable payload side. Example: {"video": video_tensor, "action": action_tensor}.
+DiffusionPayload: TypeAlias = dict[str, DiffusionPayloadValue]
+
+# Read-only payload view for code that consumes but does not mutate payload.
+DiffusionPayloadMapping: TypeAlias = Mapping[str, DiffusionPayloadValue]
+
+# Serving-facing multimodal output. Non-metadata keys hold payload values.
+# Example: {"video": video_tensor, "metadata": {"video": {"fps": 16}}}.
+DiffusionMultimodalOutput: TypeAlias = dict[str, DiffusionPayloadValue | DiffusionMetadata]
+
+
+class DiffusionOutputEnvelope(TypedDict, total=False):
+    # Canonical post_process return shape:
+    # {"payload": {"video": video_tensor}, "metadata": {"video": {"fps": 16}}}
+    payload: DiffusionPayload
+    metadata: DiffusionMetadata
+
+
+# Raw post_process return before formatter normalization. This keeps backward
+# compatibility with direct payload returns while accepting the new envelope.
+# Examples: image, {"video": video_tensor}, or
+# {"payload": {"video": video_tensor}, "metadata": {"video": {"fps": 16}}}.
+DiffusionPostprocessRawOutput: TypeAlias = DiffusionPayloadValue | DiffusionPayload | DiffusionOutputEnvelope
+
+
+def strip_internal_metadata(metadata: DiffusionMetadataMapping) -> DiffusionMetadata:
+    """Return public metadata, dropping model-private postprocess state."""
+    return {key: value for key, value in metadata.items() if key != "internal"}
+
+
+def validate_public_diffusion_metadata(metadata: DiffusionMetadataMapping) -> None:
+    """Validate metadata that is safe to expose outside model postprocess."""
+    if "internal" in metadata:
+        raise ValueError("Diffusion metadata field 'internal' must not escape public formatting.")
+    validate_diffusion_metadata(metadata)
+
+
+def validate_diffusion_metadata(metadata: DiffusionMetadataMapping) -> None:
+    """Validate canonical diffusion output metadata fields.
+
+    Unknown model-specific groups and keys are intentionally tolerated. This
+    validator only guards fields that are part of the shared output protocol or
+    consumed by serving code.
+    """
+    _validate_video_metadata(_optional_mapping(metadata, "video"))
+    _validate_audio_metadata(_optional_mapping(metadata, "audio"))
+    _validate_actions_metadata(_optional_mapping(metadata, "actions"))
+    _validate_transfer_metadata(_optional_mapping(metadata, "transfer"))
+    _validate_text_metadata(_optional_mapping(metadata, "text"))
+    _validate_common_metadata(_optional_mapping(metadata, "common"))
+
+
+def _validate_video_metadata(metadata: DiffusionMetadataMapping | None) -> None:
+    if metadata is None:
+        return
+    _check_optional_positive_number(metadata, "fps", group="video")
+    _check_optional_int(metadata, "video_fps_multiplier", group="video", min_value=1)
+    _check_optional_shape(metadata, "shape", group="video")
+
+
+def _validate_audio_metadata(metadata: DiffusionMetadataMapping | None) -> None:
+    if metadata is None:
+        return
+    _check_optional_int(metadata, "sample_rate", group="audio", min_value=1)
+    _check_optional_non_empty_str(metadata, "audiox_task", group="audio")
+
+
+def _validate_actions_metadata(metadata: DiffusionMetadataMapping | None) -> None:
+    if metadata is None:
+        return
+    _check_optional_int(metadata, "raw_action_dim", group="actions", min_value=1)
+    _check_optional_non_empty_str(metadata, "action_mode", group="actions")
+    _check_optional_int(metadata, "domain_id", group="actions")
+
+
+def _validate_transfer_metadata(metadata: DiffusionMetadataMapping | None) -> None:
+    if metadata is None:
+        return
+    controls = metadata.get("controls")
+    if controls is not None and not _is_tensor_mapping(controls):
+        raise TypeError("Diffusion metadata field 'transfer.controls' must be a mapping of str to torch.Tensor.")
+    hints = metadata.get("hints")
+    if hints is not None and not _is_str_list(hints):
+        raise TypeError("Diffusion metadata field 'transfer.hints' must be a list of strings.")
+
+
+def _validate_text_metadata(metadata: DiffusionMetadataMapping | None) -> None:
+    if metadata is None:
+        return
+    _check_optional_str(metadata, "text_output", group="text")
+    _check_optional_str(metadata, "think_text", group="text")
+
+
+def _validate_common_metadata(metadata: DiffusionMetadataMapping | None) -> None:
+    if metadata is None:
+        return
+    value = metadata.get("action_only_output")
+    if value is not None and not isinstance(value, bool):
+        raise TypeError("Diffusion metadata field 'common.action_only_output' must be bool.")
+
+
+def _optional_mapping(metadata: DiffusionMetadataMapping, group: str) -> DiffusionMetadataMapping | None:
+    value = metadata.get(group)
+    if value is None:
+        return None
+    if not isinstance(value, Mapping):
+        raise TypeError(f"Diffusion metadata group '{group}' must be a mapping.")
+    return value
+
+
+def _check_optional_positive_number(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None:
+    value = metadata.get(key)
+    if value is None:
+        return
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise TypeError(f"Diffusion metadata field '{group}.{key}' must be a positive number.")
+    if value <= 0:
+        raise ValueError(f"Diffusion metadata field '{group}.{key}' must be > 0.")
+
+
+def _check_optional_int(
+    metadata: DiffusionMetadataMapping,
+    key: str,
+    *,
+    group: str,
+    min_value: int | None = None,
+) -> None:
+    value = metadata.get(key)
+    if value is None:
+        return
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise TypeError(f"Diffusion metadata field '{group}.{key}' must be int.")
+    if min_value is not None and value < min_value:
+        raise ValueError(f"Diffusion metadata field '{group}.{key}' must be >= {min_value}.")
+
+
+def _check_optional_str(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None:
+    value = metadata.get(key)
+    if value is not None and not isinstance(value, str):
+        raise TypeError(f"Diffusion metadata field '{group}.{key}' must be str.")
+
+
+def _check_optional_non_empty_str(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None:
+    value = metadata.get(key)
+    if value is None:
+        return
+    if not isinstance(value, str):
+        raise TypeError(f"Diffusion metadata field '{group}.{key}' must be str.")
+    if not value:
+        raise ValueError(f"Diffusion metadata field '{group}.{key}' must be non-empty.")
+
+
+def _check_optional_shape(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None:
+    value = metadata.get(key)
+    if value is None:
+        return
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)):
+        raise TypeError(f"Diffusion metadata field '{group}.{key}' must be a list or tuple of positive ints.")
+    if any(isinstance(dim, bool) or not isinstance(dim, int) or dim <= 0 for dim in value):
+        raise ValueError(f"Diffusion metadata field '{group}.{key}' must contain positive ints.")
+
+
+def _is_tensor_mapping(value: object) -> bool:
+    return isinstance(value, Mapping) and all(
+        isinstance(key, str) and isinstance(tensor, torch.Tensor) for key, tensor in value.items()
+    )
+
+
+def _is_str_list(value: object) -> bool:
+    return isinstance(value, list) and all(isinstance(item, str) for item in value)
diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py
new file mode 100644
index 000000000..2536a8ade
--- /dev/null
+++ b/vllm_omni/outputs/output_modality.py
@@ -0,0 +1,120 @@
+"""Output modality types for vLLM-Omni.
+
+This module defines the OutputModality enum and TensorAccumulationStrategy
+for type-safe multimodal output routing and tensor merging.
+
+"""
+
+from __future__ import annotations
+
+import re
+from enum import Enum, Flag, StrEnum, auto
+
+_MODALITY_ALIASES: dict[str, str] = {
+    "speech": "audio",
+    "images": "image",
+    "latents": "latent",
+    "wav": "audio",
+    "waveform": "audio",
+    "pixel_values": "image",
+    "pixels": "image",
+    "token_ids": "text",
+    "tokens": "text",
+}
+
+
+class OutputModalityNames(StrEnum):
+    """Keys for output modalities.
+
+    TODO: (Alex) Integrate this with the big-flag enum below + throughout the code
+    for better type safety (currently only used for output processor).
+    """
+
+    TEXT = "text"
+    IMAGE = "image"
+    AUDIO = "audio"
+    LATENT = "latent"
+
+
+# Specify which output modalities may be drained when handling delta messages.
+# For some types, e.g., latents, we need to be careful to ensure the full context
+# is passed as the stream yields due to assumptions in the I/O processing and model
+# when async chunk isn't enabled.
+NON_DRAINABLE_MODALITIES = {OutputModalityNames.TEXT, OutputModalityNames.LATENT}
+DRAINABLE_MODALITIES = {mod for mod in OutputModalityNames if mod not in NON_DRAINABLE_MODALITIES}
+
+
+class OutputModality(Flag):
+    """Bit-flag enum for output modalities.
+
+    Compose freely with ``|`` — no need to enumerate every combination.
+
+    Single:   ``OutputModality.TEXT``, ``OutputModality.IMAGE``, ...
+    Compound: ``OutputModality.TEXT | OutputModality.IMAGE``  (text+image)
+
+    Note: POOLING is intentionally excluded. Pooling/embedding is vLLM's
+    native path (pooling_output → PoolingRequestOutput), handled entirely
+    by the base OutputProcessor. vLLM-Omni's layer does not participate.
+    """
+
+    TEXT = auto()
+    IMAGE = auto()
+    AUDIO = auto()
+    LATENT = auto()
+
+    @classmethod
+    def from_string(cls, s: str | None) -> OutputModality:
+        """Parse a free-text modality string into an OutputModality flag.
+
+        Handles common aliases and compound strings separated by + or ,.
+
+        Examples::
+
+            OutputModality.from_string("text+image")
+            # → OutputModality.TEXT | OutputModality.IMAGE
+        """
+        if not s or not s.strip():
+            return cls.TEXT
+
+        parts = [p.strip().lower() for p in re.split(r"[+,]", s.strip())]
+        result = cls(0)
+        for p in parts:
+            p = _MODALITY_ALIASES.get(p, p)
+            try:
+                result |= cls[p.upper()]
+            except KeyError:
+                raise ValueError(f"Unknown modality: {p!r}. Supported: {[m.name.lower() for m in cls]}")
+        return result
+
+    @property
+    def has_text(self) -> bool:
+        return OutputModality.TEXT in self
+
+    @property
+    def has_multimodal(self) -> bool:
+        return bool(self & ~OutputModality.TEXT)
+
+
+class TensorAccumulationStrategy(Enum):
+    """Strategy for merging incremental multimodal tensors."""
+
+    CONCAT_DIM0 = "concat_dim0"
+    """Concatenate along dimension 0. Used for image/latent tensors."""
+
+    CONCAT_LAST = "concat_last"
+    """Concatenate along the last dimension. Used for audio waveforms."""
+
+    APPEND_LIST = "append_list"
+    """Append to a list (no tensor concatenation)."""
+
+    REPLACE = "replace"
+    """Replace previous tensor entirely with the latest one."""
+
+
+def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy:
+    """Determine tensor merge strategy from the multimodal flags."""
+    if OutputModality.AUDIO in modality:
+        return TensorAccumulationStrategy.CONCAT_LAST
+    if OutputModality.IMAGE in modality or OutputModality.LATENT in modality:
+        return TensorAccumulationStrategy.CONCAT_DIM0
+    return TensorAccumulationStrategy.CONCAT_DIM0  # default
diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py
new file mode 100644
index 000000000..e39b995b0
--- /dev/null
+++ b/vllm_omni/outputs/output_processor.py
@@ -0,0 +1,633 @@
+from dataclasses import fields as dataclass_fields
+from typing import Any
+
+import torch
+from vllm.logger import init_logger
+from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput
+from vllm.sampling_params import RequestOutputKind
+from vllm.tokenizers import TokenizerLike
+from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
+from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
+from vllm.v1.engine.output_processor import (
+    OutputProcessorOutput,
+    RequestOutputCollector,
+    RequestState,
+)
+from vllm.v1.engine.parallel_sampling import ParentRequest
+from vllm.v1.metrics.stats import IterationStats, RequestStateStats
+
+from vllm_omni.data_entry_keys import unflatten_payload
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput, MultimodalPayload
+from vllm_omni.outputs.output_modality import (
+    DRAINABLE_MODALITIES,
+    OutputModality,
+    get_accumulation_strategy,
+)
+
+logger = init_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Module-level helpers
+# ---------------------------------------------------------------------------
+
+
+def _modality_to_type_string(modality: OutputModality) -> str:
+    """Convert an OutputModality flag to a lowercase type string."""
+    try:
+        if OutputModality.AUDIO in modality:
+            return "audio"
+        if OutputModality.IMAGE in modality:
+            return "image"
+        if OutputModality.LATENT in modality:
+            return "latent"
+    except TypeError:
+        # Flag identity mismatch (e.g. after module reload in tests).
+        name = getattr(modality, "name", "") or ""
+        lowered = name.lower()
+        if "audio" in lowered:
+            return "audio"
+        if "image" in lowered:
+            return "image"
+        if "latent" in lowered:
+            return "latent"
+    return "text"
+
+
+class OmniRequestState(RequestState):
+    """Request state for omni models, tracking multimodal outputs.
+
+    Extends the base RequestState with support for accumulating
+    multimodal tensor outputs (e.g., images, audio, latents) that
+    are produced incrementally during generation.
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+        # arrival_time is always passed as a keyword argument by the sole caller
+        # (from_new_request).  The *args fallback via positional index was removed
+        # because upstream RequestState.__init__ reorders parameters between
+        # releases, making args[12] fragile.
+        arrival_time = kwargs.get("arrival_time")
+        super().__init__(*args, **kwargs)
+        self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0))
+        # Omni-specific: multimodal output accumulation
+        # TODO: mm_type is per-request, not per-key. If a model ever produces
+        # both audio and latent outputs, the modality type would flip on each
+        # add_multimodal_tensor() call. Consider tracking per-key modality
+        # types (e.g. dict[str, str]) for future multi-output models.
+        self.mm_type: str | None = None
+        self.mm_accumulated: MultimodalPayload = MultimodalPayload()
+
+    def apply_streaming_update(self, update) -> None:
+        super().apply_streaming_update(update)
+        self.native_text_stats.arrival_time = update.arrival_time
+
+    def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None:
+        """Accumulate a multimodal tensor payload into the request state.
+
+        Normalizes incoming payloads (dict or raw tensor) into a
+        MultimodalPayload and merges with any previously accumulated data.
+        Uses list-based deferred concatenation to avoid O(n²) repeated
+        torch.cat calls.
+        """
+        if payload is None:
+            return
+        try:
+            if mm_type:
+                self.mm_type = mm_type.lower()
+            modality_key = self.mm_type or "hidden"
+
+            incoming = MultimodalPayload.from_raw(payload, modality_key)
+            if incoming is not None:
+                self.mm_accumulated = self.mm_accumulated.merged_with(incoming)
+        except (ValueError, TypeError, RuntimeError):
+            logger.exception("Error accumulating multimodal tensor")
+
+    def _consolidate_multimodal_tensors(self) -> None:
+        """Consolidate accumulated tensor lists into single tensors.
+
+        Uses TensorAccumulationStrategy derived from the output modality
+        to determine concatenation behavior. Metadata values always use
+        REPLACE (keep latest).
+        """
+        if self.mm_accumulated.is_empty:
+            return
+
+        try:
+            modality = OutputModality.from_string(self.mm_type)
+        except (ValueError, KeyError):
+            modality = OutputModality.TEXT
+        strategy = get_accumulation_strategy(modality)
+
+        try:
+            self.mm_accumulated.consolidate_tensors(strategy)
+            self.mm_accumulated.consolidate_metadata()
+        except (RuntimeError, TypeError, KeyError):
+            logger.exception("Error consolidating multimodal tensors")
+
+    # Override: do not route to pooling-only path; always create completion
+    # outputs, and attach pooling_result into the CompletionOutput.
+    def make_request_output(
+        self,
+        new_token_ids: list[int],
+        pooling_output: torch.Tensor | None,
+        finish_reason: FinishReason | None,
+        stop_reason: int | str | None,
+        kv_transfer_params: dict[str, Any] | None = None,
+        routed_experts: Any = None,
+    ) -> OmniRequestOutput | PoolingRequestOutput | None:
+        """Create a request output from generation results.
+
+        Creates a RequestOutput or PoolingRequestOutput from the generated
+        tokens and accumulated multimodal outputs. Attaches multimodal
+        tensors to the completion output if available.
+
+        Args:
+            new_token_ids: List of newly generated token IDs
+            pooling_output: Optional pooling output tensor
+            finish_reason: Optional finish reason indicating why generation stopped
+            stop_reason: Optional stop reason (token ID or stop string)
+            kv_transfer_params: Optional KV cache transfer parameters
+            routed_experts: Optional MoE routed-expert ids for this step,
+                attached to the completion output for generation stages
+
+        Returns:
+            OmniRequestOutput or PoolingRequestOutput if output should be
+            emitted (based on finish status and output kind), None otherwise
+        """
+        # Pooling-only requests should follow base behavior.
+        if self.detokenizer is None and pooling_output is not None:
+            return super().make_request_output(
+                new_token_ids,
+                pooling_output,
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+            )
+
+        finished = finish_reason is not None
+        final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
+        is_delta = self.output_kind == RequestOutputKind.DELTA
+
+        if not finished and final_only:
+            return None
+
+        # Consolidate accumulated tensors when finishing, or for every
+        # CUMULATIVE step so the consumer always sees a single tensor.
+        if finished or not is_delta:
+            self._consolidate_multimodal_tensors()
+
+        if self.stream_interval > 1 and self.detokenizer is not None:
+            # Send output request only when
+            # 1. It has finished, or
+            # 2. It is the first token, or
+            # 3. It has reached the stream interval number of tokens
+            if not (
+                finished
+                or self.sent_tokens_offset == 0
+                or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval
+            ):
+                return None
+
+            if self.output_kind == RequestOutputKind.DELTA:
+                # Send tokens from the offset in DELTA mode, otherwise all
+                # tokens are sent.
+                new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :]
+                self.sent_tokens_offset = self.detokenizer.num_output_tokens()
+
+        external_req_id = self.external_req_id
+
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts)
+
+        if self.parent_req is None:
+            outputs = [output]
+        else:
+            outputs, finished = self.parent_req.get_outputs(self.request_id, output)
+            if not outputs:
+                return None
+            external_req_id = self.parent_req.external_req_id
+
+        return self._new_request_output(
+            external_req_id,
+            outputs,
+            finished,
+            kv_transfer_params,
+        )
+
+    def _new_completion_output(
+        self,
+        token_ids: list[int],
+        finish_reason: FinishReason | None,
+        stop_reason: int | str | None,
+        routed_experts: Any = None,
+    ) -> MultimodalCompletionOutput | CompletionOutput:
+        """Create a completion output with multimodal data attached.
+
+        Returns a MultimodalCompletionOutput when multimodal data has been
+        accumulated, otherwise returns the base CompletionOutput.  Snapshots
+        the accumulated payload before draining DELTA-mode modality keys so
+        that callers receive an immutable view.
+        """
+        # When there is no detokenizer (generation stages), build a minimal
+        # CompletionOutput directly since upstream asserts detokenizer != None.
+        if self.detokenizer is None:
+            finished = finish_reason is not None
+            base_output = CompletionOutput(
+                index=self.request_index,
+                text="",
+                token_ids=token_ids,
+                logprobs=None,
+                cumulative_logprob=None,
+                finish_reason=str(finish_reason) if finished else None,
+                stop_reason=stop_reason if finished else None,
+                routed_experts=routed_experts,
+            )
+        else:
+            base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
+
+        # Always provide cumulative token IDs for inter-stage processors.
+        if self.detokenizer is not None:
+            base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids)
+        else:
+            base_output.cumulative_token_ids = list(token_ids)
+
+        # Attach cumulative_text only at the final step for inter-stage use.
+        if finish_reason is not None and hasattr(self.detokenizer, "output_text"):
+            base_output.cumulative_text = self.detokenizer.output_text
+
+        try:
+            if self.mm_accumulated and not self.mm_accumulated.is_empty:
+                # Snapshot: copy current tensors/metadata so drain doesn't
+                # mutate the payload already handed to the caller.
+                # Unflatten dotted keys (e.g. "hidden_states.layer_0") back
+                # to nested dicts so downstream consumers (thinker2talker etc.)
+                # can access with .get("hidden_states", {}).get("layers", {}).
+                merged = {**self.mm_accumulated.tensors, **self.mm_accumulated.metadata}
+                nested = unflatten_payload(merged)
+                snapshot = MultimodalPayload.from_dict(nested) or MultimodalPayload()
+                kwargs = {f.name: getattr(base_output, f.name) for f in dataclass_fields(CompletionOutput)}
+                output = MultimodalCompletionOutput(
+                    multimodal_output=snapshot,
+                    **kwargs,
+                )
+                output.cumulative_token_ids = base_output.cumulative_token_ids
+                if hasattr(base_output, "cumulative_text"):
+                    output.cumulative_text = base_output.cumulative_text
+
+                # DELTA mode: drain modality keys (e.g. audio) so the next
+                # step only sees freshly accumulated data for those keys.
+                if self.output_kind == RequestOutputKind.DELTA:
+                    for modality_key in DRAINABLE_MODALITIES:
+                        self.mm_accumulated.tensors.pop(modality_key, None)
+
+                return output
+        except (RuntimeError, TypeError, AttributeError):
+            logger.exception("Error creating MultimodalCompletionOutput")
+        return base_output
+
+    def _new_request_output(
+        self,
+        external_req_id: str,
+        outputs: list,
+        finished: bool,
+        kv_transfer_params: dict[str, Any] | None = None,
+    ) -> RequestOutput | PoolingRequestOutput:
+        """Create request output, handling no-detokenizer generation stages.
+
+        Upstream asserts ``self.logprobs_processor is not None`` which fails
+        for generation stages (talker, code2wav) that have no tokenizer.
+        When the logprobs processor is absent we build the RequestOutput
+        directly with ``prompt_logprobs=None``.
+        """
+        if self.logprobs_processor is not None:
+            return super()._new_request_output(
+                external_req_id,
+                outputs,
+                finished,
+                kv_transfer_params,
+            )
+
+        # No-detokenizer path: build RequestOutput directly.
+        prompt_token_ids = self.prompt_token_ids
+        if prompt_token_ids is None and self.prompt_embeds is not None:
+            prompt_token_ids = [0] * len(self.prompt_embeds)
+        if prompt_token_ids is None:
+            prompt_token_ids = []
+
+        return RequestOutput(
+            request_id=external_req_id,
+            lora_request=self.lora_request,
+            prompt=self.prompt,
+            prompt_token_ids=prompt_token_ids,
+            prompt_logprobs=None,
+            outputs=outputs,
+            finished=finished,
+            kv_transfer_params=kv_transfer_params,
+            num_cached_tokens=self.num_cached_tokens,
+            metrics=self.stats,
+        )
+
+
+class MultimodalOutputProcessor(VLLMOutputProcessor):
+    """Handles multimodal output processing.
+
+    Captures multimodal outputs from OmniEngineCoreOutput and accumulates
+    them as MultimodalPayload in OmniRequestState, before delegating to
+    the base vLLM OutputProcessor for text handling.
+
+    The data flow is:
+    1. For each EngineCoreOutput with multimodal_output:
+       - Capture into OmniRequestState.add_multimodal_tensor()
+    2. Base vLLM OutputProcessor handles text detokenization
+    3. On finish, _consolidate_multimodal_tensors() concatenates accumulated
+       tensors using strategy-based dispatch
+    4. _new_completion_output() returns MultimodalCompletionOutput
+    """
+
+    def __init__(
+        self,
+        tokenizer: TokenizerLike | None,
+        *,
+        log_stats: bool,
+        stream_interval: int = 1,
+        tracing_enabled: bool = False,
+        engine_core_output_type: str | None = None,
+        output_modality: OutputModality = OutputModality.TEXT,
+    ):
+        """Initialize the multimodal output processor.
+
+        Args:
+            tokenizer: Tokenizer for detokenizing text outputs
+            log_stats: Whether to log statistics
+            stream_interval: Stream interval for output generation
+            engine_core_output_type: Optional output type string (e.g.,
+                "image", "audio", "latent"). Converted to OutputModality
+                internally. Kept for backward compatibility with
+                stage_init_utils.
+            output_modality: Type-safe output modality flag. Used to tag
+                multimodal outputs with the correct modality key when
+                per-output type info is unavailable.
+        """
+        super().__init__(
+            tokenizer=tokenizer,
+            log_stats=log_stats,
+            stream_interval=stream_interval,
+            tracing_enabled=tracing_enabled,
+        )
+        # Convert string-based engine_core_output_type to OutputModality
+        if engine_core_output_type is not None:
+            self.output_modality = OutputModality.from_string(engine_core_output_type)
+        else:
+            self.output_modality = output_modality
+        self.engine_core_output_type = engine_core_output_type
+        self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {}
+
+    def _native_text_metric_record(self, request_id: str) -> dict[str, Any]:
+        return self._native_text_metrics_by_request.setdefault(
+            request_id,
+            {
+                "vllm_ttft_ms": 0.0,
+                "vllm_tpot_ms": 0.0,
+                "vllm_itl_ms": 0.0,
+                "vllm_itls_ms": [],
+            },
+        )
+
+    def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]:
+        return self._native_text_metrics_by_request.pop(request_id, {})
+
+    def abort_requests(self, request_ids, internal: bool) -> list[str]:
+        request_ids = list(request_ids)
+        for request_id in request_ids:
+            if internal:
+                req_state = self.request_states.get(request_id)
+                if req_state is not None:
+                    self._native_text_metrics_by_request.pop(req_state.external_req_id, None)
+            else:
+                self._native_text_metrics_by_request.pop(request_id, None)
+        return super().abort_requests(request_ids, internal)
+
+    def add_request(
+        self,
+        request: EngineCoreRequest,
+        prompt: str | None,
+        parent_req: ParentRequest | None = None,
+        request_index: int = 0,
+        queue: RequestOutputCollector | None = None,
+    ) -> None:
+        """Add a new request to be processed.
+
+        Creates an OmniRequestState for the request and registers it
+        for output processing.
+
+        Args:
+            request: Engine core request to add
+            prompt: Optional prompt string for the request
+            parent_req: Optional parent request for parallel sampling
+            request_index: Index of the request in the batch
+            queue: Optional queue for collecting outputs
+
+        Raises:
+            ValueError: If the request ID is already registered
+        """
+        request_id = request.request_id
+        req_state = self.request_states.get(request_id)
+        if req_state is not None:
+            self._update_streaming_request_state(req_state, request, prompt)
+            return
+
+        req_state = OmniRequestState.from_new_request(
+            tokenizer=self.tokenizer,
+            request=request,
+            prompt=prompt,
+            parent_req=parent_req,
+            request_index=request_index,
+            queue=queue,
+            log_stats=self.log_stats,
+            stream_interval=self.stream_interval,
+        )
+        self.request_states[request_id] = req_state
+        if parent_req:
+            self.parent_requests[parent_req.request_id] = parent_req
+        self.external_req_ids[req_state.external_req_id].append(request_id)
+
+    def remove_request(self, request_id: str) -> None:
+        """Rollback one previously registered request if it was never submitted."""
+        req_state = self.request_states.pop(request_id, None)
+        if req_state is None:
+            return
+
+        external_req_id = getattr(req_state, "external_req_id", None)
+        if external_req_id is not None:
+            self._native_text_metrics_by_request.pop(external_req_id, None)
+            request_ids = self.external_req_ids.get(external_req_id)
+            if request_ids is not None:
+                self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id]
+                if not self.external_req_ids[external_req_id]:
+                    self.external_req_ids.pop(external_req_id, None)
+
+        parent_req = getattr(req_state, "parent_req", None)
+        if parent_req is not None:
+            self.parent_requests.pop(parent_req.request_id, None)
+
+    def process_outputs(
+        self,
+        engine_core_outputs: list[EngineCoreOutput],
+        engine_core_timestamp: float | None = None,
+        iteration_stats: IterationStats | None = None,
+    ) -> OutputProcessorOutput:
+        default_mm_type = _modality_to_type_string(self.output_modality)
+
+        # Separate outputs that upstream can handle (has detokenizer or
+        # pooling) from multimodal-only outputs (no detokenizer, no pooling)
+        # that would trigger upstream's `assert detokenizer is not None`.
+        upstream_outputs: list[EngineCoreOutput] = []
+        mm_only_outputs: list[EngineCoreOutput] = []
+
+        for eco in engine_core_outputs:
+            req_state = self.request_states.get(eco.request_id)
+            if req_state is None:
+                continue
+
+            # Accumulate multimodal tensors regardless of path.
+            if isinstance(req_state, OmniRequestState):
+                mm_output = getattr(eco, "multimodal_output", None)
+                if mm_output is not None:
+                    mm_type = getattr(eco, "output_type", None) or default_mm_type
+                    req_state.add_multimodal_tensor(mm_output, mm_type)
+
+            # Route: if no detokenizer and no pooling output, handle locally
+            # to avoid upstream's assert on detokenizer.
+            if req_state.detokenizer is None and eco.pooling_output is None:
+                mm_only_outputs.append(eco)
+            else:
+                upstream_outputs.append(eco)
+
+        # Handle multimodal-only outputs (generation stages) locally.
+        self._process_mm_only_outputs(mm_only_outputs)
+
+        # Delegate text/pooling outputs to upstream.
+        return super().process_outputs(
+            upstream_outputs,
+            engine_core_timestamp=engine_core_timestamp,
+            iteration_stats=iteration_stats,
+        )
+
+    def _process_mm_only_outputs(
+        self,
+        engine_core_outputs: list[EngineCoreOutput],
+    ) -> None:
+        """Handle outputs from generation stages that have no detokenizer.
+
+        These cannot go through upstream process_outputs because it asserts
+        detokenizer is not None when pooling_output is None.
+        """
+        for eco in engine_core_outputs:
+            req_state = self.request_states.get(eco.request_id)
+            if req_state is None or not isinstance(req_state, OmniRequestState):
+                continue
+
+            new_token_ids = eco.new_token_ids
+            finish_reason = eco.finish_reason
+            stop_reason = eco.stop_reason
+            kv_transfer_params = eco.kv_transfer_params
+            routed_experts = eco.routed_experts
+            req_state.num_cached_tokens = eco.num_cached_tokens
+            req_state.is_prefilling = False
+
+            if request_output := req_state.make_request_output(
+                new_token_ids,
+                None,  # pooling_output
+                finish_reason,
+                stop_reason,
+                kv_transfer_params,
+                routed_experts,
+            ):
+                if req_state.queue is not None:
+                    req_state.queue.put(request_output)
+
+            if finish_reason is not None:
+                self._finish_request(req_state)
+
+    def _update_stats_from_output(
+        self,
+        req_state: RequestState,
+        engine_core_output: EngineCoreOutput,
+        engine_core_timestamp: float | None,
+        iteration_stats: IterationStats | None,
+    ):
+        was_prefilling = req_state.is_prefilling
+        native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
+        previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0
+
+        # NOTE: We pass ``None`` for  *iteration_stats* to the parent so that
+        # the upstream's ``_update_stats_from_output`` logs stats via its own
+        # ``RequestStateStats`` (req_state.stats) without touching the omni
+        # ``IterationStats`` object.  Below we re-apply the same update using
+        # *native_stats* (the omni-specific ``RequestStateStats`` attached to
+        # ``OmniRequestState``), which correctly feeds the omni metrics path.
+        # If the upstream adds new logic in ``_update_stats_from_output`` (e.g.
+        # additional stat fields), this override must be kept in sync.
+        super()._update_stats_from_output(
+            req_state,
+            engine_core_output,
+            engine_core_timestamp,
+            None,
+        )
+
+        if iteration_stats is None or engine_core_timestamp is None or native_stats is None:
+            return
+
+        iteration_stats.update_from_output(
+            engine_core_output,
+            engine_core_timestamp,
+            was_prefilling,
+            native_stats,
+            self.lora_states,
+            req_state.lora_name,
+        )
+        record = self._native_text_metric_record(req_state.external_req_id)
+        if was_prefilling:
+            record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0)
+            return
+
+        if previous_last_token_ts > 0:
+            itl_ms = max((engine_core_timestamp - previous_last_token_ts) * 1000.0, 0.0)
+            itls_ms = record.setdefault("vllm_itls_ms", [])
+            itls_ms.append(itl_ms)
+            record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms))
+
+    def _update_stats_from_finished(
+        self,
+        req_state: RequestState,
+        finish_reason: FinishReason | None,
+        iteration_stats: IterationStats | None,
+    ):
+        super()._update_stats_from_finished(req_state, finish_reason, None)
+
+        native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
+        if iteration_stats is None or native_stats is None or finish_reason is None:
+            return
+
+        iteration_stats.update_from_finished_request(
+            finish_reason=finish_reason,
+            request_id=req_state.external_req_id,
+            num_prompt_tokens=req_state.prompt_len,
+            max_tokens_param=req_state.max_tokens_param,
+            req_stats=native_stats,
+            num_cached_tokens=req_state.num_cached_tokens,
+        )
+        if not iteration_stats.finished_requests:
+            return
+
+        finished_request = iteration_stats.finished_requests[-1]
+        if finished_request.request_id != req_state.external_req_id:
+            return
+        self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = (
+            float(finished_request.mean_time_per_output_token) * 1000.0
+        )
diff --git a/vllm_omni/outputs/utils.py b/vllm_omni/outputs/utils.py
new file mode 100644
index 000000000..88fc11fb5
--- /dev/null
+++ b/vllm_omni/outputs/utils.py
@@ -0,0 +1,24 @@
+"""Shared helpers for multimodal output handling."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def _to_cpu(x: Any) -> Any:
+    """Move a tensor to CPU; pass through non-tensors unchanged."""
+    if isinstance(x, torch.Tensor):
+        try:
+            return x.detach().to("cpu", non_blocking=True).contiguous()
+        except (RuntimeError, AttributeError):
+            return x
+    if isinstance(x, dict):
+        return {str(sub_key): _to_cpu(sub_value) for sub_key, sub_value in x.items()}
+    return x
+
+
+def _is_tensor_list(value: Any) -> bool:
+    """Return True if *value* is a non-empty list of tensors (deferred concat)."""
+    return isinstance(value, list) and bool(value) and isinstance(value[0], torch.Tensor)
diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py
index 50eb4d991..4f2441294 100644
--- a/vllm_omni/patch.py
+++ b/vllm_omni/patch.py
@@ -417,3 +417,70 @@ def _patch_fp8_use_quack_fused_bias():
 
 
 _patch_fp8_use_quack_fused_bias()
+
+
+# =============================================================================
+# Patch CuMemAllocator._python_free_callback to fix CUDA double-free on shutdown
+# =============================================================================
+# WHY: CuMemAllocator._python_free_callback guards the asleep-entry double-free
+# skip with ``if data.is_asleep and current_platform.is_rocm():`` — only ROCm
+# gets the safe empty-handle return.  On CUDA, the callback falls through and
+# returns the original handle, causing cuMemRelease on already-freed memory
+# (CUDA_ERROR_INVALID_VALUE) during EngineCore subprocess atexit cleanup.
+#
+# This happens because ``sleep()`` calls ``unmap_and_release()`` on ALL
+# platforms, then sets ``is_asleep = True``.  When the atexit handler
+# (``_shutdown_singleton`` -> ``release_pools()`` -> GC) triggers the free
+# callback, the asleep entries must return an empty chunk list so the C
+# extension skips ``cuMemRelease`` — exactly the same logic already used by
+# the ROCm guard.
+#
+# The fix removes the ``current_platform.is_rocm()`` condition so the guard
+# applies to CUDA (and any future platform that implements cumem).
+#
+# FRAGILITY: Relies on ``_python_free_callback`` being a regular method
+# (not a slot or C extension).  If CuMemAllocator is rewritten in C/Cython,
+# this monkey-patch will silently become a no-op.
+def _patch_cumem_free_callback_cuda() -> None:
+    try:
+        from vllm.device_allocator.cumem import CuMemAllocator
+    except ImportError:
+        _PATCH_LOGGER.debug("[cumem-cuda] CuMemAllocator not available; skipping patch")
+        return
+
+    _original_free_callback = CuMemAllocator._python_free_callback
+
+    if getattr(_original_free_callback, "_omni_cumem_cuda_patched", False):
+        return
+
+    # The upstream bug: `_python_free_callback` only skips the double-free
+    # for ROCm (line ~206).  We wrap the method to extend the guard to all
+    # platforms.
+    def _patched_free_callback(self, ptr: int) -> tuple:
+        data = self.pointer_to_data.pop(ptr)
+        if data.cpu_backup_tensor is not None:
+            data.cpu_backup_tensor = None
+        if data.is_asleep:
+            # sleep() already called unmap_and_release() on this allocation.
+            # Return an empty chunk list so the C extension skips
+            # cuMemRelease, avoiding a double-free.  Same logic as the
+            # existing ROCm guard, but applied to all platforms.
+            device, size, d_mem, _ = data.handle
+            result = (device, size, d_mem, [])
+            _PATCH_LOGGER.debug(
+                "[cumem-cuda] Free callback: asleep entry %s -> empty handle",
+                ptr,
+            )
+            return result
+        # Drain pending kernels before the C extension's cuMemUnmap.
+        torch.accelerator.synchronize(data.handle[0])
+        return data.handle
+
+    _patched_free_callback._omni_cumem_cuda_patched = True
+    CuMemAllocator._python_free_callback = _patched_free_callback
+    _PATCH_LOGGER.info(
+        "[cumem-cuda] CuMemAllocator._python_free_callback patched: asleep guard extended to all platforms."
+    )
+
+
+_patch_cumem_free_callback_cuda()
diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py
index 2127509eb..a05bb289a 100644
--- a/vllm_omni/platforms/cuda/platform.py
+++ b/vllm_omni/platforms/cuda/platform.py
@@ -121,6 +121,31 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
+                try:
+                    importlib.import_module("kernels")
+                    logger.info("Using HuggingFace kernels-backed attention backend '%s'", backend_upper)
+                except ImportError:
+                    if backend_upper == "FLASH_ATTN_HUB":
+                        logger.warning(
+                            "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN."
+                        )
+                        backend_upper = "FLASH_ATTN"
+                    elif backend_upper == "FLASH_ATTN_3_HUB":
+                        logger.warning(
+                            "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN."
+                        )
+                        backend_upper = "FLASH_ATTN"
+
+            if backend_upper == "FLASH_ATTN_3_HUB":
+                fa3_hub_supported = compute_capability is not None and compute_capability.major >= 9
+                if not fa3_hub_supported:
+                    logger.warning(
+                        "FLASH_ATTN_3_HUB requires a Hopper-class GPU with compute capability >= 9.0. "
+                        "Falling back to FLASH_ATTN_HUB."
+                    )
+                    backend_upper = "FLASH_ATTN_HUB"
+
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
                 if not compute_supported:
                     logger.warning(
@@ -224,16 +249,8 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
     @classmethod
     def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
-        """Copied from upstream CudaPlatformBase with inductor-aware logic.
-
-        When inductor is active (compiling) use native as the default;
-        otherwise prefer vllm_c kernels where available.
-        """
-        from vllm.config.compilation import CompilationMode
-
-        cc = vllm_config.compilation_config
-        using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
-        default = ["native"] if using_inductor else ["vllm_c", "native"]
+        """Prefer ``vllm_c`` CUDA kernels over ``native`` for diffusion IR ops."""
+        default = ["vllm_c", "native"]
 
         # Use oink if enabled for rms_norm
         # TODO(Laurawly/luka): remove this env var,
diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py
index 96e27edf3..92a194d59 100644
--- a/vllm_omni/platforms/musa/platform.py
+++ b/vllm_omni/platforms/musa/platform.py
@@ -78,6 +78,14 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
+                logger.warning(
+                    "HuggingFace kernels-backed FlashAttention is "
+                    "not supported on MUSA. Falling back to local "
+                    "FLASH_ATTN."
+                )
+                backend_upper = "FLASH_ATTN"
+
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
                 if not compute_supported:
                     logger.warning(
diff --git a/vllm_omni/platforms/npu/_310p/patch/__init__.py b/vllm_omni/platforms/npu/_310p/patch/__init__.py
index b3ab637ff..5bfd42429 100644
--- a/vllm_omni/platforms/npu/_310p/patch/__init__.py
+++ b/vllm_omni/platforms/npu/_310p/patch/__init__.py
@@ -7,6 +7,7 @@ from __future__ import annotations
 
 _WORKER_PATCHED = False
 _QWEN3_TTS_TALKER_ARCH = "Qwen3TTSTalkerForConditionalGeneration"
+_QWEN3_TTS_CODE2WAV_ARCH = "Qwen3TTSCode2Wav"
 
 
 def apply_patches() -> None:
@@ -22,9 +23,13 @@ def apply_patches() -> None:
 
 
 def apply_model_patches(model_config) -> None:
-    if getattr(model_config, "model_arch", None) != _QWEN3_TTS_TALKER_ARCH:
+    model_arch = getattr(model_config, "model_arch", None)
+    if model_arch == _QWEN3_TTS_TALKER_ARCH:
+        from vllm_omni.platforms.npu._310p.patch.qwen3_tts import apply_talker_patches
+
+        apply_talker_patches()
         return
+    elif model_arch == _QWEN3_TTS_CODE2WAV_ARCH:
+        from vllm_omni.platforms.npu._310p.patch.qwen3_tts import apply_code2wav_patches
 
-    from vllm_omni.platforms.npu._310p.patch.qwen3_tts import apply_talker_patches
-
-    apply_talker_patches()
+        apply_code2wav_patches()
diff --git a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py
index c43cdd3cf..713b67f81 100644
--- a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py
+++ b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py
@@ -5,26 +5,30 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 import numpy as np
 import torch
-from transformers.models.mimi import modeling_mimi
+import torch.nn as nn
+import torch.nn.functional as F
+import torch_npu
 from vllm.multimodal.audio import AudioResampler
+from vllm_ascend._310p.attention.attention_mask import AttentionMaskBuilder310
+from vllm_ascend.sample.sampler import apply_top_k_top_p, random_sample
+from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, aligned_16, maybe_trans_nz, nd_to_nz_2d
 
 from vllm_omni.model_executor.models.common import qwen3_code_predictor
-from vllm_omni.model_executor.models.qwen3_tts import prompt_embeds_builder, qwen3_tts_talker
+from vllm_omni.model_executor.models.qwen3_tts import (
+    prompt_embeds_builder,
+    qwen3_tts_code2wav,
+    qwen3_tts_code_predictor_vllm,
+    qwen3_tts_talker,
+)
 
 _RUNTIME_DTYPE = torch.float16
 _CPU_DEVICE = torch.device("cpu")
-
-
-class _MimiEuclideanCodebook310P(modeling_mimi.MimiEuclideanCodebook):
-    def quantize(self, hidden_states):
-        # 310P does not support torch.cdist on NPU.
-        device = hidden_states.device
-        dists = torch.cdist(
-            hidden_states[None].to(_CPU_DEVICE, torch.float32), self.embed[None].to(_CPU_DEVICE, torch.float32), p=2
-        )[0]
-        return dists.argmin(dim=-1).to(device=device)
+_PATCHED = False
+_CODE2WAV_PATCHED = False
 
 
 class _Qwen3TTSTalker310P(qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneration):
@@ -32,10 +36,18 @@ class _Qwen3TTSTalker310P(qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneratio
         super().__init__(vllm_config=vllm_config, prefix=prefix)
         self._embedding_dtype = _RUNTIME_DTYPE
         self._prompt_builder._embedding_dtype = _RUNTIME_DTYPE
+        # 310P random sampling operators cannot be captured by ACL graphs.
+        # Keep Talker-MTP eager so CodePredictor can replay its own NPU graphs
+        # without changing sampling semantics.
+        self.talker_mtp_graph_safe = False
+        self.talker_mtp_accepts_per_row_generators = True
 
     def load_weights(self, weights):
         loaded = super().load_weights(weights)
-        self.encoder.to(dtype=_RUNTIME_DTYPE)
+        # The Mimi tokenizer encoder is used only while building ref_audio
+        # prompts.  Keep this preprocessing-only module on CPU because the 310P
+        # NPU path is sensitive to changing reference-audio shapes.
+        self.encoder.to(device=_CPU_DEVICE, dtype=torch.float32)
         return loaded
 
     def _encode_ref_audio_batch(
@@ -51,7 +63,7 @@ class _Qwen3TTSTalker310P(qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneratio
             resampler = AudioResampler(target_sr=target_sr)
             wavs = [resampler.resample(w.astype(np.float32), orig_sr=int(sr)) for w in wavs]
 
-        inputs = fe(raw_audio=wavs, sampling_rate=target_sr, return_tensors="pt").to(device).to(_RUNTIME_DTYPE)
+        inputs = fe(raw_audio=wavs, sampling_rate=target_sr, return_tensors="pt").to(torch.float32)
 
         with torch.inference_mode():
             encoded = self.encoder.encode(
@@ -100,10 +112,56 @@ class _Qwen3TTSPromptEmbedsBuilder310P(prompt_embeds_builder.Qwen3TTSPromptEmbed
         return spk.to(dtype=dtype)
 
 
+# ===================================================================
+#  Code2Wav layer patches
+# ===================================================================
+#
+# Code2Wav runs under the 310P graph path after the stage-0 Talker has
+# produced codec tokens.  Common NPU weight packing and fused tokenizer ops
+# live in the shared model code; the 310P patch only selects the runtime dtype.
+
+
+class _Qwen3TTSCode2Wav310P(qwen3_tts_code2wav.Qwen3TTSCode2Wav):
+    """Qwen3-TTS Code2Wav specialized for the 310P NPU path."""
+
+    def _npu_decoder_runtime_dtype(self, device: torch.device) -> torch.dtype:
+        return _RUNTIME_DTYPE
+
+
+# ===================================================================
+#  CodePredictor layer patches
+# ===================================================================
+#
+# Keep the portable implementation in common/qwen3_code_predictor.py.
+# The overrides below are installed only by the 310P platform patch because
+# the short CodePredictor loop is graph-captured on 310P and needs the 310P
+# flash-attention mask layout plus loop-local projection and sampling changes.
+
+
 class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttention):
+    """Attention override using 310P RoPE and flash-attention kernels.
+
+    The shared attention path is written in portable PyTorch.  This override
+    keeps the 310P-specific RoPE op, token alignment, and FRACTAL_NZ mask path
+    scoped to the platform patch.
+    """
+
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self._buffers.pop("_fusion_causal_mask", None)
+        self._q_size = self.num_heads * self.head_dim
+        self._kv_size = self.num_kv_heads * self.head_dim
+        self._fused_qkv_weight = None
+        self._fused_qkv_bias = None
+
+    def prepare_qkv_weights(self) -> None:
+        # Pack QKV once so each graph replay uses one matmul and consumes the
+        # weight directly in the 310P matmul layout.
+        self._fused_qkv_weight = maybe_trans_nz(
+            torch.cat((self.q_proj.weight, self.k_proj.weight, self.v_proj.weight), dim=0).contiguous()
+        )
+        if self.q_proj.bias is not None:
+            self._fused_qkv_bias = torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias), dim=0)
 
     def forward(
         self,
@@ -111,32 +169,29 @@ class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttenti
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        if hidden_states.device.type != "npu" or attention_mask is None:
-            return super().forward(hidden_states, position_embeddings)
-
         bsz, seq_len, _ = hidden_states.shape
-        hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim)
-        hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim)
-
-        q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2)
-        k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2)
-        v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2)
+        qkv = F.linear(hidden_states, self._fused_qkv_weight, self._fused_qkv_bias)
+        q, k, v = qkv.split((self._q_size, self._kv_size, self._kv_size), dim=-1)
+        q = self.q_norm(q.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2)
+        k = self.k_norm(k.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2)
+        v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
 
         cos, sin = position_embeddings
         cos = cos.unsqueeze(1)
         sin = sin.unsqueeze(1)
-        q = (q * cos) + (qwen3_code_predictor._rotate_half(q) * sin)
-        k = (k * cos) + (qwen3_code_predictor._rotate_half(k) * sin)
+        # Use the fused Ascend RoPE op instead of expanding RoPE into
+        # elementwise mul/add/rotate-half kernels.
+        q = torch_npu.npu_rotary_mul(q, cos, sin)
+        k = torch_npu.npu_rotary_mul(k, cos, sin)
 
         real_tokens = int(bsz) * int(seq_len)
         output_dtype = q.dtype
 
-        import torch_npu
-        from vllm_ascend.utils import aligned_16
-
-        q_f = aligned_16(q.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_heads, self.head_dim))
-        k_f = aligned_16(k.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim))
-        v_f = aligned_16(v.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim))
+        # 310P flash attention consumes token-major fp16 inputs with 16-token
+        # alignment; seq_lens carries the padding information.
+        q_f = aligned_16(q.transpose(1, 2).reshape(real_tokens, self.num_heads, self.head_dim))
+        k_f = aligned_16(k.transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim))
+        v_f = aligned_16(v.transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim))
 
         aligned_tokens = int(q_f.shape[0])
         seq_lens = torch.full((int(bsz),), int(seq_len), dtype=torch.int32, device="cpu")
@@ -160,23 +215,34 @@ class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttenti
 
 
 class _Qwen3CodePredictorDecoderLayer310P(qwen3_code_predictor.CodePredictorDecoderLayer):
+    """Decoder layer override that passes the 310P attention mask."""
+
     def forward(
         self,
         hidden_states: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        attention_mask: torch.Tensor | None = None,
+        attention_mask: torch.Tensor,
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
         hidden_states = self.self_attn(hidden_states, position_embeddings, attention_mask=attention_mask)
-        hidden_states = residual + hidden_states
-
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states, _, residual = torch_npu.npu_add_rms_norm(
+            hidden_states,
+            residual,
+            self.post_attention_layernorm.weight,
+            self.post_attention_layernorm.variance_epsilon,
+        )
         return residual + self.mlp(hidden_states)
 
 
 class _Qwen3CodePredictorBaseModel310P(qwen3_code_predictor.CodePredictorBaseModel):
+    """Base model override with a cached 310P causal mask.
+
+    The 310P flash-attention path consumes the additive causal mask in
+    FRACTAL_NZ format.  The CodePredictor sequence length is fixed and short,
+    so the mask is built once and reused across graph replays.
+    """
+
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self._attention_mask_310p = None
@@ -188,14 +254,9 @@ class _Qwen3CodePredictorBaseModel310P(qwen3_code_predictor.CodePredictorBaseMod
         inputs_embeds: torch.Tensor,
         position_ids: torch.Tensor,
     ) -> torch.Tensor:
-        if inputs_embeds.device.type != "npu":
-            return super().forward(inputs_embeds, position_ids)
-
         if self._attention_mask_310p is None or self._attention_mask_310p_device != inputs_embeds.device:
-            import torch_npu
-            from vllm_ascend._310p.attention.attention_mask import AttentionMaskBuilder310
-            from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, nd_to_nz_2d
-
+            # Store the additive causal mask in the format consumed by the
+            # 310P flash-attention kernel.
             mask = AttentionMaskBuilder310.gen_causal_additive_mask(
                 self._attention_mask_310p_max_seq,
                 inputs_embeds.device,
@@ -205,23 +266,239 @@ class _Qwen3CodePredictorBaseModel310P(qwen3_code_predictor.CodePredictorBaseMod
 
         input_dtype = inputs_embeds.dtype
         hidden_states = inputs_embeds
-        with torch.amp.autocast(inputs_embeds.device.type, enabled=False, dtype=torch.float32):
-            position_embeddings = self.rotary_emb(hidden_states, position_ids)
-            for layer in self.layers:
-                hidden_states = layer(
-                    hidden_states,
-                    position_embeddings,
-                    attention_mask=self._attention_mask_310p,
-                )
-            hidden_states = self.norm(hidden_states)
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        for layer in self.layers:
+            hidden_states = layer(
+                hidden_states,
+                position_embeddings,
+                attention_mask=self._attention_mask_310p,
+            )
+        hidden_states = self.norm(hidden_states)
         return hidden_states.to(input_dtype)
 
 
+class _Qwen3TTSTalkerCodePredictor310P(
+    qwen3_tts_code_predictor_vllm.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
+):
+    """Qwen3-TTS code predictor specialized for the 310P NPU path."""
+
+    def __init__(
+        self,
+        *,
+        vllm_config,
+        config,
+        talker_config,
+        prefix: str = "code_predictor",
+    ) -> None:
+        super().__init__(
+            vllm_config=vllm_config,
+            config=config,
+            talker_config=talker_config,
+            prefix=prefix,
+        )
+        self._projected_codec_embed_weight = None
+
+    def _prepare_npu_weights(self) -> None:
+        qkv_projections = set()
+        with torch.no_grad():
+            for layer in self.model.layers:
+                attention = layer.self_attn
+                attention.prepare_qkv_weights()
+                qkv_projections.update((attention.q_proj, attention.k_proj, attention.v_proj))
+
+            for module in self.modules():
+                if isinstance(module, nn.Linear) and module not in qkv_projections:
+                    module.weight.data = maybe_trans_nz(module.weight.data)
+
+    def load_weights(self, weights):
+        loaded = super().load_weights(weights)
+        with torch.no_grad():
+            self._projected_codec_embed_weight = torch.stack(
+                [self.small_to_mtp_projection(embed.weight).detach() for embed in self.model.codec_embedding],
+                dim=0,
+            ).contiguous()
+        return loaded
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        layer0_code: torch.Tensor,
+        layer0_embed: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        do_sample: bool = True,
+        temperature: float = 0.9,
+        top_k: int = 50,
+        top_p: float = 1.0,
+        generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
+    ) -> torch.Tensor:
+        bsz = int(layer0_code.shape[0])
+        if generators is not None and len(generators) != bsz:
+            raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
+        num_groups = self._num_groups
+        device = layer0_code.device
+
+        self._setup_compile()
+        dtype = self._model_dtype
+
+        padded_bsz = self._padded_bsz(bsz)
+        self._ensure_buffers(device, dtype, padded_bsz)
+
+        proj_buf = self._proj_buf
+        max_seq = num_groups + 1
+        projection = self.small_to_mtp_projection
+        model_fwd = self._compiled_model_fwd
+        lm_heads = self._lm_heads_list
+        if generators is not None:
+            npu_generators = {i: row_generator for i, row_generator in enumerate(generators) if row_generator}
+        elif generator is not None:
+            npu_generators = {i: generator for i in range(bsz)}
+        else:
+            npu_generators = {}
+
+        proj_buf[:padded_bsz].zero_()
+        initial_embeds = torch.cat(
+            (
+                last_talker_hidden.reshape(bsz, 1, -1),
+                layer0_embed.reshape(bsz, 1, -1),
+            ),
+            dim=1,
+        )
+        if initial_embeds.dtype != dtype:
+            initial_embeds = initial_embeds.to(dtype)
+        proj_buf[:bsz, :2, :].copy_(projection(initial_embeds))
+
+        stored_mode = self._wrapper_config.sampling_mode == "stored"
+        if stored_mode:
+            s_top_k = self._top_k
+            s_top_p = self._top_p
+        else:
+            use_sampling = do_sample and temperature > 0
+            inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0
+            if use_sampling and top_p != 1.0:
+                raise NotImplementedError(
+                    "top_p sampling is not implemented for the vLLM-native code predictor; please set top_p=1.0."
+                )
+
+        top_k_tensor = None
+        top_p_tensor = None
+        if stored_mode:
+            top_k_hint = s_top_k if s_top_k > 0 else None
+            if s_top_k > 0:
+                top_k_tensor = torch.full((bsz,), s_top_k, dtype=torch.int32, device=device)
+            if s_top_p < 1.0:
+                top_p_tensor = torch.full((bsz,), s_top_p, dtype=dtype, device=device)
+        elif use_sampling:
+            top_k_hint = top_k if top_k > 0 else None
+            if top_k > 0:
+                top_k_tensor = torch.full((bsz,), top_k, dtype=torch.int32, device=device)
+        else:
+            top_k_hint = None
+
+        all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device)
+        all_codes[:, 0] = layer0_code.reshape(bsz)
+
+        for step in range(1, num_groups):
+            graph_key: int | tuple[int, int] = padded_bsz
+            seq_len = max_seq
+            if self._prefix_graphs_enabled:
+                prefix_key = (padded_bsz, step + 1)
+                if prefix_key in self._device_graphs:
+                    graph_key = prefix_key
+                    seq_len = step + 1
+            pos_ids = self._bucket_pos_ids.get(graph_key)
+            if pos_ids is None:
+                pos_ids = (
+                    torch.arange(seq_len, device=device, dtype=torch.long)
+                    .unsqueeze(0)
+                    .expand(padded_bsz, -1)
+                    .contiguous()
+                )
+
+            device_graph_entry = self._device_graphs.get(graph_key)
+            if device_graph_entry is not None:
+                device_graph_entry[0].replay()
+                hidden_out = device_graph_entry[1]
+            else:
+                hidden_out = model_fwd(proj_buf[:padded_bsz, :seq_len, :], pos_ids)
+
+            logits = lm_heads[step - 1](hidden_out[:bsz, step, :])
+
+            if stored_mode:
+                if top_k_tensor is not None or top_p_tensor is not None:
+                    logits = apply_top_k_top_p(logits, p=top_p_tensor, k=top_k_tensor, top_k=top_k_hint)
+                candidate_indices = None
+                if isinstance(logits, tuple):
+                    logits, candidate_indices = logits
+                probs = F.softmax(logits, dim=-1, dtype=torch.float32)
+                code = random_sample(probs, npu_generators)
+                if candidate_indices is not None:
+                    code = candidate_indices.gather(1, code.unsqueeze(1)).squeeze(1)
+            else:
+                if use_sampling:
+                    scaled = logits * inv_temperature
+                    if top_k_tensor is not None:
+                        scaled = apply_top_k_top_p(scaled, p=None, k=top_k_tensor, top_k=top_k_hint)
+                    candidate_indices = None
+                    if isinstance(scaled, tuple):
+                        scaled, candidate_indices = scaled
+                    probs = F.softmax(scaled, dim=-1, dtype=torch.float32)
+                    code = random_sample(probs, npu_generators)
+                    if candidate_indices is not None:
+                        code = candidate_indices.gather(1, code.unsqueeze(1)).squeeze(1)
+                else:
+                    code = logits.argmax(dim=-1, keepdim=True)
+
+            all_codes[:, step] = code.reshape(bsz)
+            if step < num_groups - 1:
+                proj_buf[:bsz, step + 1, :].copy_(
+                    F.embedding(code.reshape(-1), self._projected_codec_embed_weight[step - 1])
+                )
+
+        return all_codes
+
+
+# ===================================================================
+#  Patch registration
+# ===================================================================
+
+
 def apply_talker_patches() -> None:
-    modeling_mimi.MimiEuclideanCodebook = _MimiEuclideanCodebook310P
+    """Install Qwen3-TTS Talker and CodePredictor 310P patches.
+
+    The generic model modules stay unchanged.  Patch registration swaps in the
+    310P-specialized CodePredictor classes and wrapper methods only when the
+    310P platform applies the Talker patch.
+    """
+
+    global _PATCHED
+
+    if _PATCHED:
+        return
+
     qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneration = _Qwen3TTSTalker310P
     qwen3_tts_talker.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilder310P
+    qwen3_tts_talker.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM = _Qwen3TTSTalkerCodePredictor310P
+    qwen3_tts_code_predictor_vllm.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM = (
+        _Qwen3TTSTalkerCodePredictor310P
+    )
+    qwen3_tts_code_predictor_vllm.CodePredictorBaseModel = _Qwen3CodePredictorBaseModel310P
+    qwen3_tts_code_predictor_vllm.Qwen3TTSTalkerCodePredictorModelVLLM = _Qwen3CodePredictorBaseModel310P
     prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilder310P
     qwen3_code_predictor.CodePredictorAttention = _Qwen3CodePredictorAttention310P
     qwen3_code_predictor.CodePredictorDecoderLayer = _Qwen3CodePredictorDecoderLayer310P
     qwen3_code_predictor.CodePredictorBaseModel = _Qwen3CodePredictorBaseModel310P
+
+    _PATCHED = True
+
+
+def apply_code2wav_patches() -> None:
+    """Install the 310P Code2Wav runtime patch."""
+    global _CODE2WAV_PATCHED
+
+    if _CODE2WAV_PATCHED:
+        return
+
+    qwen3_tts_code2wav.Qwen3TTSCode2Wav = _Qwen3TTSCode2Wav310P
+
+    _CODE2WAV_PATCHED = True
diff --git a/vllm_omni/platforms/npu/_310p/patch/worker.py b/vllm_omni/platforms/npu/_310p/patch/worker.py
index 3116b54da..49ee6052c 100644
--- a/vllm_omni/platforms/npu/_310p/patch/worker.py
+++ b/vllm_omni/platforms/npu/_310p/patch/worker.py
@@ -12,6 +12,9 @@ compile-mode setup is applied here after device initialization.
 
 from __future__ import annotations
 
+from vllm.v1.sample.sampler import Sampler
+from vllm_ascend._310p.sample.sampler import AscendSampler310
+
 from vllm_omni.platforms.npu._310p import disable_jit_compile
 from vllm_omni.platforms.npu.worker import base as worker_base
 
@@ -24,4 +27,6 @@ class _OmniNPUWorkerBase310P(worker_base.OmniNPUWorkerBase):
 
 
 def apply_patch() -> None:
+    # Triton-Ascend does not target 310P; use vLLM's native penalty path.
+    AscendSampler310.apply_penalties = staticmethod(Sampler.apply_penalties)
     worker_base.OmniNPUWorkerBase = _OmniNPUWorkerBase310P
diff --git a/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py b/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py
index 282689bca..bd74e9dc6 100644
--- a/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py
+++ b/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py
@@ -1,15 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Monkey-patch ``Qwen3TTSCode2Wav.__init__`` for NPU Code2Wav runtime knobs."""
+"""Patch Qwen3-TTS Code2Wav NPU runtime setup and weight preparation."""
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
 import torch
+import torch.nn as nn
+import torch_npu
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
+from vllm_ascend.utils import maybe_trans_nz
 
 if TYPE_CHECKING:
     pass
@@ -18,6 +21,8 @@ logger = init_logger(__name__)
 
 _PATCHED = False
 _original_init = None
+_original_load_weights = None
+_ACL_FORMAT_FRACTAL_Z = 4
 
 
 def _prepare_npu_code2wav_runtime() -> None:
@@ -35,14 +40,43 @@ def _patched_init(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
     _original_init(self, vllm_config=vllm_config, prefix=prefix)
 
 
+def _prepare_npu_decoder_weights(decoder: nn.Module) -> None:
+    linear_count = 0
+    conv_count = 0
+    with torch.no_grad():
+        for module in decoder.modules():
+            if isinstance(module, nn.Linear):
+                module.weight.data = maybe_trans_nz(module.weight.data)
+                linear_count += 1
+            elif isinstance(module, (nn.Conv1d, nn.ConvTranspose1d)) and module.groups == 1:
+                module.weight.data = torch_npu.npu_format_cast(module.weight.data.contiguous(), _ACL_FORMAT_FRACTAL_Z)
+                conv_count += 1
+
+    logger.info("Prepared NPU Code2Wav weights: linear=%d conv=%d", linear_count, conv_count)
+
+
+def _patched_load_weights(self, weights):
+    assert _original_load_weights is not None
+    loaded = _original_load_weights(self, weights)
+    device = self.vllm_config.device_config.device
+    runtime_dtype = getattr(self, "_npu_decoder_runtime_dtype", lambda _: torch.float32)(device)
+    self.decoder.to(device=device, dtype=runtime_dtype)
+    _prepare_npu_decoder_weights(self.decoder)
+    if runtime_dtype != torch.float32 and hasattr(self.decoder, "precompute_snake_caches"):
+        self.decoder.precompute_snake_caches()
+    return loaded
+
+
 def apply_qwen3_tts_code2wav_patch() -> None:
-    global _PATCHED, _original_init
+    global _PATCHED, _original_init, _original_load_weights
     if _PATCHED:
         return
 
     from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import Qwen3TTSCode2Wav
 
     _original_init = Qwen3TTSCode2Wav.__init__
+    _original_load_weights = Qwen3TTSCode2Wav.load_weights
     Qwen3TTSCode2Wav.__init__ = _patched_init  # type: ignore[method-assign]
+    Qwen3TTSCode2Wav.load_weights = _patched_load_weights  # type: ignore[method-assign]
     _PATCHED = True
-    logger.debug("Applied NPU patch for Qwen3TTSCode2Wav.__init__")
+    logger.debug("Applied NPU patch for Qwen3TTSCode2Wav")
diff --git a/vllm_omni/platforms/npu/models/qwen3_tts_tokenizer_v2.py b/vllm_omni/platforms/npu/models/qwen3_tts_tokenizer_v2.py
new file mode 100644
index 000000000..f89ecded9
--- /dev/null
+++ b/vllm_omni/platforms/npu/models/qwen3_tts_tokenizer_v2.py
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""NPU patches for the Qwen3-TTS 12Hz tokenizer decoder."""
+
+from __future__ import annotations
+
+import torch
+import torch_npu
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_PATCHED = False
+
+
+def _apply_rotary_pos_emb_npu(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    position_ids=None,
+    unsqueeze_dim: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    del position_ids
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    return torch_npu.npu_rotary_mul(q, cos, sin), torch_npu.npu_rotary_mul(k, cos, sin)
+
+
+def _rms_norm_forward_npu(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    return torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.variance_epsilon)[0]
+
+
+def apply_qwen3_tts_tokenizer_v2_patch() -> None:
+    global _PATCHED
+    if _PATCHED:
+        return
+
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import (
+        modeling_qwen3_tts_tokenizer_v2,
+    )
+
+    modeling_qwen3_tts_tokenizer_v2.apply_rotary_pos_emb = _apply_rotary_pos_emb_npu
+    modeling_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2DecoderRMSNorm.forward = _rms_norm_forward_npu
+    _PATCHED = True
+    logger.debug("Applied NPU patch for Qwen3-TTS 12Hz tokenizer decoder")
diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py
index f03fdb1bf..fe49455d5 100644
--- a/vllm_omni/platforms/npu/platform.py
+++ b/vllm_omni/platforms/npu/platform.py
@@ -37,8 +37,12 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
         from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import (
             apply_qwen3_tts_code2wav_patch,
         )
+        from vllm_omni.platforms.npu.models.qwen3_tts_tokenizer_v2 import (
+            apply_qwen3_tts_tokenizer_v2_patch,
+        )
 
         apply_qwen3_tts_code2wav_patch()
+        apply_qwen3_tts_tokenizer_v2_patch()
         apply_310p_patches()
 
     @classmethod
@@ -107,6 +111,14 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
+                logger.warning(
+                    "HuggingFace kernels-backed FlashAttention is "
+                    "not supported on NPU. Falling back to local "
+                    "FLASH_ATTN."
+                )
+                backend_upper = "FLASH_ATTN"
+
             backend = DiffusionAttentionBackendEnum[backend_upper]
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
index c66d493c2..9a8a55709 100644
--- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py
@@ -46,6 +46,7 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTran
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner
 from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
+from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
 
 
 def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]:
@@ -479,7 +480,6 @@ class NPUARModelRunner(OmniNPUModelRunner):
                     logits_indices,
                     spec_decode_metadata,
                     total_num_scheduled_tokens,
-                    num_scheduled_tokens_compressed_list,
                 ) = self._prepare_inputs(
                     scheduler_output,
                     num_scheduled_tokens_np,
@@ -576,7 +576,12 @@ class NPUARModelRunner(OmniNPUModelRunner):
                     # Another possible condition is num_tokens_padded != num_tokens_unpadded
                     # but this scope is way too big and the consequences are unpredictable
                     num_reqs_padded = self._pad_query_start_loc_for_fia(
-                        num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_mode, batch_desc.num_reqs
+                        self.query_start_loc,
+                        num_tokens_padded,
+                        num_reqs_padded,
+                        num_reqs,
+                        cudagraph_mode,
+                        batch_desc.num_reqs,
                     )
 
                 (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
@@ -617,8 +622,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
                     positions=positions,
                     inputs_embeds=inputs_embeds,
                     req_ids=req_ids[:num_reqs],
-                    num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)],
-                    num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)],
+                    num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs],
+                    num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
                     input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],
                 )
             #  -------------------------------------- Omni-new -------------------------------------------------
@@ -686,8 +691,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
         with record_function_or_nullcontext("post process"):
             #  -------------------------------------- Omni-new -------------------------------------------------
             # [Omni] Map pending ropes metadata to req_ids.
-            if hasattr(self.model, "flush_pending_metadata"):
-                self.model.flush_pending_metadata(list(req_ids))
+            flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None)
+            if callable(flush_pending_metadata):
+                flush_pending_metadata(req_ids[:num_reqs])
 
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
 
@@ -897,7 +903,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
         ) = self.execute_model_state
         # Clear ephemeral state.
         self.execute_model_state = None
-        seq_len = hidden_states.shape[0]
+        hidden_seq_len = int(hidden_states.shape[0])
+        scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
 
         # Apply structured output bitmasks if present.
         if grammar_output is not None:
@@ -916,6 +923,15 @@ class NPUARModelRunner(OmniNPUModelRunner):
                 logits_vocab = logits.shape[-1]
                 if self.input_batch.vocab_size > logits_vocab:
                     smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab)
+
+        # Drop min-tokens stop ids the head cannot emit (e.g. the text
+        # tokenizer EOS folded into all_stop_token_ids on a narrow codec
+        # talker head); they would index_put_ out of bounds (#4962).
+        if logits is not None:
+            sanitize_min_tokens_stop_ids(
+                self.input_batch.sampling_metadata.logitsprocs,
+                logits.shape[-1],
+            )
         #  -------------------------------------- Omni-new -------------------------------------------------
 
 
@@ -1136,7 +1152,8 @@ class NPUARModelRunner(OmniNPUModelRunner):
                                 start=start,
                                 end=end,
                                 pass_lists_through=False,
-                                seq_len=seq_len,
+                                seq_len=hidden_seq_len,
+                                scheduled_seq_len=scheduled_seq_len,
                             )
                     payload.update(mm_payload)
                 pooler_output.append(flatten_payload(payload))
diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
index df8b92bf5..ea03842b4 100644
--- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py
@@ -164,7 +164,6 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                     logits_indices,
                     spec_decode_metadata,
                     total_num_scheduled_tokens,
-                    num_scheduled_tokens_compressed_list,
                 ) = self._prepare_inputs(
                     scheduler_output,
                     num_scheduled_tokens_np,
@@ -261,7 +260,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
                     # Another possible condition is num_tokens_padded != num_tokens_unpadded
                     # but this scope is way too big and the consequences are unpredictable
                     num_reqs_padded = self._pad_query_start_loc_for_fia(
-                        num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_mode, batch_desc.num_reqs
+                        self.query_start_loc,
+                        num_tokens_padded,
+                        num_reqs_padded,
+                        num_reqs,
+                        cudagraph_mode,
+                        batch_desc.num_reqs,
                     )
 
                 (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata(
@@ -729,7 +733,12 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
 
             if not profile_cpp:
                 num_reqs_padded = self._pad_query_start_loc_for_fia(
-                    num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_runtime_mode, batch_desc.num_reqs
+                    self.query_start_loc,
+                    num_tokens_padded,
+                    num_reqs_padded,
+                    num_reqs,
+                    cudagraph_runtime_mode,
+                    batch_desc.num_reqs,
                 )
 
             pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py
index 3b9b4937a..671c7b7e5 100644
--- a/vllm_omni/platforms/npu/worker/npu_model_runner.py
+++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py
@@ -193,7 +193,12 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner):
 
             if not profile_cpp:
                 num_reqs_padded = self._pad_query_start_loc_for_fia(
-                    num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_runtime_mode, batch_desc.num_reqs
+                    self.query_start_loc,
+                    num_tokens_padded,
+                    num_reqs_padded,
+                    num_reqs,
+                    cudagraph_runtime_mode,
+                    batch_desc.num_reqs,
                 )
 
             pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py
index b932d3794..f905fd73c 100644
--- a/vllm_omni/platforms/rocm/platform.py
+++ b/vllm_omni/platforms/rocm/platform.py
@@ -88,6 +88,14 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
+                logger.warning(
+                    "HuggingFace kernels-backed FlashAttention is "
+                    "not supported on ROCm. Falling back to local "
+                    "FLASH_ATTN."
+                )
+                backend_upper = "FLASH_ATTN"
+
             if backend_upper == "FLASH_ATTN" and not aiter_supported:
                 logger.warning(
                     "Flash Attention requires `aiter` library which is only supported "
diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py
index 8be01939f..0bc782e03 100644
--- a/vllm_omni/platforms/xpu/platform.py
+++ b/vllm_omni/platforms/xpu/platform.py
@@ -42,6 +42,14 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
+                logger.warning(
+                    "HuggingFace kernels-backed FlashAttention is "
+                    "not supported on XPU. Falling back to local "
+                    "FLASH_ATTN."
+                )
+                backend_upper = "FLASH_ATTN"
+
             backend = DiffusionAttentionBackendEnum[backend_upper]
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
diff --git a/vllm_omni/plugins/__init__.py b/vllm_omni/plugins/__init__.py
index 41bf398ca..240394b2c 100644
--- a/vllm_omni/plugins/__init__.py
+++ b/vllm_omni/plugins/__init__.py
@@ -68,6 +68,10 @@ def load_omni_general_plugins() -> None:
         return
     omni_plugins_loaded = True
 
+    from vllm.plugins import load_general_plugins
+
+    load_general_plugins()
+
     plugins = load_omni_plugins_by_group(group=OMNI_DEFAULT_PLUGINS_GROUP)
     # General plugins: we only need to execute the loaded functions.
     for func in plugins.values():
diff --git a/vllm_omni/quantization/__init__.py b/vllm_omni/quantization/__init__.py
index d8cad3a98..8896903e8 100644
--- a/vllm_omni/quantization/__init__.py
+++ b/vllm_omni/quantization/__init__.py
@@ -12,10 +12,10 @@ Adds per-component quantization for multi-stage models.
 """
 
 from .component_config import ComponentQuantizationConfig
-from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config
+from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config, register_quantization_override
 from .inc_config import OmniINCConfig
 
-# Heavy configs (GGUF, MXFP8) are NOT imported here to avoid pulling in
+# Heavy configs are NOT imported here to avoid pulling in
 # optional dependencies (pynvml, torch_npu) at module load time.
 # Import them directly when needed:
 #   from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
@@ -25,4 +25,5 @@ __all__ = [
     "ComponentQuantizationConfig",
     "OmniINCConfig",
     "SUPPORTED_QUANTIZATION_METHODS",
+    "register_quantization_override",
 ]
diff --git a/vllm_omni/quantization/bitsandbytes_config.py b/vllm_omni/quantization/bitsandbytes_config.py
new file mode 100644
index 000000000..7d8e7a684
--- /dev/null
+++ b/vllm_omni/quantization/bitsandbytes_config.py
@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""BitsAndBytes 4-bit quantization config for diffusion transformers.
+
+Supports online (dynamic) NF4/FP4 weight-only quantization from BF16/FP16
+checkpoints on CUDA GPUs.
+"""
+
+from typing import TYPE_CHECKING, Any, Optional
+
+import torch
+from torch.nn import Module
+from vllm.model_executor.layers.linear import (
+    LinearBase,
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from vllm.model_executor.layers.quantization import QuantizationMethods
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from vllm.model_executor.layers.quantization.fp8 import _copy_missing_attrs
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    is_layer_skipped,
+)
+from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight
+from vllm.model_executor.parameter import ModelWeightParameter
+from vllm.model_executor.utils import replace_parameter
+
+from vllm_omni.platforms import current_omni_platform
+from vllm_omni.quantization.int8_config import LazyWeightMixin
+
+if TYPE_CHECKING:
+    from vllm.model_executor.models.utils import WeightsMapper
+
+
+class DiffusionBitsAndBytesConfig(QuantizationConfig):
+    """BitsAndBytes 4-bit weight-only config for diffusion transformers.
+
+    Supports online (dynamic) quantization from BF16/FP16 checkpoints.
+    Works on CUDA GPUs with the optional ``bitsandbytes`` package installed.
+    """
+
+    def __init__(
+        self,
+        quant_type: str = "nf4",
+        compress_statistics: bool = True,
+        ignored_layers: list[str] | None = None,
+    ) -> None:
+        super().__init__()
+
+        if quant_type not in ("nf4", "fp4"):
+            raise ValueError(f"Unsupported quant_type {quant_type!r}; expected 'nf4' or 'fp4'")
+        self.quant_type = quant_type
+        self.compress_statistics = compress_statistics
+        self.ignored_layers = ignored_layers or []
+
+    @classmethod
+    def get_name(cls) -> QuantizationMethods:
+        return "bitsandbytes"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 75
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []
+
+    def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
+        if self.ignored_layers is not None:
+            self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers)
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "DiffusionBitsAndBytesConfig":
+        quant_type = cls.get_from_keys_or(config, ["quant_type"], "nf4")
+        compress_statistics = cls.get_from_keys_or(config, ["compress_statistics"], True)
+        ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)
+
+        if not ignored_layers:
+            ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None)
+        return cls(
+            quant_type=quant_type,
+            compress_statistics=compress_statistics,
+            ignored_layers=ignored_layers,
+        )
+
+    def get_quant_method(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+    ) -> Optional["QuantizeMethodBase"]:
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped(
+                prefix=prefix,
+                ignored_layers=self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedLinearMethod()
+            if current_omni_platform.is_cuda():
+                return BnBOnlineLinearMethod(self)
+            raise NotImplementedError("BitsAndBytes online quantization is only supported on CUDA.")
+        return None
+
+
+class BnBOnlineLinearMethod(LazyWeightMixin, LinearMethodBase):
+    """Online BitsAndBytes 4-bit linear method.
+
+    Loads BF16/FP16 checkpoint weights and quantizes them during loading.
+    """
+
+    def __init__(self, quant_config: DiffusionBitsAndBytesConfig):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        if layer.weight.device == torch.device("meta"):
+            weight = ModelWeightParameter(
+                data=torch.empty_like(layer.weight, device=layer._load_device),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=layer.weight.weight_loader,
+            )
+            _copy_missing_attrs(layer.weight, weight)
+            layer.register_parameter("weight", weight)
+            initialize_single_dummy_weight(layer.weight)
+
+        import bitsandbytes.functional as bnb_F
+
+        weight = layer.weight.data.contiguous()
+        if not weight.is_cuda:
+            weight = weight.cuda()
+
+        original_shape = tuple(weight.shape)
+        qweight, quant_state = bnb_F.quantize_4bit(
+            weight,
+            quant_type=self.quant_config.quant_type,
+            compress_statistics=self.quant_config.compress_statistics,
+        )
+
+        replace_parameter(
+            layer,
+            "weight",
+            torch.nn.Parameter(qweight, requires_grad=False),
+        )
+        layer.quant_state = quant_state
+        layer.bnb_shape = original_shape
+
+        layer._already_called_process_weights_after_loading = True
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        import bitsandbytes as bnb
+
+        ori_shape = x.shape
+        ori_dtype = x.dtype
+        x_2d = x.reshape(-1, ori_shape[-1])
+
+        out = bnb.matmul_4bit(
+            x_2d,
+            layer.weight.t(),
+            quant_state=layer.quant_state,
+            bias=bias,
+        )
+        return out.reshape(*ori_shape[:-1], -1).to(ori_dtype)
diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py
index 87564ef46..28c33d637 100644
--- a/vllm_omni/quantization/component_config.py
+++ b/vllm_omni/quantization/component_config.py
@@ -80,6 +80,7 @@ class ComponentQuantizationConfig(QuantizationConfig):
         component_configs: dict[str, QuantizationConfig | None],
         default_config: QuantizationConfig | None = None,
     ) -> None:
+        super().__init__()
         self._components = component_configs
         self._default = default_config
         self._sorted_prefixes = sorted(self._components.keys(), key=len, reverse=True)
diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py
index d9a7a31bf..b5d0a8ad7 100644
--- a/vllm_omni/quantization/factory.py
+++ b/vllm_omni/quantization/factory.py
@@ -3,8 +3,8 @@
 """Factory for building quantization configs.
 
 build_quant_config() delegates to vLLM's quantization registry.
-The only extension point is _OVERRIDES for methods that need a
-different QuantizationConfig subclass in the OMNI context (e.g. GGUF).
+Out-of-tree integrations can register Omni-specific builders with
+register_quantization_override().
 """
 
 from __future__ import annotations
@@ -81,13 +81,6 @@ from .component_config import ComponentQuantizationConfig  # noqa: E402
 logger = init_logger(__name__)
 
 
-def _build_gguf(**kw: Any) -> QuantizationConfig:
-    """Lazy import to avoid pulling in CUDA/pynvml at module load time."""
-    from .gguf_config import DiffusionGGUFConfig
-
-    return DiffusionGGUFConfig(**kw)
-
-
 def _build_int8(**kw: Any) -> QuantizationConfig:
     """Lazy import for Int8 diffusion config (supports CUDA + NPU)."""
     from .int8_config import DiffusionInt8Config
@@ -95,6 +88,13 @@ def _build_int8(**kw: Any) -> QuantizationConfig:
     return DiffusionInt8Config(**kw)
 
 
+def _build_bitsandbytes(**kw: Any) -> QuantizationConfig:
+    """Lazy import for BitsAndBytes 4-bit diffusion config (CUDA only)."""
+    from .bitsandbytes_config import DiffusionBitsAndBytesConfig
+
+    return DiffusionBitsAndBytesConfig(**kw)
+
+
 def _build_mxfp8(**kw: Any) -> QuantizationConfig:
     """Lazy import for W8A8 MXFP8 diffusion config (NPU only)."""
     from .mxfp8_config import DiffusionMXFP8Config
@@ -140,8 +140,8 @@ def _build_inc(**kw: Any) -> QuantizationConfig:
 
 
 _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
-    "gguf": _build_gguf,
     "int8": _build_int8,
+    "bitsandbytes": _build_bitsandbytes,
     "mxfp8": _build_mxfp8,
     "mxfp4": _build_mxfp4,
     "mxfp4_dualscale": _build_mxfp4_dualscale,
@@ -150,7 +150,24 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = {
     "auto_round": _build_inc,
 }
 
-SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))
+
+def _compute_supported_quantization_methods() -> list[str]:
+    return list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys())))
+
+
+SUPPORTED_QUANTIZATION_METHODS: list[str] = _compute_supported_quantization_methods()
+
+
+def _refresh_registered_methods() -> None:
+    SUPPORTED_QUANTIZATION_METHODS[:] = _compute_supported_quantization_methods()
+    global _CACHED_ALIAS_MAP
+    _CACHED_ALIAS_MAP = None
+
+
+def register_quantization_override(method: str, builder: Callable[..., QuantizationConfig]) -> None:
+    """Register an Omni-specific quantization config builder."""
+    _OVERRIDES[_normalize_method_name(method)] = builder
+    _refresh_registered_methods()
 
 
 def _build_reverse_alias_map() -> dict[str, str]:
diff --git a/vllm_omni/quantization/gguf_config.py b/vllm_omni/quantization/gguf_config.py
deleted file mode 100644
index 186ac4ee7..000000000
--- a/vllm_omni/quantization/gguf_config.py
+++ /dev/null
@@ -1,323 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""GGUF quantization config for diffusion transformers.
-
-Uses dequant+GEMM instead of the fused kernel path (which expects 2D inputs).
-
-The GGUF quantization module was migrated to an external plugin in upstream
-vLLM (commit 6635279d8).  This file inlines the base classes that
-DiffusionGGUFConfig / DiffusionGGUFLinearMethod depend on, reproduced from
-the last upstream version before the migration.
-"""
-
-from __future__ import annotations
-
-from collections.abc import Mapping
-from types import MappingProxyType
-from typing import TYPE_CHECKING, Any
-
-import gguf
-import torch
-from gguf import GGMLQuantizationType as WeightType
-from torch.nn.parameter import Parameter, UninitializedParameter
-from vllm import _custom_ops as ops
-from vllm.logger import init_logger
-from vllm.model_executor.layers.linear import (
-    LinearBase,
-    LinearMethodBase,
-    UnquantizedLinearMethod,
-)
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig,
-    QuantizeMethodBase,
-)
-from vllm.model_executor.utils import set_weight_attrs
-
-if TYPE_CHECKING:
-    from vllm.model_executor.layers.quantization import QuantizationMethods
-
-logger = init_logger(__name__)
-
-# ---------------------------------------------------------------------------
-# GGUF base types (inlined from upstream vllm before migration to plugin)
-# ---------------------------------------------------------------------------
-
-UNQUANTIZED_TYPES: set[WeightType] = {WeightType.F32, WeightType.F16, WeightType.BF16}
-
-STANDARD_QUANT_TYPES: set[WeightType] = {
-    WeightType.Q4_0,
-    WeightType.Q4_1,
-    WeightType.Q5_0,
-    WeightType.Q5_1,
-    WeightType.Q8_0,
-    WeightType.Q8_1,
-}
-KQUANT_TYPES: set[WeightType] = {
-    WeightType.Q2_K,
-    WeightType.Q3_K,
-    WeightType.Q4_K,
-    WeightType.Q5_K,
-    WeightType.Q6_K,
-}
-IMATRIX_QUANT_TYPES: set[WeightType] = {
-    WeightType.IQ1_M,
-    WeightType.IQ1_S,
-    WeightType.IQ2_XXS,
-    WeightType.IQ2_XS,
-    WeightType.IQ2_S,
-    WeightType.IQ3_XXS,
-    WeightType.IQ3_S,
-    WeightType.IQ4_XS,
-    WeightType.IQ4_NL,
-}
-DEQUANT_TYPES: set[WeightType] = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES
-
-
-def is_layer_skipped_gguf(
-    prefix: str,
-    unquantized_modules: list[str],
-    fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
-) -> bool:
-    # Fused layers like gate_up_proj or qkv_proj will not be fused
-    # in the safetensors checkpoint. So, we convert the name
-    # from the fused version to unfused + check to make sure that
-    # each shard of the fused layer has the same scheme.
-    proj_name = prefix.split(".")[-1]
-    if proj_name in fused_mapping:
-        shard_prefixes = [prefix.replace(proj_name, shard_proj_name) for shard_proj_name in fused_mapping[proj_name]]
-
-        is_skipped = None
-        for shard_prefix in shard_prefixes:
-            is_shard_skipped = any(shard_prefix in module_name for module_name in unquantized_modules)
-
-            if is_skipped is None:
-                is_skipped = is_shard_skipped
-            elif is_shard_skipped != is_skipped:
-                raise ValueError(
-                    f"Detected some but not all shards of {prefix} "
-                    "are quantized. All shards of fused layers "
-                    "to have the same precision."
-                )
-    else:
-        is_skipped = any(module_name in prefix for module_name in unquantized_modules)
-
-    assert is_skipped is not None
-    return is_skipped
-
-
-class _GGUFUninitializedParameter(UninitializedParameter):
-    """UninitializedParameter subclass that stores GGUF weight data in a list.
-
-    The actual tensor is materialized later via _create_padded_weight_param.
-    """
-
-    cls_to_become = Parameter
-    data_container: list[torch.Tensor]
-
-
-class GGUFConfig(QuantizationConfig):
-    """Config class for GGUF (inlined from upstream vllm)."""
-
-    def __init__(self, unquantized_modules: list[str] | None = None) -> None:
-        super().__init__()
-        self.unquantized_modules = unquantized_modules or []
-
-    def __repr__(self) -> str:
-        return "GGUFConfig()"
-
-    def get_name(self) -> QuantizationMethods:
-        return "gguf"
-
-    def get_supported_act_dtypes(self) -> list[torch.dtype]:
-        return [torch.half, torch.bfloat16, torch.float32]
-
-    @classmethod
-    def get_min_capability(cls) -> int:
-        return 60
-
-    @classmethod
-    def get_config_filenames(cls) -> list[str]:
-        return []  # no extra configs.
-
-    @classmethod
-    def from_config(cls, config: dict[str, Any]) -> GGUFConfig:
-        return cls()
-
-    @classmethod
-    def override_quantization_method(
-        cls, hf_quant_cfg: dict[str, Any], user_quant: str | None, hf_config=None
-    ) -> QuantizationMethods | None:
-        if user_quant == "gguf":
-            return "gguf"
-        return None
-
-    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
-        if isinstance(layer, LinearBase):
-            if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
-                return UnquantizedLinearMethod()
-            return GGUFLinearMethod(self)
-        return None
-
-
-class GGUFLinearMethod(LinearMethodBase):
-    """Linear method for GGUF (inlined from upstream vllm)."""
-
-    def __init__(self, quant_config: GGUFConfig):
-        self.quant_config = quant_config
-
-    def create_weights(
-        self,
-        layer: torch.nn.Module,
-        input_size_per_partition: int,
-        output_partition_sizes: list[int],
-        input_size: int,
-        output_size: int,
-        params_dtype: torch.dtype,
-        **extra_weight_attrs,
-    ):
-        self.params_dtype = params_dtype
-        output_size_per_partition = sum(output_partition_sizes)
-
-        tensor_shape = (output_size_per_partition, input_size_per_partition)
-        qweight = _GGUFUninitializedParameter(requires_grad=False)
-        set_weight_attrs(
-            qweight,
-            {
-                "input_dim": 1,
-                "output_dim": 0,
-                "tensor_shape": tensor_shape,
-                "is_gguf_weight": True,
-                "data_container": [],
-                "shard_id": [],
-                "shard_id_map": {},
-            },
-        )
-        set_weight_attrs(qweight, extra_weight_attrs)
-        layer.register_parameter("qweight", qweight)
-
-        qweight_type = Parameter(
-            torch.empty(len(output_partition_sizes), dtype=torch.uint8),
-            requires_grad=False,
-        )
-        set_weight_attrs(
-            qweight_type,
-            {
-                "is_gguf_weight_type": True,
-                "weight_type": 0,
-                "shard_weight_type": {},
-                "ignore_warning": True,
-            },
-        )
-        set_weight_attrs(qweight_type, extra_weight_attrs)
-        layer.register_parameter("qweight_type", qweight_type)
-
-    def process_weights_after_loading(self, layer: torch.nn.Module):
-        qweight_type = layer.qweight_type.weight_type
-        if not (qweight_type in UNQUANTIZED_TYPES or qweight_type in DEQUANT_TYPES):
-            qweight_type = WeightType(qweight_type)
-            raise ValueError(f"Unsupported GGUF quantization type {qweight_type} in layer {layer}.")
-        self._create_padded_weight_param(layer)
-
-    def _create_padded_weight_param(self, layer: torch.nn.Module):
-        """Create padded weight parameter for GGUF MergedLinear layer."""
-        qweight = layer.qweight
-        shard_id_map = qweight.shard_id_map
-        shard_id = qweight.shard_id
-        if len(data_container := qweight.data_container) > 1:
-            dtype = {data.dtype for data in data_container}
-            assert len(dtype) == 1, ValueError(f"Data container has mixed dtypes: {dtype}")
-            dtype = next(iter(dtype))
-            padded_side = max(x.size(1) for x in data_container)
-            concat_side = sum(x.size(0) for x in data_container)
-            padded_data = torch.zeros((concat_side, padded_side), dtype=dtype, device=qweight.device)
-            shard_offset_map = dict[str, tuple[int, int, int]]()
-            for idx in shard_id:
-                id_in_container = shard_id_map[idx]
-                start = sum(x.size(0) for x in data_container[:id_in_container])
-                end = start + data_container[id_in_container].size(0)
-                size = data_container[id_in_container].size(1)
-                padded_data[start:end, :size] = data_container[id_in_container]
-                shard_offset_map[idx] = (start, end, size)
-            qweight.data_container.clear()
-            padded_param = Parameter(padded_data, requires_grad=False)
-            set_weight_attrs(padded_param, vars(qweight))
-            set_weight_attrs(padded_param, {"shard_offset_map": shard_offset_map})
-            layer.register_parameter("qweight", padded_param)
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        # Base apply: overridden by DiffusionGGUFLinearMethod.
-        raise NotImplementedError("Use DiffusionGGUFLinearMethod for diffusion GGUF")
-
-
-# ---------------------------------------------------------------------------
-# Diffusion-specific GGUF (dequant+GEMM, replaces fused kernel path)
-# ---------------------------------------------------------------------------
-
-
-def dequant_gemm_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor:
-    if qweight_type in UNQUANTIZED_TYPES:
-        return x @ qweight.T
-    block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type]
-    shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size)
-    weight = ops.ggml_dequantize(qweight, qweight_type, *shape, x.dtype)
-    return x @ weight.T
-
-
-class DiffusionGGUFLinearMethod(GGUFLinearMethod):
-    """GGUF linear method using dequant+GEMM for N-D diffusion tensors."""
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        shard_id = getattr(layer.qweight, "shard_id", None)
-
-        if shard_id:
-            shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id
-            qweight = layer.qweight
-            result = []
-            for idx in shard_id:
-                start, end, offset = layer.qweight.shard_offset_map[idx]
-                qweight_type = layer.qweight_type.shard_weight_type[idx]
-                result.append(
-                    dequant_gemm_gguf(
-                        x,
-                        qweight[start:end, :offset].contiguous(),
-                        qweight_type,
-                    )
-                )
-            out = torch.cat(result, axis=-1)
-        else:
-            qweight = layer.qweight
-            qweight_type = layer.qweight_type.weight_type
-            out = dequant_gemm_gguf(x, qweight, qweight_type)
-
-        if bias is not None:
-            out.add_(bias)
-        return out
-
-
-class DiffusionGGUFConfig(GGUFConfig):
-    """GGUF config that carries gguf_model path and uses dequant+GEMM."""
-
-    def __init__(
-        self,
-        gguf_model: str | None = None,
-        unquantized_modules: list[str] | None = None,
-    ) -> None:
-        super().__init__(unquantized_modules=unquantized_modules or [])
-        self.gguf_model = gguf_model
-
-    def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None:
-        if isinstance(layer, LinearBase):
-            if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping):
-                return UnquantizedLinearMethod()
-            return DiffusionGGUFLinearMethod(self)
-        return None
diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py
index 6a02d5edb..a6b29ea79 100644
--- a/vllm_omni/quantization/mxfp4_config.py
+++ b/vllm_omni/quantization/mxfp4_config.py
@@ -7,6 +7,8 @@ Architecture mirrors mxfp8_config.py:
   MXFPLinearMethodBase                    – platform-agnostic skeleton (imported from mxfp8_config)
     NPUMxfp4LinearMethod                  – NPU single-scale offline (W4A4 MXFP4)
       NPUMxfp4OnlineLinearMethod          – NPU single-scale online (BF16 → FP4)
+    ROCmMxfp4LinearMethod                 – ROCm base class (AITER quant + shuffle; online-only)
+      ROCmMxfp4OnlineLinearMethod         – ROCm online (BF16 → FP4 via AITER)
     NPUMxfp4DualScaleLinearMethod         – NPU dual-scale offline (W4A4 MXFP4 DualScale)
       NPUMxfp4DualScaleOnlineLinearMethod – NPU dual-scale online (BF16 → FP4)
 
@@ -153,9 +155,16 @@ class DiffusionMXFP4Config(QuantizationConfig):
                 if self.is_checkpoint_mxfp4_serialized:
                     return NPUMxfp4LinearMethod(self)
                 return NPUMxfp4OnlineLinearMethod(self)
+            if current_omni_platform.is_rocm():
+                gcn_arch = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).gcnArchName
+                if "gfx950" not in gcn_arch:
+                    raise NotImplementedError(f"MXFP4 on ROCm requires gfx950 (MI355X). Detected: {gcn_arch}")
+                if self.is_checkpoint_mxfp4_serialized:
+                    raise NotImplementedError("Pre-quantized MXFP4 checkpoints are not yet supported on ROCm.")
+                return ROCmMxfp4OnlineLinearMethod(self)
             raise NotImplementedError(
                 "DiffusionMXFP4Config (W4A4 MXFP4) is currently only supported "
-                "on NPU (Ascend) platforms. CUDA support is not yet implemented."
+                "on NPU (Ascend) and ROCm (AMD, gfx950) platforms."
             )
         return None
 
@@ -332,6 +341,151 @@ class NPUMxfp4OnlineLinearMethod(_LazyWeightMixin, NPUMxfp4LinearMethod):
         layer._already_called_process_weights_after_loading = True
 
 
+# ---------------------------------------------------------------------------
+# ROCm MXFP4 base method (AITER) — used as base class for the online subclass.
+# Online-only for now; offline (pre-quantized checkpoint) not yet supported.
+# ---------------------------------------------------------------------------
+
+
+def _register_rocm_mxfp4_op() -> None:
+    """Register the vllm_omni::rocm_mxfp4_gemm custom op for torch.compile.
+
+    Wraps activation quantization + GEMM in a single opaque op so that
+    torch.compile/inductor doesn't try to trace through AITER internals.
+    """
+    import aiter
+
+    @torch.library.custom_op("vllm_omni::rocm_mxfp4_gemm", mutates_args=())
+    def _rocm_mxfp4_gemm(
+        a: torch.Tensor,
+        w_quant: torch.Tensor,
+        w_scale: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        quant_func = aiter.get_hip_quant(aiter.QuantType.per_1x32)
+        a_quant, a_scale = quant_func(a, shuffle=True)
+        return aiter.gemm_a4w4(
+            a_quant,
+            w_quant,
+            a_scale,
+            w_scale,
+            bpreshuffle=True,
+            bias=bias,
+        )
+
+    @_rocm_mxfp4_gemm.register_fake
+    def _rocm_mxfp4_gemm_fake(
+        a: torch.Tensor,
+        w_quant: torch.Tensor,
+        w_scale: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        M, _ = a.shape
+        N, _ = w_quant.shape
+        return torch.empty(M, N, dtype=a.dtype, device=a.device)
+
+
+class ROCmMxfp4LinearMethod(MXFPLinearMethodBase):
+    """ROCm W4A4 MXFP4 linear method using AITER (gemm_a4w4).
+
+    Weight buffers after process_weights_after_loading:
+      weight_shuffle : FP4 quantized + shuffled via shuffle_weight(layout=(16,16))
+      weight_scale   : per-group-of-32 scales from AITER per_1x32
+
+    Forward path:
+      _quantize_activation: pass-through (activation quant is inside the custom op)
+      _quant_matmul: torch.ops.vllm_omni.rocm_mxfp4_gemm wraps activation quant
+                     then gemm_a4w4 in one custom op
+    """
+
+    def __init__(self, quant_config: DiffusionMXFP4Config) -> None:
+        self.quant_config = quant_config
+        self.out_dtype = torch.get_default_dtype()
+        if not hasattr(torch.ops.vllm_omni, "rocm_mxfp4_gemm"):
+            _register_rocm_mxfp4_op()
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        import aiter
+        from aiter.ops.shuffle import shuffle_weight
+
+        quant_func = aiter.get_hip_quant(aiter.QuantType.per_1x32)
+        weight_quant, weight_scale = quant_func(layer.weight.data, shuffle=True)
+        weight_shuffled = shuffle_weight(weight_quant, layout=(16, 16))
+
+        # Store quantized tensors as non-parameter buffers; delete original weight.
+        layer.register_buffer("weight_shuffle", weight_shuffled, persistent=True)
+        layer.register_buffer("weight_scale", weight_scale, persistent=True)
+
+        # Remove the original weight parameter to free memory.
+        if hasattr(layer, "weight") and isinstance(layer.weight, torch.nn.Parameter):
+            delattr(layer, "weight")
+            layer.register_parameter("weight", None)
+
+        layer._already_called_process_weights_after_loading = True
+
+    def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # Activation quantization is inside the custom op called by _quant_matmul,
+        # so pass through the raw activation here.
+        return x, None
+
+    def _quant_matmul(
+        self,
+        x_q: torch.Tensor,
+        x_scale: torch.Tensor,
+        layer: torch.nn.Module,
+        bias: torch.Tensor | None,
+        ori_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        output = torch.ops.vllm_omni.rocm_mxfp4_gemm(
+            x_q,
+            layer.weight_shuffle,
+            layer.weight_scale,
+            bias,
+        )
+        if output.dtype != ori_dtype:
+            output = output.to(ori_dtype)
+        return output
+
+
+# ---------------------------------------------------------------------------
+# ROCm MXFP4 online method (BF16 checkpoint → quantize at load time)
+# ---------------------------------------------------------------------------
+
+
+class ROCmMxfp4OnlineLinearMethod(_LazyWeightMixin, ROCmMxfp4LinearMethod):
+    """ROCm W4A4 MXFP4 online linear method using AITER.
+
+    MRO: ROCmMxfp4OnlineLinearMethod → _LazyWeightMixin → ROCmMxfp4LinearMethod
+         → MXFPLinearMethodBase → LinearMethodBase
+
+      create_weights  : _LazyWeightMixin          (meta device + patched loader)
+      process_weights : ROCmMxfp4OnlineLinearMethod  (meta → materialize, then AITER quant + shuffle)
+      apply / ops     : ROCmMxfp4LinearMethod / MXFPLinearMethodBase (shared)
+    """
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(layer, "_already_called_process_weights_after_loading", False):
+            return
+
+        # Materialise from meta device if needed (same pattern as NPU online).
+        if layer.weight is not None and layer.weight.device == torch.device("meta"):
+            weight = ModelWeightParameter(
+                data=torch.empty_like(layer.weight, device=layer._load_device),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=layer.weight.weight_loader,
+            )
+            _copy_missing_attrs(layer.weight, weight)
+            layer.register_parameter("weight", weight)
+            initialize_single_dummy_weight(layer.weight)
+
+        # Delegate to the base class which does AITER quant + shuffle.
+        ROCmMxfp4LinearMethod.process_weights_after_loading(self, layer)
+
+
 # ---------------------------------------------------------------------------
 # NPU MXFP4 dual-scale offline method (W4A4_MXFP4_DUALSCALE checkpoint)
 # ---------------------------------------------------------------------------
diff --git a/vllm_omni/quantization/quack_fp8.py b/vllm_omni/quantization/quack_fp8.py
index 8c7251e43..865958386 100644
--- a/vllm_omni/quantization/quack_fp8.py
+++ b/vllm_omni/quantization/quack_fp8.py
@@ -94,7 +94,7 @@ def quack_scaled_fp8_mm(
         return None
     out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=out_dtype)
     alpha = scale_a.reshape(1).float() * scale_b.reshape(1).float()
-    gemm.gemm_tuned(a, b, out, bias=bias, alpha=alpha)
+    gemm.gemm(a, b, out=out, bias=bias, alpha=alpha, tuned=True)
     return out
 
 
diff --git a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py
index 3d4f68d23..77b191738 100644
--- a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py
+++ b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py
@@ -280,7 +280,6 @@ def _build_quantization_config(
     quantization: str | None,
     quantization_config_json: str | None,
     ignored_layers: str | None,
-    gguf_model: str | None,
 ) -> dict[str, Any] | None:
     explicit = _load_json_arg(quantization_config_json)
     if explicit is not None:
@@ -292,8 +291,6 @@ def _build_quantization_config(
     ignored = _split_csv(ignored_layers)
     if ignored:
         config["ignored_layers"] = ignored
-    if gguf_model is not None:
-        config["gguf_model"] = gguf_model
     return config
 
 
@@ -305,7 +302,6 @@ def _build_variant_config(args: argparse.Namespace, variant: str) -> VariantConf
         quantization=quantization,
         quantization_config_json=getattr(args, f"{variant}_quantization_config_json"),
         ignored_layers=ignored_layers,
-        gguf_model=getattr(args, f"{variant}_gguf_model"),
     )
     return VariantConfig(
         label=variant,
@@ -554,8 +550,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--reference-ignored-layers", help="Comma-separated reference ignored layer patterns.")
     parser.add_argument("--candidate-ignored-layers", help="Comma-separated candidate ignored layer patterns.")
     parser.add_argument("--ignored-layers", help="Alias for --candidate-ignored-layers.")
-    parser.add_argument("--reference-gguf-model", help="Reference GGUF file or HF reference when using gguf.")
-    parser.add_argument("--candidate-gguf-model", help="Candidate GGUF file or HF reference when using gguf.")
     parser.add_argument("--reference-load-format", default="default", help="Reference diffusion load format.")
     parser.add_argument("--candidate-load-format", default="default", help="Candidate diffusion load format.")
     parser.add_argument("--prompt", default="a cup of coffee on the table")
diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py
index 0ca0c2c3a..189cc00aa 100644
--- a/vllm_omni/utils/mm_outputs.py
+++ b/vllm_omni/utils/mm_outputs.py
@@ -106,7 +106,13 @@ def _to_cpu(value):
 
 
 def to_payload_element(
-    element: object, idx: int, start: int, end: int, pass_lists_through: bool = False, seq_len: int | None = None
+    element: object,
+    idx: int,
+    start: int,
+    end: int,
+    pass_lists_through: bool = False,
+    seq_len: int | None = None,
+    scheduled_seq_len: int | None = None,
 ):
     """Build an mm payload element corresponding to one request index
     from an element containing 0 or more CPU tensors.
@@ -120,21 +126,39 @@ def to_payload_element(
             passthrough data; this should be False in normal cases, but True
             if we need to avoid splitting nonempty lists prior to calling
             postprocess, which is the case for prefix cache.
-        seq_len: Optional sequence length (i.e., dim 0 of hidden states).
-            When set, a tensor whose first dimension equals seq_len is
-            sliced per request. The prefix cache passthrough also passes
-            the total scheduled token count here so 1D (seq_len,) metadata
-            that is intentionally not cached is still split per request.
+        seq_len: Optional hidden-aligned batch length (``hidden_states.shape[0]``).
+            When a tensor's first dimension equals this value, it is sliced
+            per request using ``start:end``.
+        scheduled_seq_len: Optional scheduler-aligned batch length
+            (``scheduler_output.total_num_scheduled_tokens``). Some full-payload
+            mm tensors (e.g. batched ``codes.audio`` with tail-only hidden states)
+            are laid out by scheduled tokens instead of the hidden tail shape.
+            When omitted, ``seq_len`` is reused for backward compatibility.
     """
-    # Cached per-token tensors are merged elsewhere; here a first dim
-    # equal to seq_len means a per-request slice is required.
-    if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len:
+    if scheduled_seq_len is None:
+        scheduled_seq_len = seq_len
+
+    # Cached per-token tensors are merged elsewhere; here a first dim equal to
+    # either the hidden-aligned or scheduler-aligned batch length means a
+    # per-request slice is required.
+    if isinstance(element, torch.Tensor) and (
+        (seq_len is not None and element.shape[0] == seq_len)
+        or (scheduled_seq_len is not None and element.shape[0] == scheduled_seq_len)
+    ):
         return element[start:end].contiguous()
     # Every other case is shared between prefix cache (passthrough data)
     # and running a model without prefix caching.
     elif isinstance(element, dict):
         return {
-            sk: to_payload_element(sv, idx, start, end, pass_lists_through=pass_lists_through, seq_len=seq_len)
+            sk: to_payload_element(
+                sv,
+                idx,
+                start,
+                end,
+                pass_lists_through=pass_lists_through,
+                seq_len=seq_len,
+                scheduled_seq_len=scheduled_seq_len,
+            )
             for sk, sv in element.items()
         }
     elif isinstance(element, list):
diff --git a/vllm_omni/utils/tracking_parser.py b/vllm_omni/utils/tracking_parser.py
index c6dc5fee7..64849a5aa 100644
--- a/vllm_omni/utils/tracking_parser.py
+++ b/vllm_omni/utils/tracking_parser.py
@@ -9,6 +9,29 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
 UNSET = object()
 
 
+def build_shadow_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Build kwargs for the shadow argument with an ``UNSET`` default.
+
+    Actions that mutate their default in place (append/extend/append_const/
+    count) would crash on the bare ``UNSET`` sentinel, so they are remapped to
+    an equivalent store-style action; the shadow value only needs to flip away
+    from ``UNSET`` when the arg is passed explicitly.
+    """
+    shadow_kwargs = {**kwargs, "default": UNSET}
+    action = kwargs.get("action")
+
+    if action in ("append", "extend"):
+        shadow_kwargs["action"] = "store"
+
+    elif action in ("append_const", "count"):
+        shadow_kwargs["action"] = "store_const"
+
+        if action == "count":
+            shadow_kwargs["const"] = True
+
+    return shadow_kwargs
+
+
 class TrackingNamespace(argparse.Namespace):
     """Proxy that wraps an argparse namespace with explicit keys, which
     can be filtered down to a dict containing only explicitly passed values.
@@ -58,7 +81,7 @@ class TrackingGroup:
     def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action:
         """Add an argument to the real group and to the shadow group."""
         action = self._real.add_argument(*args, **kwargs)
-        default_kwargs = {**kwargs, "default": UNSET}
+        default_kwargs = build_shadow_kwargs(kwargs)
         self._shadow.add_argument(*args, **default_kwargs)
         return action
 
@@ -115,7 +138,7 @@ class TrackingArgumentParser(FlexibleArgumentParser):
     def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action:
         """Add an arg to the parser & the shadow, where the latter has UNSET for the default."""
         action = super().add_argument(*args, **kwargs)
-        shadow_kwargs = {**kwargs, "default": UNSET}
+        shadow_kwargs = build_shadow_kwargs(kwargs)
         self._shadow.add_argument(*args, **shadow_kwargs)
         return action
 
diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py
index a04a83657..6cd2b3566 100644
--- a/vllm_omni/worker/gpu_ar_model_runner.py
+++ b/vllm_omni/worker/gpu_ar_model_runner.py
@@ -8,7 +8,7 @@ from __future__ import annotations
 
 import gc
 import threading
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Sequence
 from contextlib import nullcontext
 from copy import copy
 from dataclasses import replace
@@ -22,6 +22,7 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_
 from vllm.distributed.parallel_state import get_pp_group, get_tp_group
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
+from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
 from vllm.v1.spec_decode.dflash import DFlashProposer
@@ -45,6 +46,8 @@ from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
+from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest
+from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
 
 logger = init_logger(__name__)
 
@@ -156,6 +159,11 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         async_output_copy_stream = kwargs.pop("async_output_copy_stream")
         vocab_size = kwargs.pop("vocab_size")
         routed_experts = kwargs.pop("routed_experts", None)
+        # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault
+        # for EP all2all fault tolerance (PR #43637). Omni doesn't use this
+        # feature but must consume the kwarg to prevent TypeError from stray
+        # kwargs and initialize the attribute so super().get_output() works.
+        kwargs.pop("check_ep_fault", False)
         if kwargs:
             raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}")
 
@@ -167,6 +175,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
         self.vocab_size = vocab_size
         self._logprobs_tensors = logprobs_tensors
         self._routed_experts = routed_experts
+        self._has_fault: torch.Tensor | None = None
 
         default_stream = torch.cuda.current_stream()
         with torch.cuda.stream(async_output_copy_stream):
@@ -217,6 +226,12 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput):
             if background_exception is not None:
                 raise background_exception
         self._build_model_runner_output_once()
+        # Upstream AsyncGPUModelRunnerOutput.get_output() accesses
+        # self._has_fault for EP all2all fault tolerance (PR #43637).
+        # Ensure the attribute exists even when __init__ was bypassed
+        # (e.g. unit tests using object.__new__).
+        if not hasattr(self, "_has_fault"):
+            self._has_fault = None
         with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"):
             return super().get_output()
 
@@ -587,6 +602,93 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         model = getattr(self, "model", None)
         return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True))
 
+    def _get_runner_assisted_full_attention_metadata_request(
+        self,
+        *,
+        req_ids: Sequence[str],
+        num_reqs: int,
+        num_reqs_padded: int,
+        num_scheduled_tokens_np: np.ndarray,
+        num_computed_tokens_cpu: torch.Tensor | np.ndarray,
+        max_num_scheduled_tokens: int,
+    ) -> RunnerAssistedFullAttentionMetadataRequest | None:
+        # Models without this hook keep the normal runner path. VoxCPM2 uses
+        # it to ask the runner for padded FULL attention metadata while keeping
+        # graph policy and lifecycle in the model layer.
+        hook = getattr(self.model, "get_runner_assisted_full_attention_metadata_request", None)
+        if not callable(hook):
+            return None
+        num_computed_tokens = num_computed_tokens_cpu
+        if hasattr(num_computed_tokens, "numpy"):
+            num_computed_tokens = num_computed_tokens.numpy()
+        request = hook(
+            req_ids=req_ids,
+            num_reqs=num_reqs,
+            num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
+            num_computed_tokens=num_computed_tokens,
+            max_num_scheduled_tokens=max_num_scheduled_tokens,
+        )
+        if request is None:
+            return None
+        if not isinstance(request, RunnerAssistedFullAttentionMetadataRequest):
+            raise TypeError(
+                "runner-assisted full attention metadata hook must return "
+                "RunnerAssistedFullAttentionMetadataRequest or None, got "
+                f"{type(request).__name__}"
+            )
+        padded_num_reqs = max(
+            num_reqs_padded,
+            min(int(request.num_reqs_padded), self.scheduler_config.max_num_seqs),
+        )
+        return RunnerAssistedFullAttentionMetadataRequest(
+            num_reqs_padded=padded_num_reqs,
+            for_cudagraph_capture=bool(request.for_cudagraph_capture),
+        )
+
+    def _refresh_runner_assisted_full_attention_metadata_buffers(
+        self,
+        *,
+        num_reqs: int,
+        num_reqs_padded: int,
+        num_scheduled_tokens_np: np.ndarray,
+    ) -> None:
+        num_computed_tokens = self.input_batch.num_computed_tokens_cpu[:num_reqs]
+        if hasattr(num_computed_tokens, "numpy"):
+            num_computed_tokens = num_computed_tokens.numpy()
+        np.add(
+            num_computed_tokens,
+            num_scheduled_tokens_np,
+            out=self.optimistic_seq_lens_cpu[:num_reqs].numpy(),
+        )
+        self.optimistic_seq_lens_cpu[num_reqs:num_reqs_padded].fill_(0)
+        self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True)
+
+        cum_num_tokens = self._get_cumsum_and_arange(
+            num_scheduled_tokens_np,
+            self.query_pos.np,
+        )
+        self.query_start_loc.np[0] = 0
+        self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+        self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1])
+        self.query_start_loc.copy_to_gpu()
+
+        self.input_batch.block_table.commit_block_table(num_reqs_padded)
+
+    def _set_runner_assisted_full_attention_metadata_context(
+        self,
+        *,
+        enabled: bool,
+        num_reqs: int = 0,
+    ) -> bool:
+        hook = getattr(self.model, "set_runner_assisted_full_attention_metadata_context", None)
+        if not callable(hook):
+            return False
+        hook(
+            enabled=enabled,
+            num_reqs=num_reqs,
+        )
+        return True
+
     def _deferred_prefix_cache_mm_keys(self) -> set[str]:
         """Model-declared multimodal keys whose prefix-cache writes are deferred."""
         model = getattr(self, "model", None)
@@ -613,12 +715,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         if self.omni_prefix_cache is not None and is_last_pp_rank:
             # If this happens, it generally means the model is not following the correct
             # interface yet and is therefore currently not compatible with prefix cache.
-            if multimodal_outputs is not None and not isinstance(multimodal_outputs, Mapping):
-                logger.warning_once(
-                    "prefix caching expects mm outputs to be a dict, but got %s",
-                    type(multimodal_outputs),
-                )
-
             hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
             # FIX: The .cpu attribute of slot_mapping is stale (not updated by the Triton
             # _compute_slot_mapping_kernel which only writes to .gpu). We must use .gpu and
@@ -750,7 +846,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         end: int,
         audio_sparse_output: bool,
         sparse_mm_index: dict[str, int],
-        seq_len: int,
+        hidden_seq_len: int,
+        scheduled_seq_len: int,
     ) -> dict[str, object]:
         if combined_multimodal_outputs:
             return self._build_combined_prefix_cache_mm_payload(
@@ -787,7 +884,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 start=start,
                 end=end,
                 pass_lists_through=False,
-                seq_len=seq_len,
+                seq_len=hidden_seq_len,
+                scheduled_seq_len=scheduled_seq_len,
             )
         return mm_payload
 
@@ -805,7 +903,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         mm_cpu: dict[str, object] | None,
         audio_sparse_output: bool,
         sparse_mm_index: dict[str, int],
-        seq_len: int,
+        hidden_seq_len: int,
+        scheduled_seq_len: int,
     ) -> dict[str, object]:
         payload: dict[str, object] = {}
         if not audio_sparse_output:
@@ -831,7 +930,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             end=end,
             audio_sparse_output=audio_sparse_output,
             sparse_mm_index=sparse_mm_index,
-            seq_len=seq_len,
+            hidden_seq_len=hidden_seq_len,
+            scheduled_seq_len=scheduled_seq_len,
         )
         payload.update(mm_payload)
         return payload
@@ -1027,6 +1127,26 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             )
             num_tokens_padded = batch_desc.num_tokens
             num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
+            num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu[:num_reqs]
+
+            runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request(
+                req_ids=req_ids[:num_reqs],
+                num_reqs=num_reqs,
+                num_reqs_padded=num_reqs_padded,
+                num_scheduled_tokens_np=num_scheduled_tokens_np,
+                num_computed_tokens_cpu=num_computed_tokens_cpu,
+                max_num_scheduled_tokens=max_num_scheduled_tokens,
+            )
+            runner_assisted_full_attn_capture = False
+            if runner_assisted_full_attn_request is not None:
+                num_reqs_padded = runner_assisted_full_attn_request.num_reqs_padded
+                runner_assisted_full_attn_capture = runner_assisted_full_attn_request.for_cudagraph_capture
+                num_tokens_padded = max(
+                    num_tokens_padded,
+                    num_reqs_padded * max_num_scheduled_tokens,
+                )
+            runner_assisted_full_attn = runner_assisted_full_attn_request is not None
+
             ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
                 should_ubatch,
                 num_scheduled_tokens_np,
@@ -1034,8 +1154,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 num_reqs_padded,
                 self.parallel_config.num_ubatches,
             )
-
-            pad_attn = cudagraph_mode == CUDAGraphMode.FULL
+            pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn
 
             use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
             ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
@@ -1057,11 +1176,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 num_tokens_unpadded=num_tokens_unpadded,
                 ubatch_slices=ubatch_slices_padded,
             )
+            if runner_assisted_full_attn:
+                self._refresh_runner_assisted_full_attention_metadata_buffers(
+                    num_reqs=num_reqs,
+                    num_reqs_padded=num_reqs_padded,
+                    num_scheduled_tokens_np=num_scheduled_tokens_np,
+                )
 
+            attn_num_reqs = num_reqs_padded if runner_assisted_full_attn else num_reqs
             attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata(
                 num_tokens=num_tokens_unpadded,
                 num_tokens_padded=num_tokens_padded if pad_attn else None,
-                num_reqs=num_reqs,
+                num_reqs=attn_num_reqs,
                 num_reqs_padded=num_reqs_padded if pad_attn else None,
                 max_query_len=max_num_scheduled_tokens,
                 ubatch_slices=ubatch_slices_attn,
@@ -1070,6 +1196,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
                 cascade_attn_prefix_lens=cascade_attn_prefix_lens,
                 slot_mappings=slot_mappings_by_group,
+                for_cudagraph_capture=runner_assisted_full_attn_capture,
             )
             self._maybe_attach_attention_metadata_extensions(
                 attn_metadata=attn_metadata,
@@ -1078,6 +1205,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 max_query_len=max_num_scheduled_tokens,
                 pad_attn=pad_attn,
                 num_scheduled_tokens_np=num_scheduled_tokens_np,
+                for_cudagraph_capture=runner_assisted_full_attn_capture,
             )
 
             (
@@ -1091,14 +1219,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
 
         # Let the model adjust inputs before forward (e.g. restore input_ids
         # for multimodal position detection, fix decode position offsets).
-        if hasattr(self.model, "prepare_runner_inputs"):
-            input_ids, positions = self.model.prepare_runner_inputs(
+        prepare_runner_inputs = getattr(self.model, "prepare_runner_inputs", None)
+        if callable(prepare_runner_inputs):
+            input_ids, positions = prepare_runner_inputs(
                 input_ids=input_ids,
                 positions=positions,
                 inputs_embeds=inputs_embeds,
                 req_ids=req_ids[:num_reqs],
-                num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)],
-                num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)],
+                num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs],
+                num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs],
                 input_ids_buffer=self.input_ids.gpu[:num_tokens_padded],
             )
 
@@ -1107,46 +1236,60 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         # with CUDA graph capture.
         if self.calculate_kv_scales:
             cudagraph_mode = CUDAGraphMode.NONE
+            runner_assisted_full_attn = False
             # Mark KV scales as calculated after the first forward pass
             self.calculate_kv_scales = False
 
+        runner_assisted_context_enabled = False
+        if runner_assisted_full_attn:
+            runner_assisted_context_enabled = self._set_runner_assisted_full_attention_metadata_context(
+                enabled=True,
+                num_reqs=num_reqs,
+            )
+            runner_assisted_full_attn = runner_assisted_context_enabled
+
         # Run the model.
         # Use persistent buffers for CUDA graphs.
         # When spec decode is enabled, defer connector finalization
         # (wait_for_save + clear metadata) until after draft model runs.
         defer_kv_connector_finalize = self.speculative_config is not None
-        with (
-            nullcontext(),
-            set_forward_context(
-                attn_metadata,
-                self.vllm_config,
-                num_tokens=num_tokens_padded,
-                num_tokens_across_dp=num_tokens_across_dp,
-                cudagraph_runtime_mode=cudagraph_mode,
-                batch_descriptor=batch_desc,
-                ubatch_slices=ubatch_slices_padded,
-                slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
-            ),
-            record_function_or_nullcontext("gpu_model_runner: forward"),
-            self.maybe_get_kv_connector_output(
-                scheduler_output,
-                defer_finalize=defer_kv_connector_finalize,
-            ) as kv_connector_output,
-        ):
-            model_output = self._model_forward(
-                input_ids=input_ids,
-                positions=positions,
-                intermediate_tensors=intermediate_tensors,
-                inputs_embeds=inputs_embeds,
-                **model_kwargs,
-                sampling_metadata=self.input_batch.sampling_metadata,
-                logits_index=logits_indices,
-                sampler=self.sampler,
-            )
+        try:
+            with (
+                nullcontext(),
+                set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens_padded,
+                    num_tokens_across_dp=num_tokens_across_dp,
+                    cudagraph_runtime_mode=(CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode),
+                    batch_descriptor=batch_desc,
+                    ubatch_slices=ubatch_slices_padded,
+                    slot_mapping=slot_mappings,  # OMNI: required for KV cache operations
+                ),
+                record_function_or_nullcontext("gpu_model_runner: forward"),
+                self.maybe_get_kv_connector_output(
+                    scheduler_output,
+                    defer_finalize=defer_kv_connector_finalize,
+                ) as kv_connector_output,
+            ):
+                model_output = self._model_forward(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                    **model_kwargs,
+                    sampling_metadata=self.input_batch.sampling_metadata,
+                    logits_index=logits_indices,
+                    sampler=self.sampler,
+                )
 
-            # [Omni] Map pending ropes metadata to req_ids.
-            if hasattr(self.model, "flush_pending_metadata"):
-                self.model.flush_pending_metadata(list(req_ids))
+                # [Omni] Map pending ropes metadata to req_ids.
+                flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None)
+                if callable(flush_pending_metadata):
+                    flush_pending_metadata(req_ids[:num_reqs])
+        finally:
+            if runner_assisted_context_enabled:
+                self._set_runner_assisted_full_attention_metadata_context(enabled=False)
 
         with record_function_or_nullcontext("gpu_model_runner: postprocess"):
             if self.use_aux_hidden_state_outputs:
@@ -1172,7 +1315,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 # prefix caching but the downstream pooler payload path still
                 # needs a CPU hidden-states view. Materialize it synchronously
                 # in that case; the legacy behavior is preserved.
-                if hs_for_cache is None:
+                if hs_for_cache is None and self._model_omni_pooler_payload_include_hidden():
                     hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous()
                 slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu
                 self.omni_prefix_cache.schedule_async_write(
@@ -1459,7 +1602,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     multimodal_outputs=multimodal_outputs,
                 ),
                 copy_stream=self._get_or_create_omni_payload_copy_stream(),
-                pin_memory=bool(getattr(self, "pin_memory", False)),
+                # NOTE: vLLM v0.24.0's GPUModelRunner no longer exposes a
+                # ``self.pin_memory`` attribute (it uses a module-level
+                # ``PIN_MEMORY`` constant instead), so the old
+                # ``getattr(self, "pin_memory", False)`` silently fell back to
+                # False. That allocated the async D2H snapshot destination in
+                # *pageable* host memory, which turns ``copy_(non_blocking=True)``
+                # into a fully synchronous, stream-stalling copy (~240 ms/step
+                # on the 17.5k-token Thinker prefill). Resolve pinning from the
+                # platform helper so the copy is a true async cudaMemcpyAsync.
+                pin_memory=is_pin_memory_available(),
             )
 
         payload = async_payload_snapshot.payload
@@ -1535,7 +1687,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         ec_connector_output: Any,
         cudagraph_stats: Any,
         kv_extracted_req_ids: list[str] | None,
-        seq_len: int,
         num_scheduled_tokens_np: np.ndarray,
         query_start_loc_cpu: Any,
         postprocess_already_applied: bool = False,
@@ -1591,6 +1742,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         # The actual multimodal wire transport uses multimodal_outputs instead.
         pooler_output: list[dict[str, object]] | None = None
         if needs_pooler_payload:
+            hidden_seq_len = int(hidden_states.shape[0])
+            scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens)
             mm_cpu = None
             if self.omni_prefix_cache is not None:
                 (
@@ -1646,7 +1799,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                         mm_cpu=mm_cpu,
                         audio_sparse_output=audio_sparse_output,
                         sparse_mm_index=sparse_mm_index,
-                        seq_len=seq_len,
+                        hidden_seq_len=hidden_seq_len,
+                        scheduled_seq_len=scheduled_seq_len,
                     )
                     pooler_output.append(flatten_payload(payload))
 
@@ -1732,7 +1886,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             slot_mappings,  # OMNI: unpack slot_mappings for drafter
         ) = self.execute_model_state
         self.execute_model_state = None
-        seq_len = hidden_states.shape[0]
 
         # Apply structured output bitmasks if present.
         if grammar_output is not None:
@@ -1746,6 +1899,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 if self.input_batch.vocab_size > logits_vocab:
                     smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab)
 
+        # Drop min-tokens stop ids the head cannot emit (e.g. the text
+        # tokenizer EOS folded into all_stop_token_ids on a narrow codec
+        # talker head); they would index_put_ out of bounds (#4962).
+        if logits is not None:
+            sanitize_min_tokens_stop_ids(
+                self.input_batch.sampling_metadata.logitsprocs,
+                logits.shape[-1],
+            )
+
         with record_function_or_nullcontext("gpu_model_runner: sample"):
             sampler_output = self._sample(logits, spec_decode_metadata)
 
@@ -1899,7 +2061,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     ec_connector_output=ec_connector_output,
                     cudagraph_stats=cudagraph_stats,
                     kv_extracted_req_ids=kv_extracted_req_ids,
-                    seq_len=seq_len,
                     num_scheduled_tokens_np=num_scheduled_tokens_np,
                     query_start_loc_cpu=query_start_loc_cpu,
                     postprocess_already_applied=omni_postprocess_already_applied,
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index 172121c14..ec41900b1 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/vllm_omni/worker/gpu_generation_model_runner.py
@@ -153,6 +153,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             if self.model_config.async_chunk and num_scheduled_tokens:
                 self._update_request_states(scheduler_output)
             deferred_state_corrections_fn = self._update_states(scheduler_output)
+
+            # Notify stateful models of finished requests before the
+            # zero-token early return. Request cleanup is a model lifecycle
+            # concern and does not depend on the inter-stage transfer mode.
+            if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"):
+                self.model.on_requests_finished(scheduler_output.finished_req_ids)
+
             if not scheduler_output.total_num_scheduled_tokens:
                 return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
 
@@ -789,6 +796,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             if hasattr(self.model, "get_dummy_runtime_additional_information"):
                 runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs)
                 model_kwargs["runtime_additional_information"] = runtime_addi
+            # Generation-stage code2wav models split flat input_ids using this
+            # shared contract. Dummy/profile runs use padded synthetic inputs,
+            # so expose one synthetic segment covering the actual tensor length.
+            model_kwargs["seq_token_counts"] = [int(num_tokens_padded)]
 
             if self.uses_mrope:
                 positions = self.mrope_positions.gpu[:, :num_tokens_padded]
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index a283039f3..b816a36f7 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/vllm_omni/worker/gpu_model_runner.py
@@ -601,9 +601,6 @@ class OmniGPUModelRunner(GPUModelRunner):
             # Decode additional_information payloads (dictionary)
             try:
                 if getattr(new_req_data, "additional_information", None) is not None:
-                    logger.warning_once(
-                        "additional_information on request data is deprecated, use model_intermediate_buffer"
-                    )
                     info_dict = deserialize_additional_information(new_req_data.additional_information)
                     if info_dict:
                         self.model_intermediate_buffer[req_id] = info_dict
@@ -1471,15 +1468,9 @@ class OmniGPUModelRunner(GPUModelRunner):
         for new_req in scheduler_output.scheduled_new_reqs:
             payload_info = getattr(new_req, "additional_information", None)
             if isinstance(payload_info, dict):
-                logger.warning_once(
-                    "additional_information on request data is deprecated, use model_intermediate_buffer"
-                )
                 self._update_intermediate_buffer(new_req.req_id, payload_info)
 
         if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
-            logger.warning_once(
-                "additional_information on scheduled_cached_reqs is deprecated, use model_intermediate_buffer"
-            )
             cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})
             if isinstance(cached_infos, dict):
                 for req_id, req_infos in cached_infos.items():
@@ -1707,7 +1698,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self.text_step.gpu[dst].copy_(text_step)
 
                 for req_id_b, update_dict_b in zip(req_ids_b, updates, strict=True):
-                    self._merge_additional_information_update(req_id_b, update_dict_b)
+                    self._update_intermediate_buffer(req_id_b, update_dict_b)
 
                 decode_req_ids.extend(req_ids_b)
                 decode_start_offsets.extend(start_offsets_b)
@@ -1762,7 +1753,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                     decode_start_offsets.append(s)
 
                 # TODO(Peiqi): the merge stage could move out from the critical path
-                self._merge_additional_information_update(req_id, update_dict)
+                self._update_intermediate_buffer(req_id, update_dict)
 
                 # update the inputs_embeds and input_ids
                 seg_len = min(span_len, req_embeds.shape[0])
@@ -1821,10 +1812,38 @@ class OmniGPUModelRunner(GPUModelRunner):
         def _explicit_talker_seed(req_id: str) -> int | None:
             sampling_params = getattr(self.requests[req_id], "sampling_params", None)
             extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
-            seed = extra_args.get("qwen3_tts_request_seed") if isinstance(extra_args, dict) else None
+            seed = None
+            if isinstance(extra_args, dict):
+                seed = extra_args.get("tts_local_seed")
             return int(seed) if seed is not None else None
 
-        if decode_batch_size > 1 and any(_explicit_talker_seed(req_id) is not None for req_id in decode_req_ids):
+        def _row_generator(req_id: str) -> torch.Generator | None:
+            seed = _explicit_talker_seed(req_id)
+            if seed is None:
+                return None
+            cache = getattr(self, "_talker_mtp_generators", None)
+            if cache is None:
+                cache = {}
+                self._talker_mtp_generators = cache
+            generator = cache.get(req_id)
+            if generator is None or generator.device != req_input_ids.device:
+                generator = torch.Generator(device=req_input_ids.device)
+                generator.manual_seed(seed)
+                cache[req_id] = generator
+            return generator
+
+        row_generators = [_row_generator(req_id) for req_id in decode_req_ids]
+        cache = getattr(self, "_talker_mtp_generators", None)
+        if cache:
+            # Generators live as long as their request; drop finished ones.
+            for stale_id in [rid for rid in cache if rid not in self.requests]:
+                del cache[stale_id]
+
+        if (
+            decode_batch_size > 1
+            and any(generator is not None for generator in row_generators)
+            and not getattr(self.model, "talker_mtp_accepts_per_row_generators", False)
+        ):
             # A torch.Generator is a single stream. Using one generator for a
             # multi-row batch would make explicitly-seeded requests depend on
             # other rows in the same scheduler step, so keep that path scalar.
@@ -1847,28 +1866,22 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self.text_step.gpu[:decode_batch_size].copy_(saved_text)
             return
 
-        generator = None
-        if decode_req_ids:
-            first_req_id = decode_req_ids[0]
-            seed = _explicit_talker_seed(first_req_id)
-            if seed is not None:
-                generators = getattr(self, "_talker_mtp_generators", None)
-                if generators is None:
-                    generators = {}
-                    self._talker_mtp_generators = generators
-                generator = generators.get(first_req_id)
-                if generator is None or generator.device != req_input_ids.device:
-                    generator = torch.Generator(device=req_input_ids.device)
-                    generator.manual_seed(int(seed))
-                    generators[first_req_id] = generator
         talker_kwargs = {
             "do_sample": subtalker_params.get("do_sample"),
             "temperature": subtalker_params.get("temperature"),
             "top_k": subtalker_params.get("top_k"),
             "top_p": subtalker_params.get("top_p"),
         }
-        if generator is not None:
-            talker_kwargs["generator"] = generator
+        if decode_batch_size == 1:
+            if row_generators[0] is not None:
+                talker_kwargs["generator"] = row_generators[0]
+        elif any(generator is not None for generator in row_generators):
+            talker_kwargs["generators"] = row_generators
+        if getattr(self.model, "talker_mtp_accepts_req_infos", False):
+            talker_kwargs["req_ids"] = decode_req_ids
+            talker_kwargs["req_infos"] = [
+                self.model_intermediate_buffer.setdefault(req_id, {}) for req_id in decode_req_ids
+            ]
         with current_omni_platform.set_forward_context(
             None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
         ):
@@ -1888,8 +1901,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             start_offsets = [int(self.query_start_loc.cpu[id_to_index[req_id]]) for req_id in decode_req_ids]
         for idx, (req_id, start_offset) in enumerate(zip(decode_req_ids, start_offsets, strict=True)):
             inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1]
-            update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
-            self._merge_additional_information_update(req_id, update_dict)
+            if code_predictor_codes is not None:
+                update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
+                self._update_intermediate_buffer(req_id, update_dict)
 
     def _model_forward(
         self,
@@ -1963,10 +1977,6 @@ class OmniGPUModelRunner(GPUModelRunner):
         # Backward compatible: mirror to old setattr location
         setattr(req_state, "additional_information_cpu", existing)
 
-    def _merge_additional_information_update(self, req_id, upd):
-        logger.warning_once("_merge_additional_information_update is deprecated, use _update_intermediate_buffer")
-        return self._update_intermediate_buffer(req_id, upd)
-
     def _update_streaming_input_additional_info(self, req_id):
         # For streaming input prefill case only. Set num processed tokens = 0 for new segment input
         cached_additional_info = self.model_intermediate_buffer.get(req_id, {})
diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py
index c1421a231..b173f81ac 100644
--- a/vllm_omni/worker/omni_connector_model_runner_mixin.py
+++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py
@@ -919,7 +919,7 @@ class OmniConnectorModelRunnerMixin:
         if not (finished_req_ids & pending_req_ids):
             return
 
-        logger.info(
+        logger.debug(
             "[Stage-%s] flush_full_payload_outputs: finished_req_ids=%s, pending=%s",
             self._stage_id,
             finished_req_ids,
diff --git a/vllm_omni/worker/runner_assisted_metadata.py b/vllm_omni/worker/runner_assisted_metadata.py
new file mode 100644
index 000000000..a71180a25
--- /dev/null
+++ b/vllm_omni/worker/runner_assisted_metadata.py
@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import NamedTuple, Protocol
+
+
+class RunnerAssistedFullAttentionMetadataRequest(NamedTuple):
+    num_reqs_padded: int
+    for_cudagraph_capture: bool
+
+
+class RunnerAssistedAttentionMetadataProvider(Protocol):
+    """Optional model hook for runner-built full attention metadata."""
+
+    def get_runner_assisted_full_attention_metadata_request(
+        self,
+        *,
+        req_ids: Sequence[str],
+        num_reqs: int,
+        num_scheduled_tokens: Sequence[int],
+        num_computed_tokens: Sequence[int],
+        max_num_scheduled_tokens: int,
+    ) -> RunnerAssistedFullAttentionMetadataRequest | None: ...
+
+    def set_runner_assisted_full_attention_metadata_context(
+        self,
+        *,
+        enabled: bool,
+        num_reqs: int = 0,
+    ) -> None: ...
diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py
new file mode 100644
index 000000000..c38a9bbdc
--- /dev/null
+++ b/vllm_omni/worker/sampling_utils.py
@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sampling-state guards shared by the GPU and NPU AR model runners."""
+
+import torch
+from vllm.logger import init_logger
+from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsProcessor
+
+logger = init_logger(__name__)
+
+__all__ = ["sanitize_min_tokens_stop_ids"]
+
+
+def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: int) -> None:
+    """Drop stop ids the model head cannot emit from min-tokens masking state.
+
+    vLLM's input processor unconditionally folds the stage tokenizer's EOS id
+    into ``SamplingParams.all_stop_token_ids``. AR stages whose lm_head is
+    narrower than the tokenizer vocabulary (codec talkers such as Qwen3-TTS:
+    3072 logits vs text EOS 151645) then crash on any ``min_tokens >= 1``:
+    ``MinTokensLogitsProcessor.apply`` writes ``-inf`` at an out-of-range
+    index and ``index_put_`` triggers a CUDA device-side assert (#4962).
+
+    Out-of-range ids are unreachable for the head, so dropping them never
+    changes sampling or stopping behavior. The per-request stop-id set is
+    mutated in place (it is shared with the request's ``SamplingParams``),
+    so each request is sanitized at most once; the processor's device-side
+    mask slice is rebuilt only when an out-of-range id was found.
+    """
+    for proc in logitsprocs.non_argmax_invariant:
+        if not isinstance(proc, MinTokensLogitsProcessor):
+            continue
+        min_toks = getattr(proc, "min_toks", None)
+        if not min_toks:
+            continue
+        needs_rebuild = False
+        for _, _, stop_tok_ids in min_toks.values():
+            oob = [tok for tok in stop_tok_ids if tok >= logits_vocab]
+            if not oob:
+                continue
+            stop_tok_ids.difference_update(oob)
+            needs_rebuild = True
+            logger.warning_once(
+                "min_tokens: dropped stop token ids %s that exceed the logits vocabulary (%d); "
+                "the model head cannot emit them.",
+                str(sorted(oob)),
+                logits_vocab,
+            )
+        if needs_rebuild:
+            reqs: list[int] = []
+            tok_ids: list[int] = []
+            for index, (_, _, stop_tok_ids) in min_toks.items():
+                reqs.extend([index] * len(stop_tok_ids))
+                tok_ids.extend(stop_tok_ids)
+            proc.logits_slice = (
+                proc._device_tensor(reqs, torch.int32),
+                proc._device_tensor(tok_ids, torch.int32),
+            )