[Diffusion] Document DLO mmap model opt-in

Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
[Diffusion] Address rank-local mmap review feedback
2026-08-15 06:57:38 +08:00 · 2026-08-14 17:12:05 +08:00 · 2026-08-14 14:38:06 +08:00
211 changed files with 3230 additions and 16112 deletions
@@ -1,103 +1,153 @@
 #!/bin/bash
+# vllm-omni customized version
+# Based on: vllm/.buildkite/scripts/hardware_ci/run-amd-test.sh
+# Last synced: 2025-12-15
+# Modifications: docker image name for vllm-omni

-# Run vLLM-Omni ROCm tests directly in the MI300 Kubernetes pod. The pod's
-# container image is selected by test-template-amd-omni.j2; MI300 has no DinD.
-set -euo pipefail
+# This script runs test inside the corresponding ROCm docker container.
+set -o pipefail

-: "${PYTHONFAULTHANDLER:=1}"
-: "${HF_HOME:=/home/buildkite-agent/huggingface}"
-: "${HF_HUB_DOWNLOAD_TIMEOUT:=300}"
-: "${HF_HUB_ETAG_TIMEOUT:=60}"
-: "${MIOPEN_DEBUG_CONV_DIRECT:=0}"
-: "${MIOPEN_DEBUG_CONV_GEMM:=0}"
-: "${VLLM_ROCM_USE_AITER:=0}"
-export PYTHONFAULTHANDLER HF_HOME HF_HUB_DOWNLOAD_TIMEOUT HF_HUB_ETAG_TIMEOUT
-export MIOPEN_DEBUG_CONV_DIRECT MIOPEN_DEBUG_CONV_GEMM VLLM_ROCM_USE_AITER
-export PYTORCH_ROCM_ARCH=""
-export PYTHONPATH="${PYTHONPATH:-..}"
-
-if [[ "${VLLM_CI_DOCKER_DISABLED:-0}" != "1" ]]; then
-    echo "Error: MI300 CI must run natively with Docker disabled." >&2
-    exit 1
-fi
-
-if [[ -n "${TEST_COMMANDS:-}" ]]; then
-    commands="${TEST_COMMANDS}"
-else
-    commands="$*"
-fi
-if [[ -z "${commands}" ]]; then
-    echo "Error: No test commands provided." >&2
-    exit 1
-fi
-
-echo "--- Native in-pod ROCm CI"
-
-job_id="${BUILDKITE_JOB_ID:-${BUILDKITE_PARALLEL_JOB:-local}}"
-job_id="${job_id//[^A-Za-z0-9_.-]/_}"
-native_cache_root="/tmp/vllm-omni-native-${job_id}"
-TMPDIR="${native_cache_root}/tmp"
-TORCHINDUCTOR_CACHE_DIR="${native_cache_root}/torchinductor"
-TRITON_CACHE_DIR="${native_cache_root}/triton"
-VLLM_CACHE_ROOT="${native_cache_root}/vllm"
-XDG_CACHE_HOME="${native_cache_root}/xdg"
-HF_DATASETS_CACHE="${native_cache_root}/huggingface/datasets"
-export TMPDIR TORCHINDUCTOR_CACHE_DIR TRITON_CACHE_DIR VLLM_CACHE_ROOT
-export XDG_CACHE_HOME HF_DATASETS_CACHE
-mkdir -p "${TMPDIR}" "${TORCHINDUCTOR_CACHE_DIR}" "${TRITON_CACHE_DIR}" \
-    "${VLLM_CACHE_ROOT}" "${XDG_CACHE_HOME}" "${HF_HOME}" \
-    "${HF_DATASETS_CACHE}"
-
-if [[ "${VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE:-0}" == "1" ]]; then
-    if ! command -v findmnt >/dev/null 2>&1; then
-        echo "Error: findmnt is required to verify the Hugging Face cache mount." >&2
-        exit 1
-    fi
-    hf_mount=$(findmnt -n -T "${HF_HOME}" -o TARGET 2>/dev/null || true)
-    if [[ -z "${hf_mount}" || "${hf_mount}" == "/" ]]; then
-        echo "Error: MI300 CI requires a persistent volume at or above ${HF_HOME}." >&2
-        exit 1
-    fi
-fi
+# Export Python path
+export PYTHONPATH=".."

+# Print ROCm version
+echo "--- ROCm info"
 rocminfo

-expected_gpus="${VLLM_CI_EXPECTED_GPU_COUNT:-1}"
-python3 - <<'PY'
-import os
+# cleanup older docker images
+cleanup_docker() {
+  # Get Docker's root directory
+  docker_root=$(docker info -f '{{.DockerRootDir}}')
+  if [ -z "$docker_root" ]; then
+    echo "Failed to determine Docker root directory."
+    exit 1
+  fi
+  echo "Docker root directory: $docker_root"
+  # Check disk usage of the filesystem where Docker's root directory is located
+  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
+  # Define the threshold
+  threshold=70
+  if [ "$disk_usage" -gt "$threshold" ]; then
+    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
+    # Remove dangling images (those that are not tagged and not used by any container)
+    docker image prune -f
+    # Remove unused volumes / force the system prune for old images as well.
+    docker volume prune -f && docker system prune --force --filter "until=72h" --all
+    echo "Docker images and volumes cleanup completed."
+  else
+    echo "Disk usage is below $threshold%. No cleanup needed."
+  fi
+}

-import torch
+# Call the cleanup docker function
+cleanup_docker

-expected = int(os.environ.get("VLLM_CI_EXPECTED_GPU_COUNT", "1"))
-assert torch.version.hip, "PyTorch is not a ROCm build"
-assert torch.cuda.is_available(), "ROCm GPU is not available to PyTorch"
-actual = torch.cuda.device_count()
-assert actual == expected, f"Expected {expected} ROCm GPU(s), found {actual}"
-PY
+echo "--- Pulling container"
+## Temporary change to use AMD Docker Hub to store the vllm-omni image
+# to bypass the rate limit issue with ECR Public Gallery.
+# Images are now stored in a separate repository for vllm-omni, instead of vllm-ci.
+# TODO: @tjtanaa point back to ECR Public Gallery
+# once the amd agents are configured to use ECR Public Gallery.
+# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+image_name="rocm/vllm-omni:${BUILDKITE_COMMIT}"
+container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"

-echo "Commands:${commands}"
-if /bin/bash -o pipefail -c '
-set -E
-test_status=0
-trap '\''
-    command_status=$?
-    # Keep running subsequent commands, but retain a failure for the final exit.
-    # Prefer a test failure over pytest "no tests collected" when both occur.
-    if (( test_status == 0 || (test_status == 5 && command_status != 5) )); then
-        test_status=${command_status}
-    fi
-'\'' ERR
-eval "$1"
-exit "${test_status}"
-' _ "${commands}"; then
-    exit 0
+# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
+# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
+# sudo apt-get update && sudo apt-get install -y awscli
+##  Use safe docker login helper to prevent race conditions
+# source "$(dirname "${BASH_SOURCE[0]}")/../../common/scripts/docker_login_ecr_public.sh"
+# safe_docker_login_ecr_public
+
+## Pull the container from AMD Docker Hub
+
+docker pull "${image_name}"
+
+remove_docker_container() {
+   docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true
+   docker image rm -f "${image_name}" || true
+}
+trap remove_docker_container EXIT
+
+echo "--- Running container"
+
+HF_CACHE="$(realpath ~)/huggingface"
+mkdir -p "${HF_CACHE}"
+HF_MOUNT="/root/.cache/huggingface"
+
+if [[ -n "${TEST_COMMAND:-}" ]]; then
+    commands="$TEST_COMMAND"
 else
-    exit_code=$?
+    commands="$@"
+fi
+echo "Commands:$commands"
+
+PARALLEL_JOB_COUNT=8
+MYPYTHONPATH=".."
+
+# Test that we're launching on the machine that has
+# proper access to GPUs
+render_gid=$(getent group render | cut -d: -f3)
+if [[ -z "$render_gid" ]]; then
+  echo "Error: 'render' group not found. This is required for GPU access." >&2
+  exit 1
 fi

-if [[ ${exit_code} -eq 5 && "${VLLM_CI_ALLOW_NO_TESTS:-0}" == "1" ]]; then
-    echo "Pytest collected no tests; VLLM_CI_ALLOW_NO_TESTS=1, treating exit code 5 as success."
-    exit 0
+# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
+# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
+if [[ $commands == *"--shard-id="* ]]; then
+  # assign job count as the number of shards used
+  commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
+  for GPU in $(seq 0 $(($PARALLEL_JOB_COUNT-1))); do
+    # assign shard-id for each shard
+    commands_gpu=$(echo "$commands" | sed -E "s/--shard-id[[:blank:]]*=[[:blank:]]*[0-9]*/--shard-id=${GPU} /g" | sed 's/ \\ / /g')
+    echo "Shard ${GPU} commands:$commands_gpu"
+    echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
+    docker run \
+        --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
+        --network=host \
+        --shm-size=16gb \
+        --group-add "$render_gid" \
+        -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+        -e MIOPEN_DEBUG_CONV_GEMM=0 \
+        -e VLLM_ROCM_USE_AITER=0 \
+        -e HIP_VISIBLE_DEVICES="${GPU}" \
+        -e HF_TOKEN \
+        -v "${HF_CACHE}:${HF_MOUNT}" \
+        -e "HF_HOME=${HF_MOUNT}" \
+        -e "PYTHONPATH=${MYPYTHONPATH}" \
+        --name "${container_name}_${GPU}" \
+        "${image_name}" \
+        /bin/bash -c "${commands_gpu}" \
+        |& while read -r line; do echo ">>Shard $GPU: $line"; done &
+    PIDS+=($!)
+  done
+  #wait for all processes to finish and collect exit codes
+  for pid in "${PIDS[@]}"; do
+    wait "${pid}"
+    STATUS+=($?)
+  done
+  for st in "${STATUS[@]}"; do
+    if [[ ${st} -ne 0 ]]; then
+      echo "One of the processes failed with $st"
+      exit "${st}"
+    fi
+  done
+else
+  echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
+  docker run \
+          --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
+          --network=host \
+          --shm-size=16gb \
+          --group-add "$render_gid" \
+          -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+          -e MIOPEN_DEBUG_CONV_GEMM=0 \
+          -e VLLM_ROCM_USE_AITER=0 \
+          -e HF_TOKEN \
+          -v "${HF_CACHE}:${HF_MOUNT}" \
+          -e "HF_HOME=${HF_MOUNT}" \
+          -e "PYTHONPATH=${MYPYTHONPATH}" \
+          --name "${container_name}" \
+          "${image_name}" \
+          /bin/bash -c "${commands}"
 fi
-
-exit "${exit_code}"
@@ -7,64 +7,54 @@ steps:

 - group: ":card_index_dividers: Simple Test"
  steps:
-    - label: "Simple · Model Executor Test · Shard %N/%t"
-      agent_pool: mi300_1
-      parallelism: 2
+    - label: "Simple · Diffusion & Model Executor Test (1/2)"
+      agent_pool: mi325_1
      depends_on: amd-build
+      timeout_in_minutes: 70
      mirror_hardwares: [amdproduction]
      grade: Blocking
      commands:
        - export VLLM_ROCM_USE_AITER=0
        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
        # that is not supported on CPU
-        - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml"
+        # CPU tests on ROCm gpu machines are very slow, so we give them more time.
+        - "timeout 60m pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"

-    - label: "Simple · Diffusion Test · Shard %N/%t"
-      agent_pool: mi300_1
-      parallelism: 2
+    - label: "Simple · Diffusion & Model Executor Test (2/2)"
+      agent_pool: mi325_1
      depends_on: amd-build
+      timeout_in_minutes: 70
      mirror_hardwares: [amdproduction]
      grade: Blocking
      commands:
        - export VLLM_ROCM_USE_AITER=0
-        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-        # that is not supported on CPU
-        - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-diffusion-$$BUILDKITE_PARALLEL_JOB.xml"
-
-    - label: "Simple · Engine&Entrypoints Test"
-      agent_pool: mi300_1
-      depends_on: amd-build
-      mirror_hardwares: [amdproduction]
-      grade: Blocking
-      commands:
-        - export VLLM_ROCM_USE_AITER=0
-        - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        # CPU tests on ROCm gpu machines are very slow, so we give them more time.
+        - "timeout 60m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"

    - label: "Simple · Other Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
+      timeout_in_minutes: 50
      mirror_hardwares: [amdproduction]
      grade: Blocking
      commands:
        - export VLLM_ROCM_USE_AITER=0
        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
        # that is not supported on CPU
-        - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine  --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"

 - group: ":card_index_dividers: Diffusion Test"
  steps:
    - label: "Diffusion · LoRA Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
-      artifact_paths:
-        - lora-server.log
      commands:
-        - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model" 2>&1 | tee lora-server.log
+        - timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py

    - label: "Diffusion Model CPU offloading Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -72,7 +62,7 @@ steps:
        - export VLLM_LOGGING_LEVEL=DEBUG
        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
        - |
-          bash -c '
+          timeout 20m bash -c '
            set +e
            pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
            EXIT1=\$?
@@ -82,7 +72,7 @@ steps:
          '

    - label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
-      agent_pool: mi300_4
+      agent_pool: mi325_4
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -94,7 +84,7 @@ steps:

    # merge-only tests
    - label: "Diffusion Tensor Parallelism Test"
-      agent_pool: mi300_2
+      agent_pool: mi325_2
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -102,7 +92,7 @@ steps:
        - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py

    - label: "Diffusion Model Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -110,7 +100,7 @@ steps:
        - timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py

    - label: "Diffusion GPU Worker Test"
-      agent_pool: mi300_2
+      agent_pool: mi325_2
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -118,7 +108,7 @@ steps:
        - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py

    - label: "Diffusion · Wan22 Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -131,7 +121,7 @@ steps:

 ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
 # - label: "Audio Generation Model Test"
-#   agent_pool: mi300_1
+#   agent_pool: mi325_1
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -142,7 +132,7 @@ steps:
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model

 - label: "Engine Test"
-  agent_pool: mi300_1
+  agent_pool: mi325_1
  depends_on: amd-build
  mirror_hardwares: [amdproduction]
  grade: Blocking
@@ -157,7 +147,7 @@ steps:
 - group: ":card_index_dividers: E2E Test"
  steps:
    - label: "Omni Model Test Qwen2-5-Omni"
-      agent_pool: mi300_4
+      agent_pool: mi325_4
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      commands:
@@ -165,7 +155,7 @@ steps:
        - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"

    - label: "Omni Model Test Qwen3-Omni"
-      agent_pool: mi300_2
+      agent_pool: mi325_2
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -173,23 +163,22 @@ steps:
        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
        - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'

-    - label: "Qwen3-TTS CustomVoice E2E Test · Shard %N/%t"
-      agent_pool: mi300_1
-      parallelism: 2
+    - label: "Qwen3-TTS CustomVoice E2E Test"
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
      commands:
        - |
-          bash -c '
+          timeout 50m bash -c '
            export VLLM_LOGGING_LEVEL=DEBUG
            export VLLM_WORKER_MULTIPROC_METHOD=spawn
            export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
+            pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
          '

    - label: "Qwen3-TTS Base E2E Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -207,7 +196,7 @@ steps:
 # Need to debug before reneable numerical changes across large PRs
 # # split Bagel Model Test with H100 (Real Weights) into three tests
 # - label: "Bagel Text2Img Model Test (1/3)"
-#   agent_pool: mi300_1
+#   agent_pool: mi325_1
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -219,7 +208,7 @@ steps:
 #     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"

 # - label: "Bagel Img2Img Model Test (2/3)"
-#   agent_pool: mi300_1
+#   agent_pool: mi325_1
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -231,7 +220,7 @@ steps:
 #     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"

 # - label: "Bagel Online Serving Test (3/3)"
-#   agent_pool: mi300_1
+#   agent_pool: mi325_1
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -7,9 +7,8 @@ steps:

 - group: ":card_index_dividers: Simple Test"
  steps:
-    - label: "Simple · Model Executor Test · Shard %N/%t"
-      agent_pool: mi300_1
-      parallelism: 2
+    - label: "Simple · Diffusion & Model Executor Test"
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -17,30 +16,10 @@ steps:
        - export VLLM_ROCM_USE_AITER=0
        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
        # that is not supported on CPU
-        - "pytest -sv tests/model_executor -m 'core_model and cpu' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml:coverage-model-executor-$$BUILDKITE_PARALLEL_JOB.xml"
-
-    - label: "Simple · Diffusion Test"
-      agent_pool: mi300_1
-      depends_on: amd-build
-      mirror_hardwares: [amdproduction]
-      grade: Blocking
-      commands:
-        - export VLLM_ROCM_USE_AITER=0
-        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
-        # that is not supported on CPU
-        - "pytest -sv tests/diffusion -m 'core_model and cpu'  --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
-
-    - label: "Simple · Engine&Entrypoints Test"
-      agent_pool: mi300_1
-      depends_on: amd-build
-      mirror_hardwares: [amdproduction]
-      grade: Blocking
-      commands:
-        - export VLLM_ROCM_USE_AITER=0
-        - "pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"

    - label: "Simple · Other Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -48,19 +27,19 @@ steps:
        - export VLLM_ROCM_USE_AITER=0
        # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
        # that is not supported on CPU
-        - "pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine  --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        - "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"

 - label: "Custom Pipeline Test"
-  agent_pool: mi300_1
+  agent_pool: mi325_1
  depends_on: amd-build
  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
-    - pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
+    - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"


 - label: "Voxtral TTS CUDA Unit Test"
-  agent_pool: mi300_1
+  agent_pool: mi325_1
  depends_on: amd-build
  mirror_hardwares: [amdproduction]
  grade: Blocking
@@ -70,25 +49,33 @@ steps:
 - group: ":card_index_dividers: Diffusion Test"
  steps:
    - label: "Diffusion · Batch Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
      commands:
-        - pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
+        - timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"

    - label: "Diffusion · Offloader Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
      commands:
        - export VLLM_LOGGING_LEVEL=DEBUG
        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        - pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
+        - |
+          timeout 20m bash -c '
+            set +e
+            pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
+            EXIT1=\$?
+            pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
+            EXIT2=\$?
+            exit \$((EXIT1 | EXIT2))
+          '

    - label: "Diffusion · Cache Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -96,10 +83,10 @@ steps:
        - export GPU_ARCHS=gfx942
        - export VLLM_LOGGING_LEVEL=DEBUG
        - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-        - pytest -sv tests/diffusion/cache -m 'core_model and cuda' --run-level "core_model"
+        - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"

    - label: "Diffusion Sequence Parallelism Test"
-      agent_pool: mi300_2
+      agent_pool: mi325_2
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -109,7 +96,7 @@ steps:
        - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model

    - label: "Diffusion GPU Worker Test"
-      agent_pool: mi300_2
+      agent_pool: mi325_2
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -117,7 +104,7 @@ steps:
        - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py

    - label: "Diffusion · Model Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -125,7 +112,7 @@ steps:
        - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"

    - label: "Diffusion · Wan22 Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -139,7 +126,7 @@ steps:

 ## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
 # - label: "Audio Generation Model Test"
-#   agent_pool: mi300_1
+#   agent_pool: mi325_1
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -150,7 +137,7 @@ steps:
 #     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model

 - label: "Engine Test"
-  agent_pool: mi300_1
+  agent_pool: mi325_1
  depends_on: amd-build
  mirror_hardwares: [amdproduction]
  grade: Blocking
@@ -162,7 +149,7 @@ steps:
      '

 # - label: "Omni Model Test Qwen3-Omni"
-#   agent_pool: mi300_2
+#   agent_pool: mi325_2
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -175,7 +162,7 @@ steps:
 - group: ":card_index_dividers: E2E Test"
  steps:
    - label: "Qwen3-TTS E2E Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -186,7 +173,7 @@ steps:
        - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"

    - label: "CosyVoice3-TTS E2E Test"
-      agent_pool: mi300_1
+      agent_pool: mi325_1
      depends_on: amd-build
      mirror_hardwares: [amdproduction]
      grade: Blocking
@@ -199,7 +186,7 @@ steps:
 # TODO: Bagel test on ROCm is very unstable. @tjtanaa
 # Need to debug before reneable numerical changes across large PRs
 # - label: "Bagel Text2Img Model Test"
-#   agent_pool: mi300_1
+#   agent_pool: mi325_1
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -211,7 +198,7 @@ steps:
 #     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"

 # - label: "Bagel Img2Img Model Test"
-#   agent_pool: mi300_1
+#   agent_pool: mi325_1
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -223,7 +210,7 @@ steps:
 #     - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"

 # - label: "Bagel Online Serving Test"
-#   agent_pool: mi300_1
+#   agent_pool: mi325_1
 #   depends_on: amd-build
 #   mirror_hardwares: [amdproduction]
 #   grade: Blocking
@@ -1,65 +1,11 @@
 {# vllm-omni customized version
   Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/test-template-amd.j2
-   Last synced: 2026-08-04
-   Modifications: Use the shared MI300 queues and Omni command-passing contract
+   Last synced: 2025-12-15
+   Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
 #}
 {% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
 {% set default_working_dir = "/app/vllm-omni" %}

-{# MI300 runs tests natively in Kubernetes; Docker-in-Docker is unavailable. #}
-{% macro amd_native_k8s_plugin(indent, step) -%}
-{{ indent }}plugins:
-{{ indent }}  - kubernetes:
-{{ indent }}      podSpecPatch:
-{{ indent }}        automountServiceAccountToken: false
-{{ indent }}        securityContext:
-{{ indent }}          seccompProfile:
-{{ indent }}            type: RuntimeDefault
-{{ indent }}        imagePullSecrets:
-{{ indent }}          - name: docker-config
-{{ indent }}        containers:
-{{ indent }}          - name: container-0
-{{ indent }}            image: {{ docker_image_amd }}
-{{ indent }}            imagePullPolicy: Always
-{{ indent }}            securityContext:
-{{ indent }}              allowPrivilegeEscalation: false
-{{ indent }}              capabilities:
-{{ indent }}                add:
-{{ indent }}                  - IPC_LOCK
-{{ indent }}            resources:
-{{ indent }}              limits:
-{{ indent }}                amd.com/gpu: "{{ step.agent_pool[-1:] }}"
-{{ indent }}              requests:
-{{ indent }}                amd.com/gpu: "{{ step.agent_pool[-1:] }}"
-{{ indent }}            volumeMounts:
-{{ indent }}              - name: devshm
-{{ indent }}                mountPath: /dev/shm
-{{ indent }}            env:
-{{ indent }}              - name: AMD_CI_RUNTIME
-{{ indent }}                value: native
-{{ indent }}              - name: NATIVE_CI
-{{ indent }}                value: "true"
-{{ indent }}              - name: VLLM_CI_DOCKER_DISABLED
-{{ indent }}                value: "1"
-{{ indent }}              - name: VLLM_CI_EXPECTED_GPU_COUNT
-{{ indent }}                value: "{{ step.agent_pool[-1:] }}"
-{{ indent }}              - name: HF_HOME
-{{ indent }}                value: /home/buildkite-agent/huggingface
-{{ indent }}              - name: MIOPEN_DEBUG_CONV_DIRECT
-{{ indent }}                value: "0"
-{{ indent }}              - name: MIOPEN_DEBUG_CONV_GEMM
-{{ indent }}                value: "0"
-{{ indent }}              - name: VLLM_ROCM_USE_AITER
-{{ indent }}                value: "0"
-{{ indent }}              - name: PYTORCH_ROCM_ARCH
-{{ indent }}                value: ""
-{{ indent }}        volumes:
-{{ indent }}          - name: devshm
-{{ indent }}            emptyDir:
-{{ indent }}              medium: Memory
-{{ indent }}              sizeLimit: 16Gi
-{%- endmacro %}
-
 {% if env %}
 env:
 {% for name in env %}
@@ -107,31 +53,19 @@ steps:
          {% if step.agent_pool %}
          queue: amd_{{ step.agent_pool }}
          {% else %}
-          queue: amd_mi300_1
+          queue: amd_mi325_1
          {% endif %}
-{{ amd_native_k8s_plugin("        ", step) }}
 {% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
 {% set indented_cmd = cmd_body | replace("\n", "\n            ") %}
        command: bash .buildkite/amd/scripts/run-amd-test.sh
        env:
-          AMD_CI_RUNTIME: "native"
-          NATIVE_CI: "true"
-          VLLM_CI_DOCKER_DISABLED: "1"
-          VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}"
-          VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1"
-          HF_HOME: "/home/buildkite-agent/huggingface"
-          HF_HUB_DOWNLOAD_TIMEOUT: "300"
-          HF_HUB_ETAG_TIMEOUT: "60"
-          MIOPEN_DEBUG_CONV_DIRECT: "0"
-          MIOPEN_DEBUG_CONV_GEMM: "0"
-          VLLM_ROCM_USE_AITER: "0"
-          PYTORCH_ROCM_ARCH: ""
+          DOCKER_BUILDKIT: "1"
 {% if step.env %}
 {% for name in step.env %}
          {{ name }}: "{{ step.env[name] }}"
 {% endfor %}
 {% endif %}
-          TEST_COMMANDS: |-
+          TEST_COMMAND: |-
            (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
 {% if env %}
 {% for name in env %}
@@ -148,15 +82,6 @@ steps:
 {% endif %}
            {{ indented_cmd | safe }}
        priority: 100
-{% if step.parallelism %}
-        parallelism: {{ step.parallelism }}
-{% endif %}
-{% if step.artifact_paths %}
-        artifact_paths:
-{% for path in step.artifact_paths %}
-          - "{{ path }}"
-{% endfor %}
-{% endif %}
        {% if step.timeout_in_minutes %}
        timeout_in_minutes: {{ step.timeout_in_minutes }}
        {% endif %}
@@ -175,31 +100,19 @@ steps:
      {% if step.agent_pool %}
      queue: amd_{{ step.agent_pool }}
      {% else %}
-      queue: amd_mi300_1
+      queue: amd_mi325_1
      {% endif %}
-{{ amd_native_k8s_plugin("    ", step) }}
 {% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
 {% set indented_cmd = cmd_body | replace("\n", "\n        ") %}
    command: bash .buildkite/amd/scripts/run-amd-test.sh
    env:
-      AMD_CI_RUNTIME: "native"
-      NATIVE_CI: "true"
-      VLLM_CI_DOCKER_DISABLED: "1"
-      VLLM_CI_EXPECTED_GPU_COUNT: "{{ step.agent_pool[-1:] }}"
-      VLLM_CI_REQUIRE_PERSISTENT_HF_CACHE: "1"
-      HF_HOME: "/home/buildkite-agent/huggingface"
-      HF_HUB_DOWNLOAD_TIMEOUT: "300"
-      HF_HUB_ETAG_TIMEOUT: "60"
-      MIOPEN_DEBUG_CONV_DIRECT: "0"
-      MIOPEN_DEBUG_CONV_GEMM: "0"
-      VLLM_ROCM_USE_AITER: "0"
-      PYTORCH_ROCM_ARCH: ""
+      DOCKER_BUILDKIT: "1"
 {% if step.env %}
 {% for name in step.env %}
      {{ name }}: "{{ step.env[name] }}"
 {% endfor %}
 {% endif %}
-      TEST_COMMANDS: |-
+      TEST_COMMAND: |-
        (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
 {% if env %}
 {% for name in env %}
@@ -216,15 +129,6 @@ steps:
 {% endif %}
        {{ indented_cmd | safe }}
    priority: 100
-{% if step.parallelism %}
-    parallelism: {{ step.parallelism }}
-{% endif %}
-{% if step.artifact_paths %}
-    artifact_paths:
-{% for path in step.artifact_paths %}
-      - "{{ path }}"
-{% endfor %}
-{% endif %}
    {% if step.timeout_in_minutes %}
    timeout_in_minutes: {{ step.timeout_in_minutes }}
    {% endif %}
@@ -9,11 +9,6 @@
 #   a2b3_npu_1  — ascend-a2b3, resource_class npu-1 (A2/B3 CI image)
 #   a2b3_npu_4  — ascend-a2b3, resource_class npu-4
 #   a3_npu_2    — ascend-a3, resource_class npu-2
-#
-# The ``l4_*`` presets run under the ``docker`` plugin, which does not put
-# ``buildkite-agent`` on PATH inside the container the way the ``kubernetes``
-# presets do, so they set ``mount-buildkite-agent: true``; without it any step
-# calling ``buildkite-agent`` (artifact upload, annotate, meta-data) fails there.

 x-npu-k8s-plugins: &npu_k8s_plugins
  - kubernetes:
@@ -31,7 +26,6 @@ mirror_hardwares:
          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
          always-pull: true
          propagate-environment: true
-          mount-buildkite-agent: true
          shm-size: "8gb"
          environment:
            - HF_HOME=/fsx/hf_cache
@@ -47,7 +41,6 @@ mirror_hardwares:
          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
          always-pull: true
          propagate-environment: true
-          mount-buildkite-agent: true
          shm-size: "8gb"
          environment:
            - HF_HOME=/fsx/hf_cache
@@ -1,82 +0,0 @@
-#!/usr/bin/env bash
-# Run a model's E2E tests once per entry mode as separate coverage runs and upload
-# the reports, so per-mode attribution survives (see vllm-project/vllm-omni#5332).
-#
-# Usage: run_cov_split.sh --model-id <id> [--offline <paths>] [--online <paths>] \
-#                         -- <pytest args...>
-#
-#   --model-id  Names the artifacts: coverage-<id>-{offline,online}-<step id>.xml.gz.
-#               Use the model's directory name under vllm_omni/*/models/.
-#   --offline   Pytest path(s) for the offline half. Repeatable.
-#   --online    Pytest path(s) for the online half. Repeatable.
-#   --          Everything after this is forwarded verbatim to every pytest run,
-#               quoting preserved: the job's own -m expression, --run-level, and
-#               anything else it needs (e.g. --test-config-file).
-#
-# At least one of --offline/--online is required, so a job with only one mode runs
-# just that half. Total runtime is bounded by the step's timeout_in_minutes.
-#
-# Exits non-zero if any half or the upload failed.
-set -uo pipefail
-
-MODEL_ID=""
-OFFLINE=()
-ONLINE=()
-PYTEST_ARGS=()
-
-while [[ $# -gt 0 ]]; do
-    case "$1" in
-        --model-id) MODEL_ID="$2"; shift 2 ;;
-        --offline)  OFFLINE+=("$2"); shift 2 ;;
-        --online)   ONLINE+=("$2"); shift 2 ;;
-        --)         shift; PYTEST_ARGS=("$@"); break ;;
-        *) echo "run_cov_split.sh: unknown argument: $1" >&2; exit 2 ;;
-    esac
-done
-
-if [[ -z "${MODEL_ID}" ]]; then
-    echo "run_cov_split.sh: missing --model-id" >&2
-    exit 2
-fi
-
-if ((${#OFFLINE[@]} == 0 && ${#ONLINE[@]} == 0)); then
-    echo "run_cov_split.sh: need at least one of --offline / --online" >&2
-    exit 2
-fi
-
-# Two steps can test the same model, and artifacts are scoped to the build (a
-# nightly one carries the ready, merge and nightly tiers at once), so unsuffixed
-# reports would share a path -- a glob download then races to rename onto one file.
-# Not BUILDKITE_JOB_ID: the step id survives retries, so attempts collapse under
-# Buildkite's latest-attempt lookup instead of accumulating. A parallel step would
-# need BUILDKITE_PARALLEL_JOB as well; no pilot is one.
-REPORT_SUFFIX="${BUILDKITE_STEP_ID:-local}"
-OFFLINE_REPORT="coverage-${MODEL_ID}-offline-${REPORT_SUFFIX}.xml"
-ONLINE_REPORT="coverage-${MODEL_ID}-online-${REPORT_SUFFIX}.xml"
-
-run_half() {
-    local mode="$1"
-    shift
-    local report="coverage-${MODEL_ID}-${mode}-${REPORT_SUFFIX}.xml"
-    echo "--- coverage: ${MODEL_ID} ${mode}"
-    pytest -s -v "$@" "${PYTEST_ARGS[@]}" --cov=vllm_omni --cov-report="xml:${report}"
-    local status=$?
-    # Cobertura lists every statement in vllm_omni whether or not it ran, so the
-    # report is ~7MB no matter how little the job covers. The repetition is what
-    # dominates it, and gzip takes that to ~400KB.
-    gzip -9 -f "${report}" || status=1
-    return "${status}"
-}
-
-# Clear both modes, not just the ones being run: the upload glob matches both, so
-# a report left by an earlier run would otherwise be published as if the mode it
-# belongs to had run this time.
-rm -f "${OFFLINE_REPORT}" "${ONLINE_REPORT}" "${OFFLINE_REPORT}.gz" "${ONLINE_REPORT}.gz"
-
-EXIT=0
-((${#OFFLINE[@]})) && { run_half offline "${OFFLINE[@]}" || EXIT=1; }
-((${#ONLINE[@]})) && { run_half online "${ONLINE[@]}" || EXIT=1; }
-
-buildkite-agent artifact upload "coverage-${MODEL_ID}-*-${REPORT_SUFFIX}.xml.gz" || EXIT=1
-
-exit "${EXIT}"
@@ -76,9 +76,6 @@ steps:
          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
          - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
          - vllm_omni/model_executor/models/qwen3_tts/
-          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
-          - vllm_omni/model_executor/models/common/snake_activation.py
-          - vllm_omni/model_executor/models/common/alias_free_activation.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
@@ -95,24 +92,31 @@ steps:
          - tests/e2e/online_serving/test_qwen3_tts_base.py
          - tests/e2e/offline_inference/test_qwen3_tts_base.py
          - vllm_omni/model_executor/models/qwen3_tts/
-          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
-          - vllm_omni/model_executor/models/common/snake_activation.py
-          - vllm_omni/model_executor/models/common/alias_free_activation.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
-          - .buildkite/common/scripts/run_cov_split.sh
-          - pyproject.toml
-        # 40 -> 60: the coverage split runs the model twice.
-        timeout_in_minutes: 60
        commands:
          - |
-            export VLLM_LOGGING_LEVEL=DEBUG
-            export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
-            .buildkite/common/scripts/run_cov_split.sh \
-              --model-id qwen3_tts \
-              --offline tests/e2e/offline_inference/test_qwen3_tts_base.py \
-              --online tests/e2e/online_serving/test_qwen3_tts_base.py \
-              -- -m 'advanced_model and cuda' --run-level 'advanced_model'
+            timeout 40m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
+              pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        mirror_hardwares: l4_1
+
+      - label: "TTS · Higgs-Audio-V3 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_higgs_audio_v3.py
+          - vllm_omni/model_executor/models/higgs_audio_v3/
+          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
+          - vllm_omni/deploy/higgs_multimodal_qwen3.yaml
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_USE_DEEP_GEMM=0
+              export VLLM_MOE_USE_DEEP_GEMM=0
+              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'advanced_model' --run-level 'advanced_model'
+            "
        mirror_hardwares: l4_1

      - label: "Omni · Qwen3-Omni Test"
@@ -120,10 +124,6 @@ steps:
          - tests/e2e/offline_inference/test_qwen3_omni.py
          - tests/e2e/online_serving/test_qwen3_omni.py
          - vllm_omni/model_executor/models/qwen3_omni/
-          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
-          - vllm_omni/model_executor/models/common/snake_activation.py
-          - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
-          - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
          - vllm_omni/deploy/qwen3_omni_moe.yaml
        timeout_in_minutes: 50
@@ -135,11 +135,6 @@ steps:
        source_file_dependencies:
          - tests/e2e/offline_inference/test_minicpmo_4_5.py
          - tests/e2e/online_serving/test_minicpmo_4_5.py
-          - vllm_omni/model_executor/models/minicpmo_4_5/
-          - vllm_omni/model_executor/models/step_audio2/
-          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
-          - vllm_omni/model_executor/models/cosyvoice3/utils.py
-          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/deploy/minicpmo_4_5.yaml
        timeout_in_minutes: 50
        commands:
@@ -161,35 +156,62 @@ steps:
          - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
-          - vllm_omni/deploy/minicpmo_4_5.yaml
          - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
          - vllm_omni/deploy/minicpmo_4_5_duplex.yaml
          - vllm_omni/experimental/fullduplex/
          - vllm_omni/model_executor/models/minicpmo_4_5/
-          - vllm_omni/model_executor/models/step_audio2/
-          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
-          - vllm_omni/model_executor/models/cosyvoice3/utils.py
          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/core/sched/
          - vllm_omni/outputs/multimodal_accumulation.py
          - vllm_omni/engine/stage_init_utils.py
          - vllm_omni/engine/async_omni_engine.py
-          - vllm_omni/engine/orchestrator.py
          - vllm_omni/entrypoints/cli/serve.py
-          - vllm_omni/entrypoints/openai/api_server.py
-          - vllm_omni/entrypoints/openai/duplex_capability.py
-          - vllm_omni/config/stage_config.py
        timeout_in_minutes: 50
        commands:
          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
        mirror_hardwares: h100_1

+      - label: "TTS · VoxCPM2 Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_voxcpm2_tts.py
+          - tests/e2e/online_serving/test_voxcpm2_tts.py
+          - vllm_omni/model_executor/models/voxcpm2/
+          - vllm_omni/deploy/voxcpm2.yaml
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        mirror_hardwares: l4_1
+
+      - label: "Diffusion · Qwen Image Edit Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_edit.py
+          - vllm_omni/diffusion/models/qwen_image/
+        timeout_in_minutes: 20
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_1
+      - label: "Diffusion · Bagel Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_bagel.py
+          - tests/e2e/online_serving/test_bagel.py
+          - vllm_omni/model_executor/models/bagel/
+          - vllm_omni/diffusion/models/bagel/
+          - vllm_omni/model_executor/stage_input_processors/bagel.py
+          - vllm_omni/deploy/bagel.yaml
+        timeout_in_minutes: 20
+        commands:
+          - pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_1
+
      - label: "Diffusion · Qwen Image Test"
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen_image.py
          - vllm_omni/diffusion/models/qwen_image/
-          - vllm_omni/diffusion/models/dmd2/
-          - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
        commands:
          - |
            timeout 40m bash -c '
@@ -198,12 +220,23 @@ steps:
            '
        mirror_hardwares: h100_1

+      - label: "Diffusion · Qwen Image Layered Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_layered.py
+          - vllm_omni/diffusion/models/qwen_image/
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+            "
+        mirror_hardwares: h100_1
+
      - label: "Diffusion · Wan22 Test"
        source_file_dependencies:
          - tests/e2e/offline_inference/test_wan22_t2v.py
          - tests/e2e/online_serving/test_wan22_t2v.py
          - vllm_omni/diffusion/models/wan2_2/
-          - vllm_omni/diffusion/models/dmd2/
        commands:
          - |
            timeout 40m bash -c "
@@ -211,3 +244,21 @@ steps:
              pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
            "
        mirror_hardwares: h100_1
+
+      - label: "Diffusion · Z Image Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/test_z_image.py
+          - tests/e2e/online_serving/test_z_image.py
+          - vllm_omni/diffusion/models/z_image/
+        timeout_in_minutes: 30
+        commands:
+          - pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+        mirror_hardwares: l4_1
+      - label: "Diffusion · FLUX.2 Klein Test"
+        source_file_dependencies:
+          - tests/e2e/offline_inference/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 -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
+        mirror_hardwares: l4_1
@@ -459,12 +459,6 @@ steps:
          - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
        mirror_hardwares: h100_2

-      - label: ":full_moon: Diffusion X2V · LTX-2.5 Online Test"
-        timeout_in_minutes: 60
-        commands:
-          - pytest -s -v tests/e2e/online_serving/test_ltx25.py -m "full_model and cuda" --run-level "full_model"
-        mirror_hardwares: h100_1
-
      - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
        timeout_in_minutes: 180
        commands:
@@ -493,13 +487,6 @@ steps:
          - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
        mirror_hardwares: h100_1

-      - label: ":full_moon: Diffusion X2V · LTX-2.5 Official Parity Test"
-        timeout_in_minutes: 180
-        commands:
-          - pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_distilled_two_stage_matches_official -m full_model --run-level full_model
-          - pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_full_matches_official -m full_model --run-level full_model
-        mirror_hardwares: h100_1
-
      - label: ":full_moon: Diffusion X2V · Perf Test"
        key: nightly-diffusion-x2v-performance
        timeout_in_minutes: 180
@@ -169,10 +169,6 @@ steps:
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen3_omni.py
          - vllm_omni/model_executor/models/qwen3_omni/
-          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
-          - vllm_omni/model_executor/models/common/snake_activation.py
-          - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py
-          - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
          - vllm_omni/deploy/qwen3_omni_moe.yaml
        commands:
@@ -185,11 +181,6 @@ steps:
      - label: "Omni · MiniCPM-o 4.5 Offline Test"
        source_file_dependencies:
          - tests/e2e/offline_inference/test_minicpmo_4_5.py
-          - vllm_omni/model_executor/models/minicpmo_4_5/
-          - vllm_omni/model_executor/models/step_audio2/
-          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
-          - vllm_omni/model_executor/models/cosyvoice3/utils.py
-          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/deploy/minicpmo_4_5.yaml
        commands:
          - |
@@ -201,11 +192,6 @@ steps:
      - label: "Omni · MiniCPM-o 4.5 Online Test"
        source_file_dependencies:
          - tests/e2e/online_serving/test_minicpmo_4_5.py
-          - vllm_omni/model_executor/models/minicpmo_4_5/
-          - vllm_omni/model_executor/models/step_audio2/
-          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
-          - vllm_omni/model_executor/models/cosyvoice3/utils.py
-          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/deploy/minicpmo_4_5.yaml
        commands:
          - |
@@ -224,24 +210,16 @@ steps:
          - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py
          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
-          - vllm_omni/deploy/minicpmo_4_5.yaml
          - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py
          - vllm_omni/deploy/minicpmo_4_5_duplex.yaml
          - vllm_omni/experimental/fullduplex/
          - vllm_omni/model_executor/models/minicpmo_4_5/
-          - vllm_omni/model_executor/models/step_audio2/
-          - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/
-          - vllm_omni/model_executor/models/cosyvoice3/utils.py
          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
          - vllm_omni/core/sched/
          - vllm_omni/outputs/multimodal_accumulation.py
          - vllm_omni/engine/stage_init_utils.py
          - vllm_omni/engine/async_omni_engine.py
-          - vllm_omni/engine/orchestrator.py
          - vllm_omni/entrypoints/cli/serve.py
-          - vllm_omni/entrypoints/openai/api_server.py
-          - vllm_omni/entrypoints/openai/duplex_capability.py
-          - vllm_omni/config/stage_config.py
        commands:
          - |
            timeout 20m bash -c "
@@ -261,13 +239,25 @@ steps:
          - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda'
        mirror_hardwares: h100_1

+      - label: "TTS · VoxCPM2 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_voxcpm2_tts.py
+          - vllm_omni/model_executor/models/voxcpm2/
+          - vllm_omni/deploy/voxcpm2.yaml
+        timeout_in_minutes: 30
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_LOGGING_LEVEL=DEBUG
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
+            "
+        mirror_hardwares: l4_1
+
      - label: "TTS · Qwen3-TTS CustomVoice Test"
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
          - vllm_omni/model_executor/models/qwen3_tts/
-          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
-          - vllm_omni/model_executor/models/common/snake_activation.py
-          - vllm_omni/model_executor/models/common/alias_free_activation.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
@@ -282,9 +272,6 @@ steps:
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen3_tts_base.py
          - vllm_omni/model_executor/models/qwen3_tts/
-          - vllm_omni/model_executor/models/common/qwen3_code_predictor.py
-          - vllm_omni/model_executor/models/common/snake_activation.py
-          - vllm_omni/model_executor/models/common/alias_free_activation.py
          - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
          - vllm_omni/deploy/qwen3_tts.yaml
        commands:
@@ -297,11 +284,52 @@ steps:
            "
        mirror_hardwares: l4_1

+      - label: "TTS · Higgs-Audio-V3 Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_higgs_audio_v3.py
+          - vllm_omni/model_executor/models/higgs_audio_v3/
+          - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
+          - vllm_omni/deploy/higgs_multimodal_qwen3.yaml
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_USE_DEEP_GEMM=0
+              export VLLM_MOE_USE_DEEP_GEMM=0
+              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
+            "
+        mirror_hardwares: l4_1
+
+      - label: "Diffusion · Qwen Image Edit Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_edit.py
+          - vllm_omni/diffusion/models/qwen_image/
+        commands:
+          - |
+            timeout 20m bash -c "
+              pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'core_model' --run-level 'core_model'
+            "
+        mirror_hardwares: h100_1
+
+      - label: "Diffusion · Bagel Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_bagel.py
+          - vllm_omni/model_executor/models/bagel/
+          - vllm_omni/diffusion/models/bagel/
+          - vllm_omni/model_executor/stage_input_processors/bagel.py
+          - vllm_omni/deploy/bagel.yaml
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model'
+            "
+        mirror_hardwares: h100_1
+
      - label: "Diffusion · Wan22 Test"
        source_file_dependencies:
          - tests/e2e/online_serving/test_wan22_t2v.py
          - vllm_omni/diffusion/models/wan2_2/
-          - vllm_omni/diffusion/models/dmd2/
        commands:
          - |
            timeout 40m bash -c "
@@ -314,9 +342,6 @@ steps:
        source_file_dependencies:
          - tests/e2e/online_serving/test_cosmos3.py
          - vllm_omni/diffusion/models/cosmos3/
-          - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py
-          - vllm_omni/experimental/world_models/
-          - vllm_omni/model_extras/cosmos3.py
        commands:
          - |
            timeout 40m bash -c "
@@ -329,8 +354,6 @@ steps:
        source_file_dependencies:
          - tests/e2e/online_serving/test_qwen_image.py
          - vllm_omni/diffusion/models/qwen_image/
-          - vllm_omni/diffusion/models/dmd2/
-          - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py
        commands:
          - |
            timeout 40m bash -c "
@@ -338,3 +361,15 @@ steps:
              pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' --run-level 'core_model'
            "
        mirror_hardwares: h100_1
+
+      - label: "Diffusion · Qwen Image Layered Test"
+        source_file_dependencies:
+          - tests/e2e/online_serving/test_qwen_image_layered.py
+          - vllm_omni/diffusion/models/qwen_image/
+        commands:
+          - |
+            timeout 40m bash -c "
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
+            "
+        mirror_hardwares: h100_1
@@ -108,22 +108,6 @@ steps:
            --run-level "full_model"
        mirror_hardwares: h100_1

-      # Own step: the image has no NanoCodec, so this file import-skips out of
-      # the mark-selected TTS steps above. The overrides keep nemo_toolkit[tts]
-      # from downgrading transformers; see
-      # examples/offline_inference/text_to_speech/README.md.
-      - label: "TTS · Gepard-1.0 · L4"
-        timeout_in_minutes: 60
-        commands:
-          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt
-          - uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt
-          - >-
-            pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py
-            -m "tts and cuda"
-            --run-level "full_model"
-        mirror_hardwares: l4_1
-
      - label: "Diffusion · L4"
        timeout_in_minutes: 60
        commands:
@@ -87,9 +87,8 @@ time timeout -k 30 30m docker run \
    echo $ZE_AFFINITY_MASK
    pip install tblib==3.1.0
    cd /workspace/vllm-omni
-    XPU_TEST_PATHS="tests/diffusion tests/dfx tests/e2e"
-    pytest -v -s $XPU_TEST_PATHS -m "core_model and xpu and B60"
+    pytest -v -s -m "core_model and xpu and B60"
    pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
-    pytest -v -s $XPU_TEST_PATHS -m "advanced_model and xpu and B60"
-    pytest -v -s $XPU_TEST_PATHS -m "omni and xpu and B60"
+    pytest -v -s -m "advanced_model and xpu and B60"
+    pytest -v -s -m "omni and xpu and B60"
 '
@@ -46,8 +46,6 @@ htmlcov/
 .cache
 nosetests.xml
 coverage.xml
-coverage-*.xml
-coverage-*.xml.gz
 *.cover
 *.py,cover
 .hypothesis/
@@ -254,7 +252,6 @@ outputs/
 results/
 generated/
 output_*/
-/tests/e2e/accuracy/artifacts/

 # Configuration overrides
 configs/local.yaml
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.27.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.26.0
 FROM ${BASE_IMAGE} AS base

 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
 #    we should swap over to use stable release ASAP.
 #    We should avoid relying on custom commits.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=v0.27.0
+ARG VLLM_VERSION_OR_COMMIT_HASH=v0.26.0
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}

@@ -79,16 +79,9 @@ RUN printf '%s\n' \
    '    os._exit(_exit_code)' \
    > ${COMMON_WORKDIR}/vllm-omni/conftest.py

-# Some optional dependencies pull in the CPU ONNX Runtime distribution. The
-# CPU and ROCm wheels own the same import path and cannot coexist, so remove
-# the CPU wheel and reinstall the ROCm wheel after dependency resolution.
 RUN cd ${COMMON_WORKDIR}/vllm-omni && \
    uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" \
        --no-cache-dir ".[dev]" --no-build-isolation && \
-    python3 -m pip uninstall -y onnxruntime && \
-    python3 -m pip install --no-cache-dir --force-reinstall --no-deps \
-        "onnxruntime-rocm>=1.22.2" && \
-    python3 -c "import importlib.metadata as m; import onnxruntime as ort; assert 'ROCMExecutionProvider' in ort.get_available_providers(), ort.get_available_providers(); assert m.version('onnxruntime-rocm'); assert not any(d.metadata['Name'] == 'onnxruntime' for d in m.distributions())" && \
    python3 -c "from torchcodec.decoders import VideoDecoder"

 RUN ln -sf /usr/bin/python3 /usr/bin/python
@@ -60,11 +60,7 @@ nav:
      - Overview and Compatibility:
        - Overview: user_guide/diffusion_features.md
        - Feature Compatibility: user_guide/feature_compatibility.md
-      - CPU Offloading:
-        - Overview: user_guide/diffusion/cpu_offload.md
-        - Model-Level: user_guide/diffusion/offloader/cpu_offload.md
-        - Layerwise: user_guide/diffusion/offloader/layerwise_offload.md
-        - Distributed Layerwise: user_guide/diffusion/offloader/distributed_layerwise_offload.md
+      - CPU Offloading: user_guide/diffusion/cpu_offload.md
      - Cache Acceleration:
        - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
        - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
@@ -76,13 +72,7 @@ nav:
        - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
        - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
        - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md
-      - Attention Backends:
-        - Overview: user_guide/diffusion/attention_backends.md
-        - Dense Backends: user_guide/diffusion/attention_backends/dense_backends.md
-        - TRTLLM: user_guide/diffusion/attention_backends/trtllm.md
-        - SageAttention: user_guide/diffusion/attention_backends/sage.md
-        - Hugging Face Hub: user_guide/diffusion/attention_backends/huggingface_hub.md
-        - RainFusion: user_guide/diffusion/attention_backends/rainfusion.md
+      - Attention Backends: user_guide/diffusion/attention_backends.md
      - Regional Compilation: user_guide/diffusion/regional_compilation.md
      - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
      - Startup and Loading: user_guide/diffusion/startup_and_loading.md
@@ -135,16 +125,11 @@ nav:
          - Tensor Parallel: design/feature/tensor_parallel.md
          - VAE Patch Parallelism: design/feature/vae_parallel.md
        - Attention Optimization:
-          - Backend Selection: design/feature/attention_backend_selection.md
          - Skip-Softmax: design/feature/skip_softmax.md
-        - CPU Offloading:
-          - Overview: design/feature/offloader/README.md
-          - Model-Level: design/feature/offloader/cpu_offload.md
-          - Layerwise: design/feature/offloader/layerwise_offload.md
-          - Distributed Layerwise: design/feature/offloader/distributed_layerwise_offload.md
        - Cache-DiT: design/feature/cache_dit.md
        - TeaCache: design/feature/teacache.md
        - Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
+        - Distributed Layerwise Offload: design/feature/distributed_layerwise_offload.md
    - Infrastructure and Performance:
      - Prometheus Metrics: design/metrics.md
      - Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md
@@ -327,99 +327,6 @@ On PR builds, `upload_pipeline.py` logs `skip '…' (no changes under …)` for
 - Use **groups** for dashboard readability; keep **E2E Test** as the group name CUDA diff filtering expects for `--e2e` nightly runs on main.
 - Prefix labels by model domain: **Omni ·**, **TTS ·**, **Diffusion ·**, **Simple ·**, etc., matching existing steps.

-### Per-model coverage
-
-Pilot (v1): one Merge-tier (L3) job in `.buildkite/cuda/test-merge.yml` uploads
-per-model, per-entry-mode coverage as Buildkite artifacts — "TTS · Qwen3-TTS Base
-Test". It runs on a single GPU so the pilot is cheap to reproduce. (A second
-diffusion pilot, "Diffusion · Bagel Test", was retired from merge; reintroduce
-another diffusion coverage job if a second pilot is needed again.) This is a
-pilot, not a rollout: no dashboard/visualization lives in this repo, and nothing
-is gated on coverage.
-
-**Naming convention**: `coverage-<model_id>-<mode>-<step_id>.xml.gz`, where
-`<model_id>` is the model's directory name under
-`vllm_omni/diffusion/models/<model_id>/` or
-`vllm_omni/model_executor/models/<model_id>/` (e.g. `z_image`, `qwen3_omni`),
-`<mode>` is `online` or `offline`, and `<step_id>` is `BUILDKITE_STEP_ID`
-(`local` outside Buildkite). `gunzip` before feeding a report to a Cobertura
-consumer.
-
-Reports are gzipped because `--cov=vllm_omni` sets coverage's `source`, which
-makes it walk the package and emit one `<line>` element per statement for every
-file — including the ones a single model's tests never import. That inventory,
-not the coverage, is the size: a Bagel report is 7.1 MB, of which 96.5% is
-`<line>` elements and 77% belongs to the 840 of 1169 files that run never
-imported. A run covering zero lines produces the same 7 MB. gzip takes it to
-roughly 410 KB and changes nothing about the data.
-
-`<step_id>` is what keeps two steps that cover the same model apart. Artifacts are
-scoped to the build, and a nightly build carries the ready, merge and nightly tiers
-at once, so several same-model steps land in one namespace — `qwen3_tts` already has
-both a Base and a CustomVoice job, and `minicpmo_4_5` a base and a duplex one. Two
-uploads on one path do not fail loudly: an exact-name download can report the path
-as ambiguous, and a glob download fetches every match concurrently and renames each
-onto the same destination, so whichever finishes last wins. The step id remains the
-logical step identifier across retries, while Buildkite's default artifact lookup
-selects only the latest attempt — `BUILDKITE_JOB_ID` would instead give every
-attempt its own filename, leaving a failed attempt's partial report to be picked up
-alongside the real one. A step using `parallelism`/matrix would need
-`BUILDKITE_PARALLEL_JOB` added, since its jobs share a step id; neither pilot is
-one.
-
-**Opting in a new model**: replace the job's combined `pytest` command with
-[`run_cov_split.sh`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/run_cov_split.sh),
-which runs each mode, names the reports, uploads them, and fails the step if
-either half or the upload failed:
-
-```yaml
-commands:
-  - |
-    .buildkite/common/scripts/run_cov_split.sh \
-      --model-id <model_id> \
-      --offline <offline test path> \
-      --online <online test path> \
-      -- <the job's existing pytest flags, e.g. -m '...' --run-level '...'>
-```
-
-Everything after `--` is forwarded verbatim to both runs with quoting preserved,
-so an expression like `-m 'advanced_model and cuda'` survives intact and a job can
-pass anything else it needs. `--offline`/`--online` are repeatable and only one is
-required, so a job with tests in a single mode runs just that half. Runtime is
-bounded solely by the step's `timeout_in_minutes` — raise it, since the split
-loads the model once per mode.
-
-Before splitting, confirm each half actually collects at least one test under the job's forwarded filters — pytest exits with code 5 ("no tests collected") if a half is empty, which the script reports as a failure, red-ing a job that used to pass as one combined invocation.
-
-Also check the tests' `num_cards` against the job's `mirror_hardwares` preset.
-`cuda_marks` turns `num_cards > 1` into `skipif(device_count() < num_cards)`, so a
-test asking for more GPUs than the preset provides is silently skipped and its
-coverage artifact comes back empty. Splitting such a job produces a file that
-looks fine and measures nothing — "Diffusion · Z Image Test" is the current
-example (online tests want 4 cards, the job runs on `l4_1`), which is why it is
-not a coverage pilot.
-
-Online-mode coverage depends on the `[tool.coverage.run]` `patch`/`sigterm`
-settings in `pyproject.toml`: online tests launch the server with
-`subprocess.Popen` and stop it with SIGTERM, and without those settings the XML
-reflects only the pytest parent process, not the server's code paths.
-
-The upload depends on `buildkite-agent` being callable inside the container. The
-`kubernetes` presets (`h100_*`, `*_npu_*`) provide it; the `docker` ones (`l4_*`)
-only do because they set `mount-buildkite-agent: true`. A new docker preset that
-runs a coverage job needs the same.
-
-List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's
-`source_file_dependencies` — both change what the job measures, so without them a
-change there is filtered out of normal PR builds and only surfaces in a later
-nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with
-a synthetic job (it does not pin real merge labels).
-Editing only the surrounding CI YAML still does not schedule them, so a PR that
-touches just the wiring needs a full E2E run (or the commands run on a GPU host)
-to produce artifacts. When checking a new model's
-artifacts, compare `lines-covered` between the online and offline XML rather than
-just confirming both files exist.
-
 ### Validation checklist

 | Check | Command / location |
@@ -406,7 +406,7 @@ the source for configuration, compatibility, and implementation details.
  stage contract.
 * **Quantization and memory efficiency:** [Quantization](feature/quantization.md)
  resolves per-pipeline or per-component quantization configurations, while
-  [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
+  [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)
  streams diffusion blocks from host memory within the existing parallel
  topology.

@@ -117,22 +117,7 @@ After (async D2H):

 ### Batch Split

-When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently. If `OUTPUT_READY` arrives before `execute_batch` registers the split map, `execute_batch` directly adopts the early output.
-
-### Reliability, Lifecycle & Timeout Behavior
-
-1. **Drain Before Memory Release (`drain_async_outputs` / `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S`)**:
-   When workers transition into sleep states (`handle_sleep_task`) or invoke memory-releasing methods, `WorkerProc.drain_async_outputs()` blocks up to `_ASYNC_OUTPUT_DRAIN_TIMEOUT_S` (10.0s) until all in-flight background D2H/SHM packing tasks finish. This acts as a synchronization barrier preventing worker sleep cycles from releasing device tensors while the side CUDA stream is still actively reading them. Draining typically completes within 10–50 ms under normal operation.
-
-2. **Timeout Diagnostics (`_ASYNC_OUTPUT_TIMEOUT`)**:
-   The engine waits for `wait_output_ready()` with a timeout of `_ASYNC_OUTPUT_TIMEOUT` (30.0s). If a timeout occurs, the executor logs detailed diagnostic bookkeeping (pending futures, cached outputs, and batch split maps) to identify whether the stall occurred in the worker, pump, or waiter.
-
-3. **Queue Serialization & Atomic Future Resolution**:
-   - `result_mq` writes are serialized with `_result_mq_lock` to protect the single-writer `MessageQueue` against concurrent writes from the worker main loop (`COMPUTE_DONE`) and background thread (`OUTPUT_READY`).
-   - `ResultPumpThread` unpacks shared-memory payloads first, and then atomically resolves waiting futures or populates `_completed_outputs` under `_futures_lock`, avoiding race conditions with concurrent `wait_output_ready()` invocations.
-
-4. **Storage-Aware IPC Packing**:
-   `pack_diffusion_output_shm()` evaluates `max(view_bytes, storage_bytes)` against `_SHM_TENSOR_THRESHOLD` (1 MB) to ensure that per-request slices sharing a large batch storage are moved to POSIX shared memory rather than serialized through pickle over the MessageQueue inline buffer.
+When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.

 ## Model Coverage

@@ -146,7 +131,6 @@ All models running in request-mode (`step_execution=False`, the default) automat
 |---|---|---|
 | **HunyuanImage-3.0** | Image | `False` |
 | **Qwen-Image** | Image | `True` |
-| **LTX-2.3** | Video / Audio | `False` |

 Other models with `step_execution=False` are also supported but not yet verified.

@@ -1,89 +0,0 @@
-# Diffusion Attention Backend Selection
-
-This document defines the selection and extension contract for diffusion
-attention backends. User-facing backend choices, installation, and tuning are
-in the [attention backend guides](../../user_guide/diffusion/attention_backends.md).
-
-## Scope
-
-The contract applies to DiT and other diffusion attention layers. It is
-separate from vLLM's autoregressive attention selector.
-
-The selection path has four responsibilities:
-
-1. normalize user configuration into `AttentionConfig` and `AttentionSpec`;
-2. resolve a spec for an attention role;
-3. ask the active platform to validate an explicit backend or choose a
-   hardware default; and
-4. load the selected `AttentionBackend` class from the registry path.
-
-## Resolution contract
-
-`get_attn_backend_for_role()` resolves configuration in this order:
-
-1. exact `per_role[role]`;
-2. category `per_role[role_category]`;
-3. `default`;
-4. platform default.
-
-An explicit resolution returns both the backend class and its
-`AttentionSpec`. A platform-default resolution returns the class and `None`.
-Layers must therefore treat the spec as optional and must not infer that a
-missing spec means SDPA.
-
-Class resolution is cached by backend name, head size, and whether TRTLLM may
-be selected automatically. Logging is cached separately by role and source so
-the startup record identifies why each role received its backend.
-
-## Registry and platform boundary
-
-`DiffusionAttentionBackendEnum` maps stable configuration names to default
-qualified class paths. `register_diffusion_backend()` may replace a path at
-runtime without changing the public enum value.
-
-The active `OmniPlatform` owns compatibility policy through
-`get_diffusion_attn_backend_cls()`. It must:
-
- validate explicit selections and fail with an actionable error when the
-  requested kernel cannot run;
- choose only an available, compatible default when no backend is explicit;
- account for head size and other platform-visible constraints; and
- return a qualified class path implementing `AttentionBackend`.
-
-The selector must not duplicate device capability or package-availability
-policy that belongs to the platform.
-
-## Typed backend options
-
-Backend-specific settings remain typed fields on `AttentionSpec` rather than
-unstructured keyword dictionaries:
-
- `quant` is consumed by FlashInfer and TRTLLM with backend-specific value
-  validation;
- `skip_softmax` is consumed by TRTLLM; and
- `block_sparse` is consumed by block-sparse backends such as RainFusion.
-
-A backend reads only the fields it owns and rejects incompatible values. New
-options should be added to a shared typed spec only when more than one backend
-shares their semantics; otherwise add a dedicated typed block.
-
-## Model integration contract
-
-Each diffusion attention call declares a stable role and, when appropriate, a
-broader category. Model-specific roles permit precise overrides; categories
-preserve common `self` or `cross` policies. A model also decides whether its
-path is compatible with automatic TRTLLM selection, because only the model
-knows whether masking and packing satisfy the kernel contract.
-
-## Adding or changing a backend
-
-An implementation change is complete only when it:
-
-1. implements the `AttentionBackend` interface and exposes a stable enum name;
-2. adds platform validation and default routing where appropriate;
-3. defines and validates any typed configuration;
-4. tests explicit selection, default selection, and incompatible paths;
-5. documents installation, fallback behavior, and quality implications in the
-   corresponding user guide; and
-6. updates the overview matrix without moving algorithm details back into the
-   landing page.
@@ -0,0 +1,207 @@
+# Distributed Layerwise Offload
+
+This document describes distributed layerwise offload (DLO) for diffusion
+models. DLO keeps only a small number of DiT blocks on the accelerator and
+streams the remaining blocks from host memory. The distributed backend can
+either shard those host-side weights across an existing parallel group or keep
+the standard loader's rank-local weights and avoid an additional collective.
+The rank-local mmap storage design is tracked in
+[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+
+For user-facing commands, see the
+[distributed layerwise offloading guide](../../user_guide/diffusion/cpu_offload.md)
+and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
+
+## Status
+
+DLO is implemented for multi-device diffusion execution. The default
+AllGather path is the primary path for DP and SP deployments. The
+`--dlo-no-use-allgather` path is a rank-local compatibility mode: it is useful
+for standard-loader sharding, workstation bring-up, and systems where an
+additional DLO collective is undesirable. For explicitly mmap-compatible
+models in a TP1, non-HSDP topology, this path shares file-backed DiT checkpoint
+pages across processes and keeps only bounded staging buffers private. TP,
+HSDP, and other models retain standard-loader rank-local host storage.
+
+The compatibility matrix below describes the current implementation. The
+unit-level guards are covered, but not every parallelism combination has a
+full model-and-hardware end-to-end test.
+
+## Design
+
+### DLO consumes the existing parallel topology
+
+DLO does not create a new DP, TP, or SP topology. It reads the configured
+`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
+standard distributed groups have been initialized.
+
+The DLO weight-sharding group is selected as follows:
+
+1. Use the existing DP group when `data_parallel_size > 1`.
+2. When DP is one and SP is greater than one, use the SP group.
+3. Otherwise, run rank-locally without a DLO process group.
+
+TP is deliberately not used as DLO's AllGather group. HSDP has its own
+parameter-sharding lifecycle and is not allowed to be sharded a second time by
+DLO's AllGather path.
+
+### AllGather path
+
+With the default `dlo_use_allgather=True`, each rank stores approximately
+`1 / group_size` of each streamable block in pinned host memory. The next
+block's shard is copied to a device buffer and reconstructed with
+`all_gather_into_tensor` on a communication stream while the current block is
+executing.
+
+Multi-rank AllGather behavior is unchanged. When the effective DLO group size
+is one, the backend now uses rank-local mmap staging; this closes a previous
+loader/backend mismatch that could leave parameters on the `meta` device.
+
+```text
+Compute:    [Block N]             [Block N+1]          [Block N+2]
+H2D:                      [shard N+1]           [shard N+2]
+AllGather:                [full N+1]             [full N+2]
+Buffers:    [current slot]       [prefetch slot]       [current slot]
+```
+
+![DLO double-buffer prefetch pipeline](../figures/dlo/dlo_pipeline.gif)
+
+The backend uses two shared device buffers, so accelerator weight residency is
+bounded by the largest streamed blocks rather than the complete model.
+
+When DP is greater than one, the engine can process one request per DP rank in
+the same denoising wave. Because AllGather is a collective, all participating
+requests must take the same execution path at every denoising step.
+
+### Rank-local path without DLO AllGather
+
+With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
+one and streams complete runtime blocks using H2D copies only. Host storage is
+selected independently from that transfer protocol:
+
+- For an explicitly mmap-compatible model with TP1 and no HSDP, each worker
+  retains immutable safetensors views backed by the node's shared OS page
+  cache. Two pinned host buffers, sized to the largest streamed block, stage
+  mmap pages for asynchronous H2D. Checkpoint-to-runtime adapters run while a
+  block is packed, so transformed weights do not become a persistent private
+  model copy.
+- With TP, HSDP, online quantization, or a model without mmap support, the
+  regular loader remains responsible for preparing rank-local tensors. This
+  preserves TP/HSDP/quantization loader callbacks.
+
+This mode means:
+
+- DP still provides independent replicas, but DLO does not shard weights
+  across DP ranks.
+- SP still performs its normal activation/attention collectives, but DLO does
+  not shard weights across SP ranks.
+- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
+  DLO's additional weight AllGather is disabled.
+- Pure DP deployments with a compatible mmap checkpoint share the streamed
+  DiT's file-backed pages. Other components and fallback paths may still keep
+  rank-private host weights.
+- The scheduler does not require a synchronized DP request wave for DLO.
+
+The mmap source remains immutable in both transfer modes. AllGather copies a
+persistent `1 / group_size` shard from it and then closes the mappings;
+rank-local mode keeps the mappings open and copies a complete block through
+bounded staging without adding a collective.
+
+The current implementation packs mmap views into a staging slot synchronously
+on the host; the subsequent pinned H2D transfer overlaps with block compute.
+This keeps host memory bounded but makes CPU packing bandwidth part of the
+rank-local path's performance profile.
+
+## Parallelism compatibility
+
+| Parallelism | DLO + AllGather | DLO without AllGather |
+|---|---|---|
+| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. DP replicas remain independent; no DLO weight collective. Compatible TP1 models share file-backed DiT pages across replicas. |
+| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
+| **TP > 1** | Unsupported in the DLO mmap path. TP-aware loader callbacks are bypassed, so the backend rejects this configuration when it enters that path. | Standard loading is retained and TP-local tensors can be streamed. This is the intended compatibility path, but it still needs broader model and hardware validation. |
+| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
+
+### Combined dimensions
+
+- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
+  than one; SP continues to use its own sequence-parallel group. If DP is one,
+  the SP group becomes DLO's sharding group in AllGather mode.
+- **DP + TP/SP without AllGather:** standard model loading defines the
+  rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
+  collective.
+- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
+  DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
+  materialization and synchronization.
+- **HSDP + DP or TP:** rejected independently by the diffusion parallel
+  configuration.
+
+## Request and loading constraints
+
+AllGather DP multi-concurrency requires:
+
+- explicit `num_inference_steps`;
+- the same `num_inference_steps` for all requests in a wave; and
+- identical request arguments that affect the collective execution path.
+
+The no-AllGather path does not impose these DLO-specific synchronized-wave
+requirements.
+
+The mmap loader is used when the model explicitly declares a compatible
+checkpoint layout. It is independent of AllGather for TP1/non-HSDP execution.
+Models can declare a per-parameter checkpoint-to-runtime adapter; the adapter
+must preserve dtype and element count so it can be applied while packing one
+staging block. MiniMax-H3 uses this mechanism for grouped QKV. Online
+quantization and TP/HSDP no-AllGather execution use the regular loader path.
+
+### Adding mmap support to another model
+
+A new pipeline should opt in only after its ordinary loader contract can be
+represented without materializing a private full-model copy:
+
+1. Populate `weights_sources` with safetensors component sources and prefixes
+   that place every DiT checkpoint key in the pipeline parameter namespace.
+2. If prefixed checkpoint keys already equal runtime parameter names, declare
+   `_supports_mmap_loading = True`. If names differ, provide
+   `_remap_ckpt_key` to map checkpoint names to runtime names (or return
+   `None` for intentionally ignored keys).
+3. If a parameter needs a TP1 layout conversion, attach a callable
+   `mmap_weight_transform` to that parameter during model construction. The
+   transform must preserve dtype and element count because it runs while one
+   bounded block is packed.
+4. Verify that mmap and the regular loader produce the same runtime tensors,
+   that every expected DiT tensor maps successfully, and that TP, HSDP, and
+   online-quantized configurations retain their standard-loader behavior.
+
+Setting `_supports_mmap_loading = True` alone is not sufficient: explicit
+opt-ins fail during initialization if any expected DiT tensor remains on the
+`meta` device after checkpoint mapping.
+
+## Validation coverage
+
+Current source-level validation includes:
+
+- HSDP + DLO + AllGather rejection;
+- HSDP + DLO without AllGather acceptance at configuration level;
+- TP rejection in the DLO+AllGather mmap path;
+- resident-layer requests requiring no-AllGather;
+- rank-local mmap source retention, layout transforms, and bounded host
+  staging;
+- DP request-wave validation for denoising-step compatibility;
+- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
+  tests.
+
+The highest-value missing coverage is end-to-end numerical comparison against
+ordinary layerwise offload for DP+SP, TP+no-AllGather, and HSDP+SP+no-AllGather
+on the target CUDA/NCCL or CANN/HCCL hardware.
+
+## Recommendations
+
+- Use **DP + DLO AllGather** for the supported throughput and host-memory
+  scaling path.
+- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
+  not the goal.
+- Use **no-AllGather** when replicas must avoid DLO collectives. TP1 models with
+  mmap support can share DiT checkpoint pages; TP/HSDP and unsupported models
+  retain the higher-memory standard-loader path.
+- Prefer **HSDP alone** for production HSDP deployments until the combined
+  HSDP + DLO no-AllGather path has broader end-to-end coverage.
@@ -1,80 +0,0 @@
-# CPU Offloading
-
-This document defines the shared architecture for diffusion CPU offloading.
-The three strategies have separate user and design pages:
-
-| Strategy | User guide | Design contract |
-| --- | --- | --- |
-| Model-level | [Guide](../../../user_guide/diffusion/offloader/cpu_offload.md) | [Design](cpu_offload.md) |
-| Layerwise | [Guide](../../../user_guide/diffusion/offloader/layerwise_offload.md) | [Design](layerwise_offload.md) |
-| Distributed layerwise | [Guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md) | [Design](distributed_layerwise_offload.md) |
-
-## Strategy selection
-
-`OffloadConfig.from_od_config()` converts diffusion configuration into one
-`OffloadStrategy`. If callers set multiple flags, selection is deterministic:
-
-1. `DISTRIBUTED_LAYERWISE`;
-2. `LAYERWISE`;
-3. `MODEL_LEVEL`; and
-4. `NONE`.
-
-The factory logs when a higher-priority strategy disables another. It derives
-parallel and HSDP state from `DiffusionParallelConfig`; callers do not provide
-a separate offload group size.
-
-`get_offload_backend()` then validates platform offload support, resolves the
-device, and creates exactly one backend. Returning `None` means offloading is
-disabled or unsupported and must not leave partially installed hooks.
-
-## Shared lifecycle
-
-Every backend implements `OffloadBackend`:
-
- `enable(pipeline)` discovers modules, establishes initial residency, and
-  installs hooks;
- `disable()` removes owned hooks and resources; and
- `is_enabled()` reports lifecycle state.
-
-`disable()` does not promise to restore the pipeline's original device
-placement. The caller owns any subsequent rematerialization.
-
-Hooks are registered through `HookRegistry` and `ModelHook`; offload backends
-must use distinct hook names and remove only hooks they own.
-
-## Discovery and topology
-
-Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
-
- `_dit_modules`;
- `_encoder_modules`;
- `_vae_modules`; and
- `_resident_modules`.
-
-Dotted paths are supported. Legacy pipelines may use the fallback scan of
-well-known attribute names, but new integrations should declare components
-explicitly.
-
-Ordinary layerwise offload obtains topology from declared
-`_layerwise_offload_blocks_attrs` (with a deprecated singular-name
-compatibility path). Distributed layerwise offload may instead use a pipeline
-`OffloadPlan`, then falls back to declared block attributes and its documented
-compatibility discovery. Discovery metadata describes structure only; the
-backend remains responsible for transfer, synchronization, and storage
-ownership.
-
-## Cross-strategy invariants
-
- At most one strategy owns offload hooks for a pipeline.
- A parameter has one authoritative host representation while offloaded.
- Device storage is not freed until dependent compute or transfer work is
-  complete.
- Non-persistent and model-specific buffers remain correct after movement or
-  checkpoint rematerialization.
- Unsupported parallel or loading combinations fail before hooks mutate the
-  model.
- Platform streams, events, synchronization, and cache management go through
-  the vLLM-Omni platform abstraction.
-
-The diffusion [Offloader module design](../../module/diffusion/offloader.md)
-describes how these feature contracts fit into the larger diffusion runtime.
@@ -1,56 +0,0 @@
-# Model-Level Offload
-
-Model-level offload enforces mutual exclusion between pipeline component
-groups. For commands and model support, see the
-[model-level user guide](../../../user_guide/diffusion/offloader/cpu_offload.md).
-
-## Generic hook path
-
-`ModelLevelOffloadBackend.enable()` discovers DiTs, encoders, and VAEs. It
-makes encoders and VAEs device resident for initialization, then installs
-`SequentialOffloadHook` instances:
-
- each DiT hook offloads all encoders and other DiTs before loading itself;
- each encoder hook offloads every DiT before loading itself; and
- VAEs are not part of the mutual-exclusion hook set.
-
-The hook moves parameters and buffers without recursively calling
-`module.to()`. This avoids recursion through modules that retain references to
-other transformer blocks. Host tensors are pinned when configured; DTensors
-are not pinned through the ordinary tensor path.
-
-Before forward returns control to the model, the hook synchronizes the active
-platform so compute cannot race the transfer. XPU uses blocking device-to-host
-copies because its allocator cannot safely evict storage while an asynchronous
-copy is pending.
-
-## Split-model delegation
-
-If a pipeline provides callable `enable_omni_model_cpu_offload`, the backend
-delegates instead of installing generic encoder/DiT hooks. This supports models
-whose mutually exclusive components are nested inside one transformer, such
-as a reasoner and generator.
-
-The delegated pipeline owns its internal contexts, but must preserve the same
-contract:
-
- only the executing component is device resident;
- enable and disable are idempotent;
- transfers use the supplied device, pinning, and HSDP settings; and
- disabling removes all model-local offload hooks or contexts.
-
-The delegation is currently duck typed. If another independent implementation
-is added, formalize the enable/disable pair as a protocol rather than expanding
-attribute-name heuristics.
-
-## Extension invariants
-
- New ordinary pipelines declare both DiT and encoder groups through
-  `SupportsComponentDiscovery`.
- Models with multiple DiTs offload the inactive DiTs as well as encoders.
- Empty or parameterless modules do not trigger transfers.
- Hook removal does not move modules; callers must not assume original
-  residency is restored.
-
-Shared selection and lifecycle behavior is defined in the
-[CPU Offloading design](README.md).
@@ -1,279 +0,0 @@
-# Distributed Layerwise Offload
-
-This document describes distributed layerwise offload (DLO) for diffusion
-models. DLO keeps only a small number of DiT blocks on the accelerator and
-streams the remaining blocks from host memory. The distributed backend can
-either shard those host-side weights across an existing parallel group or keep
-complete rank-local block sources and avoid an additional collective.
-
-For user-facing commands, see the
-[distributed layerwise offloading guide](../../../user_guide/diffusion/offloader/distributed_layerwise_offload.md)
-and the [Cosmos3 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md).
-
-## Status
-
-DLO is implemented for multi-device diffusion execution. The default
-AllGather path is the primary path for DP and SP deployments. The
-`--dlo-no-use-allgather` path streams complete blocks independently and adds no
-DLO weight collective.
-
-Host storage is selected separately from the transfer protocol. The loader can
-produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
-otherwise it uses the ordinary loader. Consequently, no-AllGather replicas on
-the same node can share immutable checkpoint pages when direct mmap is
-selected, while the ordinary-loader fallback still keeps a private runtime
-copy per process.
-
-The Phase A shared-mmap support boundary is TP1. TP greater than one is an
-ordinary-loader compatibility path: DLO can consume the resulting TP-local
-tensors, but those configurations do not use checkpoint mmap and must not be
-used to claim shared-mmap host-memory savings.
-
-The compatibility matrix below describes the current implementation. The
-unit-level guards are covered, but not every parallelism combination has a
-full model-and-hardware end-to-end test.
-
-## Design
-
-### DLO consumes the existing parallel topology
-
-DLO does not create a new DP, TP, or SP topology. It reads the configured
-`DiffusionParallelConfig` and attaches offload hooks to the DiT blocks after the
-standard distributed groups have been initialized.
-
-The DLO weight-sharding group is selected as follows:
-
-1. Use the existing DP group when `data_parallel_size > 1`.
-2. When DP is one and SP is greater than one, use the SP group.
-3. Otherwise, run rank-locally without a DLO process group.
-
-TP is deliberately not used as DLO's AllGather group. HSDP has its own
-parameter-sharding lifecycle and is not allowed to be sharded a second time by
-DLO's AllGather path.
-
-### The loader owns host-weight planning
-
-Before it decides whether ordinary weight materialization can be skipped, the
-diffusion loader builds one `HostWeightPlan`. A direct-checkpoint mmap plan is
-accepted only when preflight proves all of the following:
-
- every required DiT parameter and persistent buffer has exactly one source;
- runtime names, checkpoint keys, shapes, and dtypes match;
- the runtime topology is TP1 without HSDP or online quantization; and
- every custom loader operation is represented by a loader-owned checkpoint
-  adapter.
-
-The exact plan object is handed to DLO. The backend does not rescan checkpoint
-files, repeat the capability decision, or reconstruct names from its block
-topology. If preflight fails, the loader materializes weights normally and DLO
-consumes those runtime tensors.
-
-The plan owns only dedicated DiT component sources. If a pipeline also exposes
-ordinary sources for a text encoder or another non-DiT component, the loader
-still consumes those sources and includes their loaded names in its strict
-coverage check. Only the source prefixes covered by the plan skip ordinary
-materialization. A source that mixes DiT and non-DiT weights fails closed to
-the complete ordinary-loader path because it cannot be skipped safely as a
-unit.
-
-This boundary keeps checkpoint semantics out of DLO and avoids model-pipeline
-flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
-transforms. Model-specific direct-layout knowledge, when required, lives in a
-checkpoint adapter beside the ordinary loader.
-
-### AllGather path
-
-With the default `dlo_use_allgather=True`, each rank stores approximately
-`1 / group_size` of each streamable block in pinned host memory. The next
-block's shard is copied to a device buffer and reconstructed with
-`all_gather_into_tensor` on a communication stream while the current block is
-executing.
-
-```text
-Compute:    [Block N]             [Block N+1]          [Block N+2]
-H2D:                      [shard N+1]           [shard N+2]
-AllGather:                [full N+1]             [full N+2]
-Buffers:    [current slot]       [prefetch slot]       [current slot]
-```
-
-![DLO double-buffer prefetch pipeline](../../figures/dlo/dlo_pipeline.gif)
-
-The backend uses two shared device buffers, so accelerator weight residency is
-bounded by the largest streamed blocks rather than the complete model.
-
-When direct checkpoint mmap is selected, the checkpoint mappings are only the
-source used to prepare each rank's persistent shard. They can be closed after
-shard preparation. Across the AllGather group, those private shards total
-approximately one runtime model copy.
-
-An effective DLO group size of one performs no collective, even when
-`dlo_use_allgather=True`; it follows the rank-local transfer path described
-below.
-
-When DP is greater than one, the engine can process one request per DP rank in
-the same denoising wave. Because AllGather is a collective, all participating
-requests must take the same execution path at every denoising step.
-
-### Rank-local path without DLO AllGather
-
-With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
-one and streams complete blocks using H2D copies only. The host backing may be
-either a loader-approved checkpoint mapping or ordinary runtime tensors.
-
-For direct mmap, each process retains immutable safetensors views and uses two
-bounded pinned host staging slots. Processes on the same node that map the same
-files share physical checkpoint pages through the OS page cache. This removes
-the persistent private full-model copy per pure-DP process, but each process
-still packs and transfers every complete block. Sharing is node-local; each
-node has its own page cache.
-
-When direct mmap preflight fails, the regular model loader remains responsible
-for preparing each rank's weights, including TP-local tensors or HSDP-managed
-parameters. In that fallback, each pure-DP process keeps a private full runtime
-copy.
-
-This mode means:
-
- DP still provides independent replicas, but DLO does not shard weights
-  across DP ranks.
- SP still performs its normal activation/attention collectives, but DLO does
-  not shard weights across SP ranks.
- TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
-  DLO's additional weight AllGather is disabled.
- Pure DP deployments share one checkpoint-backed copy per node when direct
-  mmap is selected; the ordinary-loader fallback keeps one private runtime
-  copy per rank.
- The scheduler does not require a synchronized DP request wave for DLO.
-
-## Parallelism compatibility
-
-| Parallelism | DLO + AllGather | DLO without AllGather |
-|---|---|---|
-| **DP** | Supported primary path. DLO shards host weights across the DP group and can run DP multi-concurrency. | Supported rank-local path. Compatible TP1 replicas can share checkpoint pages on each node; fallback runtime tensors remain private. |
-| **SP** | Supported in the implementation. With DP=1, DLO uses the SP group for host-weight sharding; SP still shards sequence/activation work. | SP remains active, but DLO keeps standard-loader rank-local weights and adds no SP weight collective. |
-| **TP > 1** | Outside the Phase A shared-mmap support scope. The loader falls back before mutation, preserves TP-local layouts, and DLO may apply DP/SP host sharding to those ordinary runtime tensors. | Outside the Phase A shared-mmap support scope. The ordinary TP-aware loader produces rank-local tensors, which DLO streams without an additional weight collective; DP replicas retain private runtime storage. |
-| **HSDP** | Rejected. HSDP has already sharded parameters, so DLO AllGather would double-shard them. | Accepted by configuration. HSDP owns parameter sharding and its own gathers; DLO only stages rank-local parameters. End-to-end coverage is limited. |
-
-### Combined dimensions
-
- **DP + SP:** DLO uses the DP group for weight sharding when DP is greater
-  than one; SP continues to use its own sequence-parallel group. If DP is one,
-  the SP group becomes DLO's sharding group in AllGather mode.
- **DP + TP/SP without AllGather:** standard model loading defines the
-  rank-local tensor layout. DLO adds no cross-DP, cross-TP, or cross-SP weight
-  collective.
- **HSDP + SP:** the general parallel configuration permits HSDP over SP, but
-  DLO must use `--dlo-no-use-allgather`. HSDP remains responsible for weight
-  materialization and synchronization.
- **HSDP + DP or TP:** rejected independently by the diffusion parallel
-  configuration.
-
-## Request and loading constraints
-
-AllGather DP multi-concurrency requires:
-
- explicit `num_inference_steps`;
- the same `num_inference_steps` for all requests in a wave; and
- identical request arguments that affect the collective execution path.
-
-The no-AllGather path does not impose these DLO-specific synchronized-wave
-requirements.
-
-Direct checkpoint mmap can back either transfer path. It is currently limited
-to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
-ordinary loader. Online quantization remains incompatible with DLO AllGather;
-use `--dlo-no-use-allgather` or disable online quantization.
-
-A normalized runtime mmap cache, built through the ordinary loader, is the
-proposed general mechanism for sharing transformed TP or quantized layouts.
-That cache and its publication/lifecycle protocol are intentionally outside
-this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
-
-## Validation coverage
-
-Current source-level validation includes:
-
- HSDP + DLO + AllGather rejection;
- HSDP + DLO without AllGather acceptance at configuration level;
- loader preflight fallback for TP, HSDP, online quantization, unknown custom
-  loaders, missing keys, and shape/dtype mismatches;
- exact loader-to-backend plan transfer and ordinary-loader fallback;
- rank-local mmap source retention, bounded two-slot staging, and adapter
-  transforms without parameter-side flags;
- resident-layer requests requiring no-AllGather;
- DP request-wave validation for denoising-step compatibility;
- sharding, double-buffer, AllGather-size, and heterogeneous-block regression
-  tests.
-
-### B300 parallel-topology smoke matrix
-
-A four-GPU B300 smoke test covered MiniMax-H3 FL2VA with the same prompt, seed,
-CUDNN attention backend, 256x256 output, two denoising steps, and
-`dlo_resident_layers=0`. The TP2 rows used DiT DP2xTP2 with the text encoder and
-VAEs at TP1. They validate the ordinary-loader fallback only, not direct mmap
-or shared-mmap host-memory savings.
-
-| Configuration | Result | Warm E2E | Peak device memory | Host PSS |
-|---|---:|---:|---:|---:|
-| DP4xTP1 AllGather | Passed, 4 concurrent requests | 2.87 s / 4 requests | 13.84 GiB | 211.99 GiB |
-| DP4xTP1 no-AllGather | Passed, 1 request | 15.02 s | 13.23 GiB | 187.77 GiB |
-| DP2xTP2 AllGather | Passed, 2 concurrent requests | 4.16 s / 2 requests | 12.50 GiB | 211.97 GiB |
-| DP2xTP2 no-AllGather | Passed, 1 request | 3.51 s | 11.88 GiB | 314.01 GiB |
-
-Within each topology, the AllGather and no-AllGather video and audio outputs
-were byte-identical. All four runs completed without an `ERROR` or traceback
-and released their device allocations. For DP4xTP1, no-AllGather direct mmap
-reduced total PSS by 24.22 GiB (11.4%) and `Private_Dirty` from 211.33 to
-125.32 GiB (40.7%) relative to AllGather. For DP2xTP2, preflight selected the
-ordinary loader as designed; no-AllGather PSS was 314.01 GiB, about 48% above
-AllGather, because DP replicas did not share checkpoint-backed runtime
-weights. This is a functional and memory smoke test, not a production-quality
-performance or output-quality benchmark.
-
-### Host-memory measurement
-
-A two-worker MiniMax-H3 FL2VA measurement on one L20X node compared the
-ordinary-loader fallback with direct mmap. Both runs used
-DP=2, TP=1, no DLO AllGather, BF16 weights, two denoising steps, and a
-256x256 four-second request. The ordinary-loader workers were sampled after
-initialization. The mmap workers were sampled after one completed request, so
-the checkpoint working set had been faulted into the page cache; this is the
-more conservative point for mmap.
-
-The values below come from `/proc/<worker>/smaps_rollup` and include the whole
-worker, not only the DiT. The stable rank-to-rank difference comes from other
-pipeline components, so each worker should be compared with the same worker in
-the other storage mode.
-
-| Worker | Ordinary RSS | mmap RSS | Ordinary PSS | mmap PSS | PSS reduction |
-|---|---:|---:|---:|---:|---:|
-| DP worker 0 | 168.27 GiB | 132.76 GiB | 167.84 GiB | 101.43 GiB | 66.40 GiB |
-| DP worker 1 | 116.19 GiB | 79.97 GiB | 115.73 GiB | 48.64 GiB | 67.09 GiB |
-| **Two-worker total** | — | — | **283.56 GiB** | **150.08 GiB** | **133.48 GiB (47.1%)** |
-
-The direct-mmap workers each reported 62.45 GiB `Shared_Clean` but only
-31.20 GiB `Pss_File`, which is the proportional charge expected when the same
-resident checkpoint pages are mapped by two workers. `Private_Dirty` also fell
-from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
-1, a reduction of about 97–98 GiB per worker. RSS understates this benefit
-because it counts a shared physical page in every process that maps it; summed
-PSS is the appropriate node-memory comparison.
-
-The highest-value missing coverage is broader end-to-end numerical and
-lifecycle comparison against ordinary layerwise offload for DP+SP,
-HSDP+SP+no-AllGather, and TP greater than one across additional models and
-target CUDA/NCCL or CANN/HCCL hardware. That broader TP coverage does not
-change the Phase A direct-mmap TP1 support boundary.
-
-## Recommendations
-
- Use **DP + DLO AllGather** for the supported throughput and host-memory
-  scaling path.
- Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
-  not the goal.
- Use **no-AllGather** when independent replica execution is required. TP1
-  direct-mmap deployments can share checkpoint pages per node; other layouts
-  retain the ordinary loader's private host memory behavior and are outside
-  the Phase A shared-mmap support scope.
- Prefer **HSDP alone** for production HSDP deployments until the combined
-  HSDP + DLO no-AllGather path has broader end-to-end coverage.
@@ -1,57 +0,0 @@
-# Layerwise Offload
-
-Layerwise offload streams transformer blocks through one accelerator. For
-commands and model support, see the
-[layerwise user guide](../../../user_guide/diffusion/offloader/layerwise_offload.md).
-
-## Hook ring
-
-`LayerWiseOffloadBackend` discovers every streamable block and registers one
-`LayerwiseOffloadHook` per block. Hooks form a ring: each hook owns the pinned
-host representation of its registered block and a reference to the next block
-it must prefetch; the final hook prefetches the first block for the next
-denoising iteration.
-
-During initialization, parameters and buffers are grouped by dtype and copied
-into contiguous pinned CPU tensors. Original tensor storage is replaced with
-an empty or meta placeholder while identity and module structure are retained.
-Metadata records names, offsets, shapes, and dtypes for rematerialization.
-
-## Forward lifecycle
-
-For each block:
-
-1. the previous hook's completion event ensures this block is materialized;
-2. the copy stream asynchronously allocates and fills the next block's device
-   tensors;
-3. parameter and buffer storage is rebound to views of those tensors;
-4. the current block computes on the normal stream; and
-5. its post-forward hook replaces device storage with placeholders after all
-   dependent work is safe.
-
-The copy stream waits on the compute stream before reusing storage, and the
-compute stream waits on the per-block prefetch event before forward. These
-dependencies, not incidental global synchronization, define correctness.
-
-## Topology contract
-
-A DiT declares ordered block-container attributes using
-`_layerwise_offload_blocks_attrs` or an `OffloadPlan`. Each resolved item must
-be an executable `nn.Module`; ordering must match forward execution. Multiple
-containers are concatenated in execution order.
-
-Non-block DiT modules, encoders, VAEs, and declared resident modules are not
-streamed by this backend.
-
-## Invariants and limitations
-
- Each offloaded tensor has one pinned host backing store.
- Parameter identity is preserved when storage is rebound.
- Blocks with heterogeneous shapes and dtypes use their recorded metadata;
-  implementations must not assume identical blocks.
- The ring is initialized before the first forward so block zero is available.
- This backend is single-device; distributed weight reconstruction belongs to
-  [Distributed Layerwise Offload](distributed_layerwise_offload.md).
-
-Shared selection and lifecycle behavior is defined in the
-[CPU Offloading design](README.md).
@@ -76,12 +76,6 @@ equivalent. `auto-round` and `auto_round` use the Omni INC/AutoRound adapter.
 ModelOpt checkpoint metadata is detected from its method, producer metadata, or
 `quant_algo`, then mapped to the corresponding vLLM ModelOpt configuration.

-Unless the input is a per-component dictionary, the resolved configuration is
-passed unchanged to every quantization-aware component constructed by the
-pipeline. It affects only layers supported by that component's quantization
-method; it does not rewrite arbitrary `torch.nn` modules. A per-component
-dictionary narrows this global scope through runtime layer-prefix routing.
-
 The factory uses Omni-specific overrides for methods whose diffusion behavior
 or platform dispatch is not provided by the vLLM registry. The current override
 set includes INT8, BitsAndBytes, MXFP8, MXFP4, MXFP4 dual-scale, and INC/
@@ -160,11 +154,9 @@ quantizable vLLM layers. Each method configuration is responsible for:
 Model code must keep precision-sensitive or unsupported modules unquantized.
 The shared `safe_quant_config()` helper is used by model integrations for
 normalization and modulation layers where a generic FP8 configuration would be
-unsafe. Encoders exposed through the quantization factory receive a global
-configuration just like the main transformer, but only their supported vLLM
-quantizable layers are affected. Ordinary `torch.nn` encoder layers, VAEs,
-schedulers, and tokenizers are not rewritten automatically. A component map can
-select a narrower scope, but it cannot add quantization support to a layer.
+unsafe. Text encoders, audio/vision encoders, VAEs, schedulers, and tokenizers
+are not implicitly quantized just because the main transformer is quantized;
+they require explicit component routing and method support.

 Weight quantization and quantized KV cache are separate paths. The former is
 resolved through `quantization_config`; runtime attention/KV-cache quantization
@@ -194,7 +186,7 @@ Distributed layerwise offload has an additional boundary: its sharded mmap
 AllGather path is not compatible with online quantization because the path
 expects already materialized checkpoint weights. Use the rank-local
 `--dlo-no-use-allgather` path or a compatible pre-quantized checkpoint when
-combining these features. See the [DLO design document](offloader/distributed_layerwise_offload.md)
+combining these features. See the [DLO design document](distributed_layerwise_offload.md)
 for the full compatibility matrix.

 ## Adding a quantization backend
@@ -2,10 +2,8 @@

 Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
 and how to pick an operating point — is in
-[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax).
-The shared selector contract is documented in
-[Diffusion Attention Backend Selection](attention_backend_selection.md). This
-page explains the algorithm.
+[Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). This page explains
+the algorithm.

 ## Motivation

@@ -53,22 +53,15 @@ implementation contract; it is not, by itself, a general support claim.
 #### Attention optimization

 The [Diffusion Attention Backends](../user_guide/diffusion/attention_backends.md)
-guides list selectable backends, platform defaults, installation, and tuning.
-The design contracts separate selection mechanics from backend algorithms:
+guide lists all selectable backends and their platform defaults. Skip-Softmax
+is the backend-related optimization with a standalone feature-design contract:

- [Attention Backend Selection](feature/attention_backend_selection.md)
 - [Skip-Softmax](feature/skip_softmax.md)

-#### CPU offloading
-
- [Overview and Shared Contracts](feature/offloader/README.md)
- [Model-Level Offload](feature/offloader/cpu_offload.md)
- [Layerwise Offload](feature/offloader/layerwise_offload.md)
- [Distributed Layerwise Offload](feature/offloader/distributed_layerwise_offload.md)
-
 - [Cache-DiT](feature/cache_dit.md)
 - [TeaCache](feature/teacache.md)
 - [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
+- [Distributed Layerwise Offload](feature/distributed_layerwise_offload.md)

 ## Infrastructure and Performance

@@ -48,10 +48,10 @@ defines the shared configuration and backend extension points.
 | Goal | User guide | Related design contract |
 | --- | --- | --- |
 | Compare acceleration methods and supported combinations | [Overview](../user_guide/diffusion_features.md), [Feature Compatibility](../user_guide/feature_compatibility.md) | [Diffusion acceleration designs](../design/index.md#diffusion-acceleration) |
-| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [CPU Offloading](../design/feature/offloader/README.md) |
+| Move weights between host and device memory | [CPU Offloading](../user_guide/diffusion/cpu_offload.md) | [Distributed Layerwise Offload](../design/feature/distributed_layerwise_offload.md) |
 | Reuse denoising computation | [Cache Acceleration](../user_guide/diffusion/cache_acceleration/cache_dit.md) | [Cache-DiT](../design/feature/cache_dit.md), [TeaCache](../design/feature/teacache.md) |
 | Distribute diffusion work across devices | [Parallelism](../user_guide/diffusion/parallelism/overview.md) | [Parallelism designs](../design/index.md#parallelism) |
-| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Attention Backend Selection](../design/feature/attention_backend_selection.md) |
+| Select dense, sparse, or quantized attention paths | [Attention Backends](../user_guide/diffusion/attention_backends.md) | [Skip-Softmax](../design/feature/skip_softmax.md) |
 | Compile repeated diffusion regions | [Regional Compilation](../user_guide/diffusion/regional_compilation.md) | User-facing optimization guide |
 | Add generated video frames | [Frame Interpolation](../user_guide/diffusion/frame_interpolation.md) | User-facing extension guide |
 | Reduce diffusion model startup time | [Startup and Loading](../user_guide/diffusion/startup_and_loading.md) | User-facing loading guide |
@@ -101,7 +101,7 @@ outputs = omni.generate(
 )

 # Access custom trajectory data
-output = outputs[0]
+output = outputs[0].request_output
 print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
 print(f"Trajectory latents shape: {output.latents.shape}")
 ```
@@ -47,7 +47,7 @@ if __name__ == "__main__":
    omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
    prompt = "a cup of coffee on the table"
    outputs = omni.generate(prompt)
-    images = outputs[0].images
+    images = outputs[0].request_output.images
    images[0].save("coffee.png")
 ```

@@ -75,7 +75,8 @@ if __name__ == "__main__":
    ]
    omni_outputs = omni.generate(prompts)
    for i_prompt, prompt_output in enumerate(omni_outputs):
-        this_images = prompt_output.images
+        this_request_output = prompt_output.request_output
+        this_images = this_request_output.images
        for i_image, image in enumerate(this_images):
            image.save(f"p{i_prompt}-img{i_image}.jpg")
            print("saved to", f"p{i_prompt}-img{i_image}.jpg")
@@ -43,12 +43,7 @@ th {
 | `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   | — |
 | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
 | `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
-| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
-| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
-| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
-| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
 | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
@@ -103,7 +98,6 @@ th {
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | | — |
 | `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
-| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | | — |
 |`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |

 The `Recipe` column links the corresponding deployment recipe. `Published` is
@@ -1,192 +1,494 @@
 # Diffusion Attention Backends

-Use this page to select and configure a diffusion attention backend. Backend-
-specific installation, tuning, and compatibility details live in separate
-guides so that the selection contract stays easy to scan.
+This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.

-Diffusion attention backend selection applies to DiT attention in image and
-video generation models. It does **not** change autoregressive LLM attention,
-which uses vLLM's attention selector.
+## Overview

-For the internal selector, registry, and platform contract, see
-[Attention Backend Selection](../../design/feature/attention_backend_selection.md).
+Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.

-## Choose a guide
+This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.

-| Need | Guide |
-| --- | --- |
-| Select a conservative or platform-native dense kernel | [Dense Backends](attention_backends/dense_backends.md) |
-| Use FlashInfer trtllm-gen FMHA, Skip-Softmax, or TRTLLM SAGE quantization | [TRTLLM Attention](attention_backends/trtllm.md) |
-| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
-| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
-| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
+The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.

-## Backend options
+## Backend Options
+
+| Value | Notes |
+|---|---|
+| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
+| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
+| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
+| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. |
+| `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. |
+| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. |

-| Value | Family | Primary use | Detail |
-| --- | --- | --- | --- |
-| `TORCH_SDPA` | Dense | Conservative reference; always available | [Dense Backends](attention_backends/dense_backends.md#torch_sdpa) |
-| `FLASH_ATTN` | Dense | FlashAttention 4/3/2 depending on the installed package and GPU | [Dense Backends](attention_backends/dense_backends.md#flash_attn) |
-| `CUDNN_ATTN` | Dense | Mask-heavy DiTs on Blackwell with cuDNN 9.5 or newer | [Dense Backends](attention_backends/dense_backends.md#cudnn_attn) |
-| `FLASHINFER_ATTN` | Dense or quantized | FlashInfer batch prefill; optional mixed Q/K and V dtypes | [Dense Backends](attention_backends/dense_backends.md#flashinfer_attn) |
-| `TRTLLM_ATTN` | Dense, sparse, or quantized | Datacenter Blackwell with `head_dim=128` and compatible packed paths | [TRTLLM Attention](attention_backends/trtllm.md) |
-| `SAGE_ATTN` | Quantized | SageAttention 2.2 INT8 attention | [SageAttention](attention_backends/sage.md#sage_attn) |
-| `SAGE_ATTN_3` | Quantized | SageAttention3 on Blackwell | [SageAttention](attention_backends/sage.md#sage_attn_3) |
-| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
-| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
-| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |

 ## Configuration

-Backend selection is resolved in this order:
+Diffusion attention backends can be configured three ways, in priority order:

-1. `--diffusion-attention-config` per-role configuration.
-2. `--diffusion-attention-backend` or `DIFFUSION_ATTENTION_BACKEND` as a
-   global default.
-3. The current platform's default.
+1. **`--diffusion-attention-config`** — structured per-role config (highest priority).
+2. **`--diffusion-attention-backend` / `DIFFUSION_ATTENTION_BACKEND` env var** — global shorthand that sets the default backend.
+3. **Platform default** — used when nothing is configured.

-`--diffusion-attention-backend` is shorthand for
-`--diffusion-attention-config.default.backend`. Do not pass it together with an
-explicit `default.backend` in the structured configuration.
+`--diffusion-attention-backend` is shorthand for `--diffusion-attention-config.default.backend`. It may be combined with `--diffusion-attention-config.per_role.*` overrides, but is mutually exclusive with `--diffusion-attention-config.default.backend`.

 ### Global default

-```bash
-vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
+Set the default backend for every diffusion attention layer:

-# Backwards-compatible environment variable
-export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
+```bash
+# CLI flag
+vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
+
+# Environment variable (also recognized for backwards compatibility)
+export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
 ```

 ### Per-role configuration

-Roles are declared by each diffusion model. Common categories are `self` and
-`cross`; a model may also use a more specific role such as
-`ltx2.audio_to_video`. Resolution order is:
+Roles are free-form strings declared by each diffusion model. The two common categories are `"self"` and `"cross"`; model-specific roles (e.g. `"ltx2.audio_to_video"`) may also be declared. A role string is matched in this order:

-1. Exact `per_role[role]` match.
-2. Category `per_role[role_category]` match.
-3. `default`.
-4. Platform default.
+1. Exact `per_role[role]` match
+2. `per_role[role_category]` fallback (e.g. `"ltx2.audio_to_video"` → `"cross"`)
+3. `default`
+4. Platform default
+
+Use vLLM-style dotted flags or one JSON blob:

 ```bash
 # Dotted flags
 vllm-omni serve <model> \
-  --diffusion-attention-config.default.backend FLASH_ATTN \
-  --diffusion-attention-config.per_role.cross.backend TORCH_SDPA
+    --diffusion-attention-config.default.backend FLASH_ATTN \
+    --diffusion-attention-config.per_role.cross.backend TORCH_SDPA

-# Equivalent JSON
+# JSON
 vllm-omni serve <model> \
-  --diffusion-attention-config \
-  '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
+    --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}'
 ```

-### Python API
+A backend that needs configuration exposes it as a typed field on the spec:
+`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and
+`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)):
+
+```bash
+--diffusion-attention-config.default.backend TRTLLM_ATTN \
+--diffusion-attention-config.default.skip_softmax.target_sparsity 0.5
+```
+
+### Programmatic API
+
+When constructing `OmniDiffusionConfig` directly:

 ```python
-from vllm_omni.diffusion.data import (
-    AttentionConfig,
-    AttentionSpec,
-    OmniDiffusionConfig,
-)
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig

 config = OmniDiffusionConfig(
    diffusion_attention_config=AttentionConfig(
        default=AttentionSpec(backend="FLASH_ATTN"),
-        per_role={"cross": AttentionSpec(backend="TORCH_SDPA")},
+        per_role={
+            "cross": AttentionSpec(backend="TORCH_SDPA"),
+        },
    ),
    ...,
 )
 ```

-Backend-specific typed blocks are documented with their consumers:
+A plain dict is also accepted and normalized to `AttentionConfig`.

- `quant`: [FlashInfer](attention_backends/dense_backends.md#flashinfer-quantized-attention)
-  and [TRTLLM SAGE](attention_backends/trtllm.md#sage-quantization).
- `skip_softmax`: [TRTLLM Skip-Softmax](attention_backends/trtllm.md#skip-softmax).
- `block_sparse`: [RainFusion](attention_backends/rainfusion.md#configuration).
+#### FlashInfer QK16/V8 quantized attention

-## Platform defaults
+`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8:
+
+```python
+config = OmniDiffusionConfig(
+    diffusion_attention_config=AttentionConfig(
+        default=AttentionSpec(
+            backend="FLASHINFER_ATTN",
+            quant=AttnQuantSpec(
+                dtype_qk="bfloat16",
+                dtype_vo="fp8_e4m3",
+            ),
+        ),
+    ),
+    ...,
+)
+```
+
+`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype.
+
+## Platform Defaults

 ### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)

-The CUDA auto-route preference is:
+Auto-route preference, in order:

-1. `TRTLLM_ATTN` on datacenter Blackwell (sm_100/sm_103) when FlashInfer is
-   available, `head_dim=128`, and the model declares a compatible packed or
-   mask-free path.
-2. `CUDNN_ATTN` when cuDNN 9.5 or newer is available.
-3. `FLASHINFER_ATTN` when FlashInfer is available but cuDNN is too old.
-4. `FLASH_ATTN` when a compatible package is installed.
-5. `TORCH_SDPA`.
+1. `TRTLLM_ATTN` — on **datacenter** Blackwell (sm_100 / sm_103) when `flashinfer` is installed, the model's `head_dim` is 128, **and the model declares a compatible packed/mask-free path**
+2. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
+3. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
+4. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
+5. `TORCH_SDPA` — last resort

-`TRTLLM_ATTN` is not auto-selected on workstation Blackwell
-(sm_120/sm_121), for other head dimensions, or for paths that require an
-unsupported mask.
+`TRTLLM_ATTN` is skipped on workstation Blackwell (sm_120 / sm_121) and for any `head_dim != 128`, so those GPUs keep the `CUDNN_ATTN` route described below.

-### Hopper, Ada, and Ampere
+`TRTLLM_ATTN` outranks `CUDNN_ATTN` on datacenter Blackwell for compatible packed/mask-free pipelines. Workstation Blackwell (sm_120 / sm_121) and pipelines that require attention masks retain their normal fallback.

-The CUDA auto-route uses `FLASH_ATTN` when available and otherwise falls back
-to `TORCH_SDPA`. `CUDNN_ATTN` and `FLASHINFER_ATTN` remain explicit options.
+The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.

-Other platforms validate an explicit backend and choose their own default
-through the platform implementation. Check the startup log to confirm the
-resolved backend.
+**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.

-## Choosing a backend manually
+### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)

-Override the platform default when you need:
+Auto-route preference:

- a correctness reference (`TORCH_SDPA`);
- a backend-specific workaround;
- training/rollout kernel alignment (Hub backends); or
- an explicitly validated sparse or quantized speedup.
+1. `FLASH_ATTN` — when `flash-attn` is installed
+2. `TORCH_SDPA` — fallback

-The startup log prints the resolved backend and whether it came from explicit
-configuration or platform defaulting. If no resolution message appears, check
-earlier logs for diffusion-stage initialization failures.
+`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.

-## Reference benchmark
+## TRTLLM_ATTN Backend and Skip-Softmax

-The following BF16 results were measured on an sm_120 RTX Pro 6000 Blackwell
-with the same prompt and seed across runs. Treat them as reference results, not
-portable guarantees.
+`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA. Dense it is a BF16 backend like the others; on top of
+that it can enable **Skip-Softmax**, a sparse-attention mode that trades a little fidelity for speed
+(algorithm: [design doc](../../design/feature/skip_softmax.md)).

-| Model | Shape | `TORCH_SDPA` | `CUDNN_ATTN` | `FLASHINFER_ATTN` |
-| --- | --- | ---: | ---: | ---: |
+Enable it through the typed `skip_softmax` block on the attention spec:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `target_sparsity` | finite, `[0, 1]` | Operating point on a calibrated curve. Needs a calibration for the model. |
+| `threshold` | finite, `≥ 0` | Direct threshold, no calibration needed. Mutually exclusive with `target_sparsity`. |
+| `disabled_until_timestep` | finite, `[0, 1]` | Holds early, high-noise steps dense; skip turns on once the normalized timestep `t ≤ D`. |
+
+```bash
+vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
+      "skip_softmax": {"target_sparsity": 0.65, "disabled_until_timestep": 0.86}}}'
+```
+
+Programmatically the same block is a typed `SkipSoftmaxSpec` (values validated at construction):
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="TRTLLM_ATTN",
+        skip_softmax=SkipSoftmaxSpec(target_sparsity=0.65, disabled_until_timestep=0.86),
+    ),
+)
+```
+
+**Start at `target_sparsity=0.65, disabled_until_timestep=0.86`** and raise `target_sparsity` only
+as far as your quality bar allows — the output diverges from dense as it climbs. The gain grows with
+sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `TRTLLM_ATTN` with no
+`skip_softmax` block.
+
+Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
+`TRTLLM_ATTN` raises rather than silently degrading.
+
+## TRTLLM_ATTN SAGE Quantization
+
+On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
+FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
+in low precision. Enable it through the typed `quant` block on the attention spec:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
+| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
+| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
+
+V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
+neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
+
+The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
+fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
+`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
+`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
+
+```bash
+vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
+      "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
+```
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="TRTLLM_ATTN",
+        quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
+    ),
+)
+```
+
+**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
+availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
+(B300); `int8` QK kernels are compiled for **SM100 only**.
+
+## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
+
+`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each
+128-token block of keys, ranks them per query block, and computes attention against only the
+top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are
+spatiotemporal neighbours rather than arbitrary rows of the packed sequence.
+
+Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the
+first-frame blocks are always kept dense, which is why the realized sparsity is lower than the
+nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an
+exempt layer, a sequence with no published video geometry, a video segment under 32 blocks —
+delegates to `FLASH_ATTN`, so a model can select this backend unconditionally.
+
+Enable it through the typed `block_sparse` block on the attention spec, shared by every
+block-sparse backend:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. |
+| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. |
+| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. |
+
+```bash
+vllm-omni serve MiniMaxAI/MiniMax-H3 \
+  --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN",
+      "block_sparse": {"sparsity": 0.8, "start_step": 0}}}'
+```
+
+Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction):
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="RAINFUSION_ATTN",
+        block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"),
+    ),
+)
+```
+
+**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the
+cheapest way to recover structure, because the early high-noise steps decide global layout while the
+later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for
+`skip_layers` last, once an A/B against dense at the same seed points at specific blocks.
+
+Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also
+incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank
+blocks — use Ulysses SP (`ring_degree=1`).
+
+### Which geometries run sparse
+
+MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when
+necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2
+does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the
+promoted rows remain visible to every query and key.
+
+For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row
+128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix.
+MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask
+and kernel tiling aligned, at the cost of retaining more blocks.
+
+8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8
+creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent
+grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release
+that implements protected video tails for rf_v2.
+
+## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
+
+Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
+
+| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
+|---|---|---|---|---|
 | HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
 | Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
 | Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
 | FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |

-Mask-heavy DiTs favored `CUDNN_ATTN`; lighter-mask or TP-saturated workloads
-were close enough that users should benchmark their exact model and shape.
+Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.

-## Compatibility anchors
+## Known Limitations

-The following headings preserve links to sections that moved into dedicated
-guides.
+### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile

-## TRTLLM_ATTN Backend and Skip-Softmax
+LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).

-See [TRTLLM Attention](attention_backends/trtllm.md#skip-softmax).
+**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:

-## TRTLLM_ATTN SAGE Quantization
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Lightricks/LTX-2 ...
+```

-See [TRTLLM Attention](attention_backends/trtllm.md#sage-quantization).
+### FlashAttention-4 on Blackwell

-## RAINFUSION_ATTN Backend and Block-Sparse Video Attention
+Install the optional CUDA 13 extra to use the CuTe-based FA4 path:

-See [RainFusion](attention_backends/rainfusion.md).
+```bash
+pip install 'vllm-omni[fa4]'
+```
+
+`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
+when it is unavailable. Version `4.0.0b18` is required; earlier beta
+wheels had known JIT failures on Blackwell.
+
+## Choosing a Backend Manually
+
+### When to override the default
+
+- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
+- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
+- **Workaround for known issues**: see Known Limitations above.
+
+### Verifying which backend is in use
+
+The startup log prints one of:
+
+```
+Using diffusion attention backend 'CUDNN_ATTN'           # explicit override
+Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
+Defaulting to diffusion attention backend SDPA           # nothing else available
+```
+
+If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.

 ## SageAttention Installation

-See [SageAttention](attention_backends/sage.md#installation).
+vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
+
+Build from source:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd SageAttention
+
+export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
+pip install . --no-build-isolation
+```
+
+Quick check:
+
+```bash
+python -c "import sageattention; print(sageattention.__file__)"
+```

 ## SageAttention3 Installation

-See [SageAttention](attention_backends/sage.md#sageattention3-installation).
+vLLM-Omni expects SageAttention3 to be installed into the same Python environment as vLLM-Omni.
+
+Build from source:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd SageAttention/sageattention3_blackwell
+python setup.py install
+```
+
+Quick check:
+
+```bash
+python -c "import sageattn3; print(sageattn3.__file__)"
+```
+
+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

-See [Hugging Face Hub Backends](attention_backends/huggingface_hub.md).
+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)
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
+    --output hv15.mp4
+```
+
+On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
+
+### Explicit backend selection
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Lightricks/LTX-2 \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
+    --output ltx2.mp4
+```
+
+### SageAttention (lossy)
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
+    --tensor-parallel-size 2 \
+    --output hv15_sage.mp4
+```
+
+Example: Wan2.2 TI2V 5B
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
+    --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 704 --width 1280 --num-frames 49 \
+    --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
+    --tensor-parallel-size 2 \
+    --output outputs/wan22_sage.mp4
+```
+
+### Enable SageAttention3
+
+Example:
+
+```bash
+DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_to_video/text_to_video.py \
+    --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
+    --prompt "A dog running across a field of golden wheat." \
+    --height 480 --width 832 --num-frames 33 \
+    --num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
+    --tensor-parallel-size 2 \
+    --output outputs/hv15_sage3.mp4
+```
@@ -1,110 +0,0 @@
-# Dense Attention Backends
-
-Use these backends when you want dense attention without backend-specific
-sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark
-the fastest compatible kernel for your model, shape, and hardware.
-
-For selection precedence and per-role configuration, see the
-[attention backend overview](../attention_backends.md).
-
-## `TORCH_SDPA`
-
-`TORCH_SDPA` calls PyTorch `scaled_dot_product_attention` and lets PyTorch's
-dispatcher choose the implementation. It is always available and is the most
-conservative reference when validating another backend.
-
-```bash
-vllm-omni serve <model> --diffusion-attention-backend TORCH_SDPA
-```
-
-## `FLASH_ATTN`
-
-`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it
-prefers the CuTe-based FlashAttention 4 path and falls back to
-FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
-Ampere it is the preferred automatic route when a compatible package is
-installed.
-
-```bash
-vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
-```
-
-### FlashAttention 4 on Blackwell
-
-Install the optional CUDA 13 extra:
-
-```bash
-pip install 'vllm-omni[fa4]'
-```
-
-Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on
-Blackwell. If the CuTe path is unavailable, the backend falls back to the
-compatible FlashAttention 3 or 2 path.
-
-## `CUDNN_ATTN`
-
-`CUDNN_ATTN` pins PyTorch SDPA to `CUDNN_ATTENTION`. It is particularly useful
-for mask-heavy DiTs and is automatically preferred on Blackwell when cuDNN
-9.5 or newer is available and the higher-priority TRTLLM route is not
-compatible.
-
-```bash
-vllm-omni serve <model> --diffusion-attention-backend CUDNN_ATTN
-```
-
-### LTX-2.0 limitation
-
-LTX-2 audio attention has a symbolic head dimension during `torch.compile`
-tracing. The cuDNN SDPA selector rejects that symbolic dimension and Dynamo
-aborts compilation. This is tracked in
-[issue #3121](https://github.com/vllm-project/vllm-omni/issues/3121).
-
-Use `FLASHINFER_ATTN` or `TORCH_SDPA` as a workaround:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
-  python examples/offline_inference/text_to_video/text_to_video.py \
-  --model Lightricks/LTX-2 ...
-```
-
-## `FLASHINFER_ATTN`
-
-`FLASHINFER_ATTN` uses FlashInfer's batch-prefill wrapper. It is an explicit
-option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
-is installed but cuDNN is too old for `CUDNN_ATTN`.
-
-```bash
-vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN
-```
-
-### FlashInfer quantized attention
-
-The backend accepts an `AttentionSpec.quant` block. For QK16/V8, keep Q and K
-in FP16 or BF16 and use FP8 E4M3 for V:
-
-```python
-from vllm_omni.diffusion.data import (
-    AttentionConfig,
-    AttentionSpec,
-    AttnQuantSpec,
-    OmniDiffusionConfig,
-)
-
-config = OmniDiffusionConfig(
-    diffusion_attention_config=AttentionConfig(
-        default=AttentionSpec(
-            backend="FLASHINFER_ATTN",
-            quant=AttnQuantSpec(
-                dtype_qk="bfloat16",
-                dtype_vo="fp8_e4m3",
-            ),
-        ),
-    ),
-    ...,
-)
-```
-
-`dtype_qk` controls Q and K; `dtype_vo` controls V. Mixed-dtype
-configurations require FlashInfer 0.6.16rc1 or newer. The shared quantization
-schema is also consumed by TRTLLM, but each backend validates its own allowed
-fields and values; see [TRTLLM SAGE quantization](trtllm.md#sage-quantization).
@@ -1,32 +0,0 @@
-# Hugging Face Hub Attention Backends
-
-Hub backends load the same published kernels commonly used with Hugging Face
-Diffusers. They are useful when training and serving must use matching kernel
-implementations to reduce numerical drift and sampling divergence.
-
-| Backend | Kernel | Platform |
-| --- | --- | --- |
-| `FLASH_ATTN_HUB` | `kernels-community/flash-attn2` | Compatible CUDA GPUs |
-| `FLASH_ATTN_3_HUB` | `kernels-community/flash-attn3` | Hopper sm_90 or newer |
-
-## Installation
-
-```bash
-pip install kernels==0.14.1
-```
-
-If `kernels` is unavailable, vLLM-Omni warns and falls back to the
-corresponding local FlashAttention implementation. On CUDA devices older than
-Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
-
-## Usage
-
-```bash
-export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB
-
-# Equivalent CLI selection
-vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB
-```
-
-For per-role kernel alignment, configure the Hub backend on only the relevant
-role; see the [attention backend overview](../attention_backends.md#per-role-configuration).
@@ -1,72 +0,0 @@
-# RainFusion Attention
-
-`RAINFUSION_ATTN` runs MindIE-SD RainFusion (`rf_v2`) block-sparse video
-attention on Ascend NPU. It pools 128-token key blocks, ranks them per query
-block, and attends to the highest-scoring blocks after arranging video tokens
-in `(t, h, w)` order.
-
-Only the video segment is sparse. Prefix rows such as text, visual conditions,
-and audio, plus first-frame blocks, remain dense. Unsupported calls—including
-warmup steps, skipped layers, missing video geometry, and video segments below
-32 blocks—delegate to `FLASH_ATTN`, so compatible models can select RainFusion
-globally.
-
-## Configuration
-
-| Key | Valid values | Meaning |
-| --- | --- | --- |
-| `sparsity` | finite, `[0, 1]` | Nominal dropped-key-block fraction; default `0.8`; `0` disables sparsity |
-| `start_step` | integer, `>= 0` | Number of early denoise steps kept dense |
-| `skip_layers` | selector such as `"0-3,38"` | DiT blocks kept dense |
-
-```bash
-vllm-omni serve MiniMaxAI/MiniMax-H3 \
-  --diffusion-attention-config '{"default":{"backend":"RAINFUSION_ATTN",\
-    "block_sparse":{"sparsity":0.8,"start_step":0}}}'
-```
-
-```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec
-
-config = AttentionConfig(
-    default=AttentionSpec(
-        backend="RAINFUSION_ATTN",
-        block_sparse=BlockSparseSpec(
-            sparsity=0.8,
-            start_step=0,
-            skip_layers="0-1",
-        ),
-    ),
-)
-```
-
-Tune `start_step`, then `sparsity`, then `skip_layers`. Increasing the dense
-early-step window is usually the cheapest way to recover global structure;
-use layer exclusions only after a same-seed dense comparison identifies
-sensitive blocks.
-
-## Requirements and compatibility
-
-RainFusion requires Ascend NPU and `mindiesd`. Selecting it on another
-platform raises. It is incompatible with ring sequence parallelism because
-the kernel needs the complete key sequence for ranking; use Ulysses sequence
-parallelism with `ring_degree=1`.
-
-## Geometry handling
-
-RainFusion handles arbitrary video grids by rearranging video spatially and,
-when necessary, promoting a real-video suffix to the always-kept prefix. It
-does not pad because `rf_v2` does not consume a padding attention mask.
-
-For MiniMax-H3 at 1344x768, grid `(62, 24, 42)` produces 62,496 video rows.
-The implementation promotes 2,976 rows and leaves 59,520 sparse rows
-(`465 x 128`) so the mask and kernel tiling remain aligned.
-
-Spatial grids aligned to 8x8 are still preferable. A latent height or width
-not divisible by 8 creates an always-kept suffix and reduces realized
-sparsity. At the image level, that means width and height divisible by 256.
-Protected video tails require a compatible MindIE-SD release.
-
-For common configuration and selector behavior, see the
-[attention backend overview](../attention_backends.md) and the
-[backend selection design](../../../design/feature/attention_backend_selection.md).
@@ -1,63 +0,0 @@
-# SageAttention
-
-SageAttention backends provide lossy low-precision attention for diffusion
-models. Validate output quality against `TORCH_SDPA` at the same seed before
-using either backend in production.
-
-## `SAGE_ATTN`
-
-`SAGE_ATTN` uses SageAttention 2.2 with INT8-quantized attention and FP16
-accumulation.
-
-### Installation
-
-Install SageAttention into the same environment as vLLM-Omni:
-
-```bash
-git clone https://github.com/thu-ml/SageAttention.git
-cd SageAttention
-export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
-pip install . --no-build-isolation
-```
-
-Verify the installation:
-
-```bash
-python -c "import sageattention; print(sageattention.__file__)"
-```
-
-Select it globally:
-
-```bash
-vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN
-```
-
-## `SAGE_ATTN_3`
-
-`SAGE_ATTN_3` uses the SageAttention3 Blackwell implementation.
-
-### SageAttention3 installation
-
-```bash
-git clone https://github.com/thu-ml/SageAttention.git
-cd SageAttention/sageattention3_blackwell
-python setup.py install
-```
-
-Verify the installation:
-
-```bash
-python -c "import sageattn3; print(sageattn3.__file__)"
-```
-
-```bash
-vllm-omni serve <model> --diffusion-attention-backend SAGE_ATTN_3
-```
-
-`SAGE_ATTN_3` requires CUDA, an importable `sageattn3`, and a Blackwell-class
-GPU. Its kernel assumes the query-head count equals the key/value-head count.
-GQA and MQA diffusion calls therefore fall back to PyTorch SDPA for
-correctness.
-
-For common configuration and platform routing, see the
-[attention backend overview](../attention_backends.md).
@@ -1,103 +0,0 @@
-# TRTLLM Attention
-
-`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
-BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
-
-For common selection and per-role configuration, see the
-[attention backend overview](../attention_backends.md).
-
-## Requirements
-
-Selecting `TRTLLM_ATTN` requires:
-
- datacenter Blackwell, sm_100 or sm_103;
- `head_dim=128`;
- FlashInfer; and
- a model-declared packed or mask-free compatible path.
-
-It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
-explicit selection raises instead of silently changing the backend.
-
-## Dense mode
-
-Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
-
-```bash
-vllm-omni serve <model> --diffusion-attention-backend TRTLLM_ATTN
-```
-
-## Skip-Softmax
-
-Skip-Softmax trades some fidelity for speed by skipping low-contribution
-attention work. See the [feature design](../../../design/feature/skip_softmax.md)
-for the algorithm and implementation contract.
-
-| Key | Valid values | Meaning |
-| --- | --- | --- |
-| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve |
-| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` |
-| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` |
-
-```bash
-vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
-    "skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}'
-```
-
-```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
-
-config = AttentionConfig(
-    default=AttentionSpec(
-        backend="TRTLLM_ATTN",
-        skip_softmax=SkipSoftmaxSpec(
-            target_sparsity=0.65,
-            disabled_until_timestep=0.86,
-        ),
-    ),
-)
-```
-
-Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase
-sparsity only after comparing quality against dense output at the same seed.
-The benefit grows with sequence length because only attention is accelerated.
-
-## SAGE quantization
-
-TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
-matrix multiplications use low precision.
-
-| Key | Valid values | Meaning |
-| --- | --- | --- |
-| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense |
-| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` |
-| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` |
-
-V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
-Only block sizes with compiled kernels are accepted.
-
-```bash
-vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
-    "quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
-```
-
-```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
-
-config = AttentionConfig(
-    default=AttentionSpec(
-        backend="TRTLLM_ATTN",
-        quant=AttnQuantSpec(
-            dtype_qk="fp8_e4m3",
-            q_block_size=1,
-            k_block_size=16,
-        ),
-    ),
-)
-```
-
-This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are
-available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only.
-The shared `AttnQuantSpec` is backend typed: values intended for
-`FLASHINFER_ATTN` are rejected here and vice versa.
@@ -1,75 +1,401 @@
 # CPU Offloading for Diffusion Models

-vLLM-Omni provides three CPU-offload strategies for diffusion models. Choose
-the coarsest strategy that meets your memory target: finer-grained streaming
-saves more device memory but adds more transfer and setup overhead.
+## Overview

-For the shared factory, discovery, and lifecycle contract, see the
-[CPU Offloading design](../../design/feature/offloader/README.md).
+vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:

-## Choose a strategy
+1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
+2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap.

-| Strategy | Device residency | Parallel scope | Guide |
-| --- | --- | --- | --- |
-| Model-level (sequential) | One pipeline component group at a time | Single device | [Model-Level Offloading](offloader/cpu_offload.md) |
-| Layerwise (blockwise) | One transformer block, with next-block prefetch | Single device | [Layerwise Offloading](offloader/layerwise_offload.md) |
-| Distributed layerwise | Fixed two-block device buffer; optional host-weight sharding and AllGather | Multiple GPU/NPU ranks | [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md) |
+Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority.

-All strategies use pinned host memory for faster transfers where applicable.
-Configuration priority is:
-
-1. Distributed layerwise offloading.
-2. Layerwise offloading.
-3. Model-level offloading.
-
-Treat the flags as mutually exclusive. If more than one is enabled, the
-higher-priority strategy is selected.
-
-## Quick selection
-
- Use [model-level offloading](offloader/cpu_offload.md) when swapping whole
-  encoders and DiTs is enough to fit and phase-boundary transfers are
-  acceptable.
- Use [layerwise offloading](offloader/layerwise_offload.md) for compute-heavy video
-  DiTs where block transfers can overlap computation.
- Use [distributed layerwise offloading](offloader/distributed_layerwise_offload.md)
-  when a multi-rank deployment also needs bounded device residency and,
-  optionally, sharded host weights.
-
-## Supported models
-
-| Architecture | Example models | DiT class | Model-level | Layerwise | Distributed layerwise | Layerwise block attributes |
-| --- | --- | --- | :---: | :---: | :---: | --- |
-| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | yes | yes | — | `transformer_blocks`, `single_transformer_blocks` |
-| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | — | yes | — | `transformer_blocks`, `single_transformer_blocks` |
-| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | — | yes | — | `layers` |
-| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | — | yes | — | `transformer` |
-| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | yes | yes | — | `transformer_blocks` |
-| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | yes | yes | — | `down_blocks`, `up_blocks` |
-| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | — | yes | — | `transformer_blocks` |
-| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
-| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | yes | yes | — | `blocks` |
-| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` | yes | yes | — | `layers` |
-| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | — | yes | — | `layers`, customized modules |
-| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | yes | yes | yes | `layers`, `gen_layers` |
-
-Model-level support requires discoverable DiT and encoder components.
-Layerwise support requires transformer block topology. Distributed support
-reuses that topology but still requires validation for the model, checkpoint,
-and parallel configuration.
-
-## Compatibility anchors
-
-These headings preserve links to sections that moved into dedicated guides.

 ## Model-level (Sequential) Offloading

-See [Model-Level Offloading](offloader/cpu_offload.md).
+### How It Works
+
+Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
+
+- **When encoders run**: DiT transformer is offloaded to CPU
+- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU.
+- **VAE**: Stays resident on GPU
+
+Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
+
+### Usage
+
+**Python API:**
+```python
+from vllm_omni import Omni
+
+m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
+```
+
+**CLI:**
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload
+```
+
+### To Support a Model
+
+Implement the `SupportsComponentDiscovery` protocol to declare which
+submodules serve as pipeline components (used by offloading, HSDP
+sharding, and other framework features):
+
+```python
+from typing import ClassVar
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+
+class MyPipeline(nn.Module, SupportsComponentDiscovery):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []  # optional
+
+    def __init__(self):
+        super().__init__()
+        self.transformer = ...     # DiT — stays on GPU during denoising
+        self.text_encoder = ...    # Encoder — offloaded to CPU during denoising
+        self.vision_model = ...    # Encoder — offloaded to CPU during denoising
+        self.vae = ...             # VAE — always on GPU
+```
+
+- `_dit_modules`: attribute names of denoising submodules (kept on GPU
+  during the diffusion loop).
+- `_encoder_modules`: attribute names of encoder/vision submodules
+  (offloaded to CPU during the diffusion loop).
+- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
+  part of the mutual exclusion hooks).
+- `_resident_modules`: attribute names of small submodules that must
+  stay on GPU during layerwise offloading (e.g. embedders, connectors).
+  Optional — defaults to `[]`.
+
+All attribute names support dotted paths for nested submodules
+(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
+
+Both DiT and encoder lists are needed because the offload hooks use
+mutual exclusion: when one group runs, the other moves to CPU.
+
+### Limitations
+- Cold start latency increases
+- Adds overhead from CPU-GPU transfers between encoder and denoising phases
+- 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

-See [Layerwise Offloading](offloader/layerwise_offload.md).
+### How It Works
+
+Layerwise offloading keeps only one transformer block on GPU at a time.
+
+As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.
+
+**Execution Flow:**
+
+| Block | Pre-forward Hook | Forward | Post-forward Hook |
+|-------|------------------|---------|-------------------|
+| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 |
+| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 |
+| ... | ... | ... | ... |
+| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) |
+
+Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward.
+
+Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.
+
+### Usage
+
+**Python API:**
+```python
+from vllm_omni import Omni
+
+# Text-to-video
+m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)
+
+# Or image-to-video
+m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)
+```
+
+**CLI:**
+```bash
+# Text-to-video
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload
+
+# Or image-to-video
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload
+```
+
+### To Support a Model
+
+Models must define the blocks attribute name for layerwise offloading:
+
+```python
+class WanTransformer3DModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute names containing transformer blocks
+
+    def __init__(self):
+        self.blocks = nn.ModuleList([...])  # Transformer blocks
+```
+
+For models with multiple block types:
+
+```python
+class Flux2Transformer2DModel(nn.Module):
+    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
+```
+
+### Limitations
+- 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
+

 ## Distributed Layerwise Offloading

-See [Distributed Layerwise Offloading](offloader/distributed_layerwise_offload.md).
+### How It Works
+
+Distributed layerwise offloading extends single-GPU layerwise offloading to
+multi-device deployments. In the default mode, each rank stores only
+**1/dp_size** of the model weights in host memory; full layer weights are
+reconstructed at runtime via **AllGather** on a dedicated communication stream,
+overlapped with computation via a fixed double-buffer scheme. The
+`--dlo-no-use-allgather` mode instead streams standard-loader rank-local
+weights independently and does not shard weights across ranks. For explicitly
+mmap-compatible TP1 models, the same mode can retain node-shared file-backed
+DiT weights and use bounded per-worker host staging instead of a private full
+DiT copy.
+
+For the implementation design and compatibility matrix, see
+[Distributed Layerwise Offload](../../design/feature/distributed_layerwise_offload.md).
+
+**Key features:**
+- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
+  reconstructed per-layer via `all_gather_into_tensor`
+- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
+  time, regardless of model size
+- **DP multi-concurrency**: N concurrent requests processed in parallel
+  (AllGather only gathers weight shards, which are request-independent)
+- **mmap weight storage**: compatible checkpoints use views into the shared OS
+  page cache. AllGather copies persistent rank shards; rank-local mode keeps
+  the views and packs complete blocks through two bounded host staging slots
+- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
+  (CANN/HCCL) via vLLM-Omni's platform abstraction
+
+**Execution Flow (per device):**
+
+```
+Compute Stream:  [Layer N]          [Layer N+1]          [Layer N+2]
+H2D Stream:      [H2D Shard N+1]   [H2D Shard N+2]
+AllGather:       [AG N+1]          [AG N+2]
+Slot Usage:      Slot0: Layer N    Slot1: Layer N+1
+```
+
+### Usage
+
+**CLI:**
+```bash
+# 4× GPU/NPU with AllGather (recommended)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4
+
+# Without DLO AllGather (compatible TP1 models use node-shared mmap storage)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4 \
+  --dlo-no-use-allgather
+
+# With SP instead of DP (long sequences)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --usp 4
+```
+
+**Python API:**
+```python
+from vllm_omni import Omni
+
+m = Omni(
+    model="/path/to/model",
+    enable_distributed_layerwise_offload=True,
+    dlo_use_allgather=True,  # default
+)
+```
+
+### CLI Flags
+
+| Flag | Description | Default |
+|------|-------------|---------|
+| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
+| `--data-parallel-size N` | Number of DP ranks; DLO weight-sharding group when AllGather is enabled | `1` |
+| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
+| `--dlo-no-use-allgather` | Use standard-loader rank-local weights; no DLO sharding or AllGather | `false` |
+
+### How model weights are loaded (mmap path)
+
+When DLO and an mmap-compatible model are active, the offloader:
+
+1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
+   from the normally-created transformer
+2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
+   initialization weights
+3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
+   which return views into the OS page cache without eagerly copying the full
+   tensor into anonymous process memory
+4. **Applies declared per-parameter layout adapters while packing**, when the
+   checkpoint and TP1 runtime layouts differ (for example MiniMax-H3 grouped
+   QKV)
+5. **Restores non-persistent buffers** from saved copies
+
+With AllGather, each rank copies only its persistent shard and then releases
+the source mappings. Without AllGather, mappings remain open and every worker
+owns two pinned staging slots sized to the largest streamed block. The latter
+shares file-backed pages, not process-local Python objects, and does not alter
+request routing or scheduling. Packing into the host slot is synchronous; the
+pinned H2D copy can overlap with computation.
+
+### OffloadPlan (declarative topology metadata)
+
+Models can optionally declare an `OffloadPlan` class variable to provide
+topology metadata (block attribute names, submodules to offload) without
+any offload-specific logic:
+
+```python
+from vllm_omni.diffusion.offloader import OffloadPlan
+
+class MyPipeline(nn.Module):
+    _dit_modules = ["transformer"]
+    _offload_plan = OffloadPlan(
+        block_attrs={"transformer": ("blocks",)},
+        offload_submodules={"context_encoder": "layers"},
+    )
+```
+
+When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
+and heuristic attribute search.  This is backward-compatible — existing models
+work without any changes.
+
+### DP Multi-concurrency
+
+When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
+up to `dp_size` requests per denoise step.  Each DP rank processes a different
+request while AllGather synchronizes only weight shards (request-independent).
+
+**Requirements for concurrent requests:**
+- `num_inference_steps` must be specified explicitly (None is not allowed)
+- All concurrent requests must have the same `num_inference_steps` value
+- These constraints exist because AllGather is a collective that requires
+  every rank to participate at each denoise step
+
+### Limitations
+- Online quantization (FP8) is incompatible with mmap loading — falls back to
+  regular loading only when the model is outside the DLO+AllGather mmap path;
+  online quantization with DLO+AllGather is rejected. Use
+  `--dlo-no-use-allgather` or disable online quantization.
+- Tensor Parallel > 1 is rejected in the DLO+AllGather mmap path because TP-aware
+  weight-loader callbacks are bypassed. The no-AllGather path retains the
+  standard TP loader instead of rank-local mmap and is experimental.
+- HSDP + AllGather is rejected because it would double-shard parameters. HSDP
+  with no-AllGather is accepted at configuration level, but has limited
+  end-to-end validation.
+- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
+
+**Module Discovery**
+
+The offloader discovers pipeline components in two ways:
+
+1. **Protocol-based** (preferred): If the pipeline implements
+    `SupportsComponentDiscovery`, its `_dit_modules`, `_encoder_modules`,
+    `_vae_modules`, and `_resident_modules` class variables are used
+    directly.  All attribute names support dotted paths (e.g.
+    `"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
+
+2. **Fallback attribute scan**: Otherwise, the offloader scans for
+    well-known attribute names:
+    - **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
+    - **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
+    - **VAE**: `vae`, `audio_vae`
+
+**Hook System**
+
+Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.
+
+**Backend Architecture**
+
+```
+OffloadBackend (base class)
+├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
+│                              (delegates to a pipeline's enable_omni_model_cpu_offload
+│                               for split models like Cosmos3)
+├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
+│                          (single-GPU, full weights on host)
+└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
+                                         (multi-GPU, sharded AllGather or rank-local mode)
+```
+
+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
+
+| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
+|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
+| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
+| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
+| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
+| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
+| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
+| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
+| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
+| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
+| MiniMaxH3Pipeline | `MiniMaxAI/MiniMax-H3` (`FL2VA`) | `MiniMaxH3DiTModel` | - | ✓ | ✓ | `"blocks"` |
+| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
+| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
+| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
+
+**Notes:**
+- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
+- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
+- Distributed Layerwise Offloading reuses the layerwise block topology, but the
+  model and parallelism combination still needs validation. AllGather mode
+  additionally requires a compatible checkpoint/mmap layout. For explicitly
+  compatible TP1 models, no-AllGather uses node-shared mmap sources; other
+  layouts retain the standard loader's rank-local tensors. See the
+  [DLO design](../../design/feature/distributed_layerwise_offload.md) and
+  [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
+  for usage examples.
@@ -1,81 +0,0 @@
-# Model-Level Offloading
-
-Model-level, or sequential, offloading keeps only the pipeline component group
-currently executing on the accelerator. It is the simplest offload strategy
-and is selected with `--enable-cpu-offload`.
-
-## How it works
-
-Pre-forward hooks enforce mutual exclusion between DiT and encoder modules:
-
- before an encoder runs, the DiT moves to CPU;
- before a DiT runs, encoders and other DiTs move to CPU; and
- VAE modules remain on the accelerator.
-
-Pinned host memory reduces transfer overhead. Transfers occur at phase
-boundaries, so cold-start and encoder-to-denoiser transitions become slower.
-
-## Usage
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    enable_cpu_offload=True,
-)
-```
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --omni --enable-cpu-offload
-```
-
-## Model integration
-
-Pipelines should implement `SupportsComponentDiscovery`:
-
-```python
-from typing import ClassVar
-
-from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
-
-
-class MyPipeline(nn.Module, SupportsComponentDiscovery):
-    _dit_modules: ClassVar[list[str]] = ["transformer"]
-    _encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
-    _vae_modules: ClassVar[list[str]] = ["vae"]
-    _resident_modules: ClassVar[list[str]] = []
-```
-
-All entries may be dotted paths. DiT and encoder lists are both required for
-mutual exclusion. VAE modules are pinned but not swapped; resident modules are
-small modules that must stay on the accelerator for layerwise paths.
-
-## Split-model components
-
-Some models, such as Cosmos3, split one transformer into mutually exclusive
-components that run in different phases. The pipeline exposes
-`enable_omni_model_cpu_offload`, and the backend delegates to the model-local
-contexts:
-
-```python
-class Cosmos3VFMTransformer(nn.Module):
-    def forward(self, ...):
-        with self._offload_context("reasoner"):
-            ...
-        with self._offload_context("generator"):
-            ...
-```
-
-This preserves the same invariant—exactly one component is device resident—
-while reusing sequential `.to()` movers.
-
-## Limitations
-
- Single device only.
- Higher cold-start latency.
- Transfers between encoder and denoising phases add latency.
-
-See the [model-level design](../../../design/feature/offloader/cpu_offload.md)
-for lifecycle and extension invariants.
@@ -1,153 +0,0 @@
-# Distributed Layerwise Offloading
-
-Distributed layerwise offloading (DLO) extends block streaming to multi-device
-deployments. With AllGather enabled, each rank stores roughly `1 / dp_size` of
-the host weights and reconstructs each layer at runtime. Without AllGather,
-each rank streams a complete block independently. Compatible TP1 deployments
-can share checkpoint-backed host pages among processes on the same node;
-otherwise DLO streams the ordinary loader's rank-local tensors.
-
-See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md)
-for the implementation contract and compatibility matrix.
-
-## Execution model
-
-DLO overlaps three operations with a fixed two-block device buffer:
-
-```text
-Compute stream:  [Layer N]          [Layer N+1]        [Layer N+2]
-H2D stream:      [H2D shard N+1]    [H2D shard N+2]
-AllGather:       [AG N+1]           [AG N+2]
-Slots:           slot 0: Layer N    slot 1: Layer N+1
-```
-
-AllGather communicates only request-independent weight shards, so data-
-parallel ranks may process different requests concurrently.
-
-## Usage
-
-```bash
-# Four ranks with sharded host weights and AllGather
-vllm serve /path/to/model --omni \
-  --enable-distributed-layerwise-offload \
-  --data-parallel-size 4
-
-# Standard-loader rank-local weights, without DLO AllGather
-vllm serve /path/to/model --omni \
-  --enable-distributed-layerwise-offload \
-  --data-parallel-size 4 \
-  --dlo-no-use-allgather
-
-# Sequence parallel deployment
-vllm serve /path/to/model --omni \
-  --enable-distributed-layerwise-offload \
-  --usp 4
-```
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="/path/to/model",
-    enable_distributed_layerwise_offload=True,
-    dlo_use_allgather=True,
-)
-```
-
-## Flags
-
-| Flag | Meaning | Default |
-| --- | --- | --- |
-| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
-| `--data-parallel-size N` | DP ranks and AllGather weight-sharding group | `1` |
-| `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
-| `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
-| `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
-
-## Host-weight loading
-
-The diffusion loader chooses host storage before DLO is enabled. It first
-attempts to build a complete, validated direct-checkpoint mmap plan. If names,
-coverage, shape, dtype, topology, or loader-callback compatibility cannot be
-proven, it runs the ordinary model loader instead. DLO consumes that result and
-does not make a second checkpoint-compatibility decision.
-
-The shared-mmap optimization in this phase is supported only with TP1. TP
-greater than one falls back before model mutation to ordinary TP-aware loading.
-DLO may still consume those TP-local tensors, but this is a compatibility path:
-it does not share checkpoint-backed runtime weights across DP replicas and
-provides no shared-mmap host-memory guarantee.
-
-The mmap plan skips only dedicated DiT weight sources. Other component sources,
-such as a text encoder loaded through the shared diffusion loader, continue to
-use their ordinary component loader. A checkpoint source that mixes DiT and
-non-DiT weights falls back completely rather than leaving an unplanned
-component uninitialized.
-
-With direct checkpoint mmap, the loader:
-
-1. saves non-persistent buffers such as RoPE frequencies;
-2. moves the normally created transformer to the meta device;
-3. loads checkpoint tensors as mmap views backed by the shared OS page cache;
-4. applies any loader-owned bounded layout adapters while packing blocks;
-5. restores saved non-persistent buffers; and
-6. preserves `post_load_weights()` and `validate_loaded_weights()` lifecycle
-   hooks.
-
-For AllGather with a group larger than one, each process copies only its
-persistent shard and then releases the source mapping. For no-AllGather, each
-process keeps the mapping open and packs complete blocks through two bounded
-pinned staging slots. Processes mapping the same files on one node share the
-immutable pages; no-AllGather still performs a complete-block H2D copy in each
-process.
-
-When the effective DLO group size is one, `dlo_use_allgather=True` does not
-perform a collective and uses the same rank-local transfer behavior.
-
-## Declarative topology
-
-Models may declare an `OffloadPlan` instead of embedding offload logic:
-
-```python
-from vllm_omni.diffusion.offloader import OffloadPlan
-
-
-class MyPipeline(nn.Module):
-    _dit_modules = ["transformer"]
-    _offload_plan = OffloadPlan(
-        block_attrs={"transformer": ("blocks",)},
-        offload_submodules={"context_encoder": "layers"},
-    )
-```
-
-When no plan exists, discovery falls back to
-`_layerwise_offload_blocks_attrs` and then heuristic attribute lookup.
-
-## Data-parallel concurrency
-
-With `data_parallel_size > 1` and AllGather enabled, the scheduler can process
-up to `dp_size` requests per denoising step. Every concurrent request must set
-the same explicit `num_inference_steps`; `None` is rejected because every rank
-must enter each collective.
-
-## Limitations
-
- Direct checkpoint mmap currently requires TP1. TP greater than one is
-  outside the Phase A shared-mmap support scope and falls back before model
-  mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
-  but it provides no shared-mmap host-memory benefit or guarantee.
- HSDP plus AllGather is rejected to avoid double sharding. HSDP without
-  AllGather has limited end-to-end validation.
- Online quantization uses the ordinary loader with no-AllGather. It remains
-  incompatible with DLO AllGather.
- Resident leading layers require `--dlo-no-use-allgather` and a model
-  `OffloadPlan` that declares eligible `resident_dit_paths`.
- DP concurrency requires an explicit, identical inference-step count.
-
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
-cache is a follow-up design in
-[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
-the direct-checkpoint path.
-
-See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
-for an end-to-end example.
@@ -1,65 +0,0 @@
-# Layerwise Offloading
-
-Layerwise, or blockwise, offloading keeps one transformer block on the
-accelerator and prefetches the next block while the current block computes.
-It is best suited to compute-heavy video DiTs whose block execution can hide
-host-to-device transfers.
-
-## Execution flow
-
-Each block has a pre-forward and post-forward hook. Parameters are consolidated
-in pinned host tensors and rematerialized for execution on a dedicated copy
-stream.
-
-| Block | Pre-forward hook | Forward | Post-forward hook |
-| --- | --- | --- | --- |
-| block 0 | Prefetch block 1 | Compute block 0 | Free block 0 |
-| block 1 | Prefetch block 2 | Compute block 1 | Free block 1 |
-| ... | ... | ... | ... |
-| last block | Prefetch block 0 | Compute last block | Free last block |
-
-Encoders, VAE modules, and non-block DiT modules such as embeddings and norms
-remain device resident.
-
-## Usage
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-    enable_layerwise_offload=True,
-)
-```
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --omni --enable-layerwise-offload
-```
-
-## Model integration
-
-Transformer classes declare containers of executable blocks:
-
-```python
-class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]
-
-
-class Flux2Transformer2DModel(nn.Module):
-    _layerwise_offload_blocks_attrs = [
-        "transformer_blocks",
-        "single_transformer_blocks",
-    ]
-```
-
-See the [layerwise design](../../../design/feature/offloader/layerwise_offload.md) for
-the discovery and hook invariants. `OffloadPlan` is a separate declarative
-topology path for distributed layerwise offload.
-
-## Limitations
-
- Single device only.
- Setup consolidates and pins block parameters, increasing cold-start time.
- Performance depends on block compute time and host-to-device bandwidth;
-  lightweight blocks may not hide transfers.
@@ -64,41 +64,6 @@ omni = Omni(
 )
 ```

-### Combined with Online FP8 Quantization
-
-HSDP can be combined with online FP8 quantization for diffusion models that
-support both features. HSDP shards the quantized weights across GPUs to reduce
-the per-GPU weight memory footprint.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-omni = Omni(
-    model="<your-model>",
-    quantization="fp8",
-    parallel_config=DiffusionParallelConfig(
-        use_hsdp=True,
-        hsdp_shard_size=2,
-    ),
-)
-```
-
-For online serving:
-
-```bash
-vllm serve <your-model> \
-    --omni \
-    --quantization fp8 \
-    --use-hsdp \
-    --hsdp-shard-size 2
-```
-
-The model must support both HSDP and online FP8 quantization. HSDP remains
-incompatible with Tensor Parallelism. See the
-[FP8 Quantization Guide](../../quantization/fp8.md) for supported models and
-hardware.
-
 ---

 ## Example Script
@@ -182,4 +147,3 @@ For detailed instructions on adding HSDP support to new models, see the [HSDP Co
 1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models
 2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism
 3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled
-4. ✅ **Combine with online FP8** - Use `quantization="fp8"` or `--quantization fp8` to reduce the sharded weight memory footprint
@@ -106,10 +106,7 @@ The following tables show which models support each feature:

 > Notes:

-> 1. CPU Offload has three strategies: model-level, layerwise, and distributed
->    layerwise. The tables below show **layerwise support** only. Split models
->    like Cosmos3 swap their reasoner/generator components for model-level
->    offload; see the [CPU Offload Guide](diffusion/cpu_offload.md).
+> 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.md).
 > 2. The **💾Quantization** column is collapsed for readability. See [Quantization](quantization/overview.md) for per-method and per-model support details.

 ### ImageGen
@@ -161,7 +158,6 @@ The following tables show which models support each feature:
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
-| **LTX-2.5**                  |     ❌     | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) |         ❓         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |    ❓ (FP8)     |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -175,7 +171,7 @@ The following tables show which models support each feature:
 **Frame Interpolation Support**

 - **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
+- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni

 ### AudioGen

@@ -218,7 +214,7 @@ The following tables show which models support each feature:
    3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
    4. The CPU Offloading (Layerwise) row describes local layerwise offload.
       Multi-device Distributed Layerwise Offload has a separate topology and
-       compatibility matrix in the [Distributed Layerwise Offloading guide](diffusion/offloader/distributed_layerwise_offload.md).
+       compatibility matrix in the [CPU Offloading guide](diffusion/cpu_offload.md#distributed-layerwise-offloading).
    5. The compatibility matrix uses FP8 as the representative quantization method.
    6. Step Execution is not compatible with any diffusion cache backend. LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).

@@ -43,7 +43,7 @@ if __name__ == "__main__":
            ),
        ],
    )
-    outputs[0].images[0].save("output.png")
+    outputs[0].request_output.images[0].save("output.png")
 ```

 ## Image-to-Image (Image Editing)
@@ -72,7 +72,7 @@ if __name__ == "__main__":
            ),
        ],
    )
-    outputs[0].images[0].save("output.png")
+    outputs[0].request_output.images[0].save("output.png")
 ```

 ## Generation Parameters
@@ -55,7 +55,7 @@ if __name__ == "__main__":
    omni = Omni(model="Qwen/Qwen-Image")
    prompt = "a cup of coffee on the table"
    outputs = omni.generate(prompt)
-    images = outputs[0].images
+    images = outputs[0].request_output.images
    images[0].save("coffee.png")
 ```

@@ -179,7 +179,7 @@ if __name__ == "__main__":
    ]
    outputs = omni.generate(prompts)
    for i, output in enumerate(outputs):
-        output.images[0].save(f"{i}.jpg")
+        output.request_output.images[0].save(f"{i}.jpg")
 ```

 !!! info
@@ -207,7 +207,7 @@ if __name__ == "__main__":
        }
    ])
    for i, output in enumerate(outputs):
-        output.images[0].save(f"{i}.jpg")
+        output.request_output.images[0].save(f"{i}.jpg")
 ```

 You can also pass a negative prompt via the CLI argument `--negative-prompt`:
@@ -165,7 +165,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \

 - **Tensor Parallelism — Text Encoder Not Sharded**: TP currently only shards the DiT blocks. Each TP rank retains a **full copy of the text encoder weights**, leading to significant GPU memory overhead proportional to TP degree. Tracked in [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).

- **CPU Offloading — Strategies Are Mutually Exclusive**: Distributed layerwise, layerwise, and model-level offload should not be enabled together. Selection priority is distributed layerwise, then layerwise, then model-level.
+- **CPU Offloading — Two Modes Are Mutually Exclusive**: Model-level offload (`enable_cpu_offload`) and layerwise offload (`enable_layerwise_offload`) cannot be used simultaneously. If both are set, layerwise takes priority and model-level is silently ignored.

 - **CPU Offloading — VAE stays on GPU**: Both offloading strategies keep the VAE on GPU at all times. For high-resolution generation, VAE decode can still cause OOM. Mitigate by combining with `vae_use_tiling=True` or VAE Patch Parallelism.

@@ -95,7 +95,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
 | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
 | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
 | Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
-| MiniMax-H3 | `MiniMaxAI/MiniMax-H3` (`FL2VA` / `Ref2VA`) | Yes | Not validated | `quantization="fp8"` quantizes eligible DiT and text-encoder linears; mixed-precision input/output heads stay FP32 | None | ✅︎ |
+| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |

 ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)

@@ -138,39 +138,6 @@ outputs = omni.generate(
 )
 ```

-### Global and per-component scope
-
-A plain, global quantization configuration is passed unchanged to every
-quantization-aware component constructed by a pipeline. This includes an
-eligible encoder when that encoder is implemented with vLLM quantizable
-layers; it does not rewrite arbitrary `torch.nn` modules. A structured
-component map is the only way to narrow that scope. Pipeline integrations that
-do not yet expose an encoder through the quantization factory remain DiT-only.
-
-For a pipeline that exposes both a transformer and a quantization-aware text
-encoder, the scope is:
-
-| Configuration | Transformer | Text encoder | Components without supported quantizable layers |
-|---------------|-------------|--------------|-------------------------------------------------|
-| `quantization="fp8"` | FP8 | FP8 | checkpoint precision |
-| `{"transformer": {"method": "fp8"}}` | FP8 | checkpoint precision | checkpoint precision |
-| `{"text_encoder": {"method": "fp8"}}` | checkpoint precision | FP8 | checkpoint precision |
-
-Use `quantization_config` for component-selective Python configuration. For
-example, quantize only the text decoder:
-
-```python
-omni = Omni(
-    model="<your-model>",
-    quantization_config={"text_encoder": {"method": "fp8"}},
-)
-```
-
-Component keys are runtime prefixes exposed by the pipeline integration; common
-keys include `transformer` and `text_encoder`. Entries may be combined, and
-`ignored_layers` can keep named eligible layers in checkpoint precision. Check
-the model recipe for supported components and their runtime prefixes.
-
 CLI:

 ```bash
@@ -96,7 +96,7 @@ if __name__ == "__main__":
        ),
    )
    from diffusers.utils import export_to_video
-    frames = outputs[0].images
+    frames = outputs[0].request_output.images
    export_to_video(frames, "quick_test_i2v.mp4", fps=16)
 ```

@@ -68,7 +68,6 @@ from vllm_omni.model_extras import (
 from vllm_omni.model_extras import (
    get_extra_body_params,
    get_model_class_name,
-    should_preserve_reference_image_size,
 )
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
@@ -455,13 +454,7 @@ def main():

    media_inputs: dict[str, Any] = {}
    if image is not None:
-        preserve_image_size = should_preserve_reference_image_size(
-            model_class_name,
-            model=args.model,
-        )
-        media_inputs["image"] = (
-            image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS)
-        )
+        media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
    if last_image is not None:
        media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
    if mask_image is not None:
@@ -335,7 +335,7 @@ def main(args):
    print(f"use_tts (tts_bos): {use_tts}")

    for stage_outputs in omni_generator:
-        output = stage_outputs
+        output = stage_outputs.request_output
        if stage_outputs.final_output_type == "text":
            request_id = output.request_id
            text_output = output.outputs[0].text
@@ -1,32 +1,28 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Validate MiniMax-H3 request-mode and DLO lifecycle behavior.
-
-The DLO modes accept any positive data-parallel size.  With AllGather enabled,
-the script submits one complete DP wave and verifies recovery after an invalid
-wave.  It also checks that every mode shuts down all worker processes cleanly.
+"""Reproduce MiniMax-H3 request-mode and DLO/DP2 lifecycle behavior.

 Examples:

    VLLM_WORKER_MULTIPROC_METHOD=spawn \
    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
    VLLM_OMNI_DLO_DP_WAVE_TIMEOUT=600 \
-    CUDA_VISIBLE_DEVICES=0,1,2,3 \
-    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
-        --model /path/to/MiniMax-H3/FL2VA --mode dlo --dp-size 4
+    CUDA_VISIBLE_DEVICES=0,1 \
+    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2

    # Rank-local mmap storage, with no DLO collective or DP wave scheduling.
    VLLM_WORKER_MULTIPROC_METHOD=spawn \
    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
    CUDA_VISIBLE_DEVICES=0,1 \
-    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
-        --model /path/to/MiniMax-H3/FL2VA --mode dlo-no-allgather --dp-size 2
+    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode dlo-dp2-no-allgather

    VLLM_WORKER_MULTIPROC_METHOD=spawn \
    VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
    CUDA_VISIBLE_DEVICES=0,1 \
-    python examples/offline_inference/minimax_h3/dlo_lifecycle.py \
-        --model /path/to/MiniMax-H3/FL2VA --mode request --tp-size 2
+    python examples/offline_inference/minimax_h3/dlo_dp2_lifecycle.py \
+        --model /path/to/MiniMax-H3/FL2VA --mode request
 """

 from __future__ import annotations
@@ -51,33 +47,14 @@ DEFAULT_PROMPTS = (
 )


-def positive_int(value: str) -> int:
-    parsed = int(value)
-    if parsed < 1:
-        raise argparse.ArgumentTypeError("must be a positive integer")
-    return parsed
-
-
 def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, help="Path to MiniMax-H3/FL2VA")
    parser.add_argument(
        "--mode",
-        choices=("dlo", "dlo-no-allgather", "request"),
+        choices=("dlo-dp2", "dlo-dp2-no-allgather", "request"),
        required=True,
-        help="DLO (with or without AllGather) or ordinary non-DLO request mode",
-    )
-    parser.add_argument(
-        "--dp-size",
-        type=positive_int,
-        default=2,
-        help="DLO data-parallel size (for example 1, 2, 4, or 8; default: 2)",
-    )
-    parser.add_argument(
-        "--tp-size",
-        type=positive_int,
-        default=2,
-        help="Tensor-parallel size used only by request mode (default: 2)",
+        help="DLO DP2 (with or without AllGather) or ordinary non-DLO TP2 request mode",
    )
    parser.add_argument("--steps", type=int, default=2)
    parser.add_argument("--duration", type=float, default=5.0)
@@ -90,11 +67,10 @@ def parse_args() -> argparse.Namespace:


 def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
-    is_dlo = args.mode in {"dlo", "dlo-no-allgather"}
    common: dict[str, Any] = {
        "model": args.model,
        "trust_remote_code": True,
-        "num_gpus": args.dp_size if is_dlo else args.tp_size,
+        "num_gpus": 2,
        "usp": 1,
        "ring": 1,
        "vae_parallel_mode": "tile",
@@ -105,22 +81,22 @@ def engine_kwargs(args: argparse.Namespace) -> dict[str, Any]:
        "stage_init_timeout": args.init_timeout,
        "init_timeout": args.init_timeout,
    }
-    if is_dlo:
+    if args.mode in {"dlo-dp2", "dlo-dp2-no-allgather"}:
        common.update(
            tensor_parallel_size=1,
-            data_parallel_size=args.dp_size,
+            data_parallel_size=2,
            text_encoder_tp_size=1,
            vae_patch_parallel_size=1,
            enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=args.mode == "dlo",
+            dlo_use_allgather=args.mode == "dlo-dp2",
            dlo_resident_layers=0,
        )
    else:
        common.update(
-            tensor_parallel_size=args.tp_size,
+            tensor_parallel_size=2,
            data_parallel_size=1,
-            text_encoder_tp_size=args.tp_size,
-            vae_patch_parallel_size=args.tp_size,
+            text_encoder_tp_size=2,
+            vae_patch_parallel_size=2,
            enable_distributed_layerwise_offload=False,
        )
    return common
@@ -203,8 +179,7 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
        "engine_kwargs": engine_kwargs(args),
    }
    try:
-        uses_dp_wave = args.mode == "dlo" and args.dp_size > 1
-        if uses_dp_wave:
+        if args.mode == "dlo-dp2":
            started = time.perf_counter()
            asymmetric = await asyncio.gather(
                generate_one(
@@ -214,15 +189,12 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
                    prompt="",
                    seed=1001,
                ),
-                *(
-                    generate_one(
-                        engine,
-                        args,
-                        request_id=f"invalid-peer-{index}",
-                        prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
-                        seed=1001 + index,
-                    )
-                    for index in range(1, args.dp_size)
+                generate_one(
+                    engine,
+                    args,
+                    request_id="invalid-peer",
+                    prompt=DEFAULT_PROMPTS[0],
+                    seed=1002,
                ),
                return_exceptions=True,
            )
@@ -230,18 +202,18 @@ async def run(args: argparse.Namespace) -> dict[str, Any]:
            summary["asymmetric_errors"] = [
                f"{type(result).__name__}: {result}" for result in asymmetric if isinstance(result, BaseException)
            ]
-            if len(summary["asymmetric_errors"]) != args.dp_size:
-                raise RuntimeError("Every request in the asymmetric DP wave must fail before dispatch")
+            if len(summary["asymmetric_errors"]) != 2:
+                raise RuntimeError("Both requests in the asymmetric DP wave must fail before dispatch")

        started = time.perf_counter()
-        request_count = args.dp_size if uses_dp_wave else 1
+        request_count = 2 if args.mode == "dlo-dp2" else 1
        outputs = await asyncio.gather(
            *(
                generate_one(
                    engine,
                    args,
                    request_id=f"recovery-{index}",
-                    prompt=DEFAULT_PROMPTS[index % len(DEFAULT_PROMPTS)],
+                    prompt=DEFAULT_PROMPTS[index],
                    seed=2000 + index,
                )
                for index in range(request_count)
@@ -74,7 +74,7 @@ if __name__ == "__main__":
    omni = Omni(model="Qwen/Qwen-Image")
    prompt = "a cup of coffee on the table"
    outputs = omni.generate(prompt)
-    images = outputs[0].images
+    images = outputs[0].request_output.images
    images[0].save("coffee.png")
 ```

@@ -265,7 +265,7 @@ if __name__ == "__main__":
    ]
    outputs = omni.generate(prompts)
    for i, output in enumerate(outputs):
-        output.images[0].save(f"{i}.jpg")
+        output.request_output.images[0].save(f"{i}.jpg")
 ```

 !!! info
@@ -301,7 +301,7 @@ if __name__ == "__main__":
        }
    ])
    for i, output in enumerate(outputs):
-        output.images[0].save(f"{i}.jpg")
+        output.request_output.images[0].save(f"{i}.jpg")
 ```

 You can also pass a negative prompt via the CLI argument `--negative-prompt`:
@@ -16,7 +16,6 @@ list of supported architectures across all modalities, see
 |---|---|---|---|---|---|---|
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
-| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
 | GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
 | Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
 | Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
@@ -396,58 +395,6 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \

 ---

-## Gepard-1.0
-
-Single-stage native AR TTS at 22.05 kHz. Pipeline: `Qwen3.5 backbone → 32 FSQ codebook heads (one frame/step) → NeMo NanoCodec`. Backbone runs under vLLM paged attention; the 32-head sampling + learned embedding feedback ride the native-AR runner hooks. `enforce_eager` (CUDA graph is a perf follow-up).
-
-### Prerequisites
-The NanoCodec decoder needs NeMo (NVIDIA Open Model License), installed separately:
-
-```bash
-uv pip install "nemo_toolkit[tts]==2.7.3" \
-    --overrides <(printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0")
-```
-
-The overrides are what make that command work. `nemo_toolkit[tts]` declares
-`transformers~=4.57.0`, so a plain install downgrades the version
-`requirements/common.txt` pins and takes `huggingface-hub` back to 0.x with it — after
-which `import vllm` fails. Both packages are named here rather than left to resolve on
-their own: transformers 5.x happens to require `huggingface-hub` 1.x, but relying on
-that would make the second repair an accident of the first. `--constraint` cannot be
-used instead, since it narrows a version range rather than overriding a dependency a
-package declares, and the two sets never intersect. NeMo 2.7.3 decodes correctly
-against transformers 5.x, so nothing is lost by holding it there.
-
-(The same clash is why this is documented here rather than declared in
-`pyproject.toml`: an extra resolves together with the base dependencies, so
-`vllm-omni[gepard]` would be unsatisfiable rather than opt-in. SoulX-Singer English
-SVS documents `nemo_toolkit[asr]` the same way.)
-
-On a host whose CUDA toolkit cannot build kernels for the GPU — no `nvcc`/`ninja`, or a
-consumer Blackwell (`sm_120`) card — also set `VLLM_USE_FLASHINFER_SAMPLER=0`. vLLM's
-sampler JIT-compiles through FlashInfer on first use, and where that build fails the engine
-dies inside `profile_run` long after every import has succeeded, which reads as a model
-failure rather than an environment one. The PyTorch sampler it falls back to costs nothing
-here.
-
-### Quick start (zero-shot, default voice)
-```bash
-python examples/offline_inference/text_to_speech/gepard/end2end.py \
-    --text "Hello, this is a Gepard demo."
-```
-
-### Voice cloning
-Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Reference-audio cloning (`ref_compressor` speaker prefix) is a follow-up PR.
-
-### Notes
- Output: 22.05 kHz mono WAV.
- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR.
- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it).
- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`.
- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced.
- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed.
-
 ## VoxCPM2

 Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
@@ -1,139 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Offline inference example for Gepard-1.0 TTS via vLLM-Omni.
-
-Single-stage native-AR pipeline: a Qwen3.5 backbone samples one 32-code FSQ
-frame per step; the NeMo NanoCodec decodes the frames to a 22.05 kHz mono
-waveform. Zero-shot: the default learned voice, no reference audio.
-
-Generation length and seed are stage settings rather than flags here. One
-output token is one audio frame, so ``max_tokens`` in the deploy YAML is the
-frame budget; pass a copy of the YAML via ``--deploy-config`` to change it.
-
-Usage:
-  python end2end.py --text "Hello, this is Gepard speaking."
-"""
-
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-import soundfile as sf
-import torch
-
-import vllm_omni
-from vllm_omni.utils.tracking_parser import TrackingArgumentParser
-
-os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
-
-from vllm_omni import Omni  # noqa: E402
-
-MODEL = "nineninesix/gepard-1.0"
-SAMPLE_RATE = 22050
-# Absolute: the loader hands the path straight to OmegaConf, which resolves a
-# bare name against the current working directory, not the package.
-DEFAULT_DEPLOY_CONFIG = str(Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml")
-
-
-def build_request(text: str) -> dict:
-    """Build an Omni request payload for Gepard (zero-shot, no ref audio).
-
-    The model only consumes the [speaker slots, SOT, text, EOT, SOS] layout, so
-    a bare text prompt would have no speaker slots and no SOS.
-    """
-    from transformers import AutoTokenizer
-
-    from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
-    from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
-
-    # From the checkpoint's own sidecar, which is what the worker builds its
-    # config from too. A layout that disagrees with the trained one does not
-    # fail, it just degrades the speech.
-    cfg = GepardConfig.from_checkpoint(MODEL)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
-    prompt_token_ids = build_gepard_prompt_ids(
-        tokenizer(text, add_special_tokens=False)["input_ids"],
-        config=cfg,
-    )
-    return {
-        "prompt_token_ids": prompt_token_ids,
-        "additional_information": {"text": [text]},
-    }
-
-
-def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = SAMPLE_RATE) -> None:
-    sf.write(path, waveform.float().numpy(), sample_rate)
-    seconds = waveform.numel() / sample_rate
-    print(
-        f"  Saved {path} ({tuple(waveform.shape)}, {sample_rate} Hz, "
-        f"{seconds:.2f}s, peak {float(waveform.abs().max()):.3f})"
-    )
-
-
-def main(args) -> None:
-    omni = Omni(
-        model=MODEL,
-        deploy_config=args.deploy_config,
-        stage_init_timeout=args.stage_init_timeout,
-        init_timeout=args.init_timeout,
-    )
-
-    # No explicit SamplingParams: one supplied here would replace the stage
-    # defaults rather than merge over them, dropping the pipeline's stop token.
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    print(f"Synthesizing: {args.text!r}")
-    inputs = build_request(text=args.text)
-
-    for i, stage_outputs in enumerate(omni.generate(inputs)):
-        for j, out in enumerate(stage_outputs.outputs):
-            mm = out.multimodal_output
-            if mm is None:
-                print(f"  [req {i}] No audio output.")
-                continue
-            # The consolidation path renames "model_outputs" to "audio".
-            audio = mm.get("audio")
-            if audio is None:
-                audio = mm.get("model_outputs")
-            if isinstance(audio, list):
-                audio = audio[0] if len(audio) else None
-            if audio is None:
-                print(f"  [req {i}] No waveform in multimodal_output (keys: {list(mm)}).")
-                continue
-            sr_tensor = mm.get("sr")
-            sr = int(sr_tensor.item()) if hasattr(sr_tensor, "item") else SAMPLE_RATE
-            save_audio(audio.cpu(), str(output_dir / f"output_{i}_{j}.wav"), sr)
-
-    print("Done.")
-
-
-def parse_args():
-    parser = TrackingArgumentParser(description="Gepard-1.0 offline TTS inference (zero-shot)")
-    parser.add_argument("--text", default="Hello, this is Gepard speaking.", help="Text to synthesize.")
-    parser.add_argument(
-        "--output-dir",
-        default=os.path.join(
-            os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
-            "gepard_output",
-        ),
-        help="Directory for WAV outputs (default: ~/.cache/gepard_output).",
-    )
-    parser.add_argument(
-        "--deploy-config",
-        default=DEFAULT_DEPLOY_CONFIG,
-        help=(
-            "Path to the deploy YAML (default: the packaged vllm_omni/deploy/gepard.yaml). This must "
-            "stay set: without the YAML's `pipeline: gepard` pin the architectures fallback routes "
-            "Qwen3_5ForCausalLM to the diffusion registry. Copy it to change max_tokens or seed."
-        ),
-    )
-    # The upstream defaults are too tight for a cold start.
-    parser.add_argument("--stage-init-timeout", type=int, default=900)
-    parser.add_argument("--init-timeout", type=int, default=1800)
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    main(parse_args())
@@ -570,8 +570,6 @@ def main():
        num_frames=args.num_frames,
        extra_args=extra_args,
    )
-    if args.frame_rate is not None:
-        sampling_kwargs["frame_rate"] = args.frame_rate
    if args.guidance_scale_high is not None:
        sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high

@@ -81,10 +81,6 @@ plugins:
        user_guide/quantization/index.md: user_guide/quantization/overview.md
        design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
        design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
-        design/feature/cpu_offloading.md: design/feature/offloader/README.md
-        design/feature/model_level_offload.md: design/feature/offloader/cpu_offload.md
-        design/feature/layerwise_offload.md: design/feature/offloader/layerwise_offload.md
-        design/feature/distributed_layerwise_offload.md: design/feature/offloader/distributed_layerwise_offload.md
  - glightbox
  - git-revision-date-localized:
      # exclude files
@@ -40,7 +40,6 @@ dev = [
    "pytest-asyncio==1.4.0",
    "pytest-cov==7.1.0",
    "pytest-mock==3.15.1",
-    "pytest-shard==0.1.2",
    "websockets==16.1",
    "pytest-subtests==0.15.0",
    "datasets==3.6.0",
@@ -302,21 +301,6 @@ filterwarnings = [
    "ignore:.*does not have '__bases__' attribute.*:UserWarning",
 ]

-[tool.coverage.run]
-core = "sysmon"
-source = ["vllm_omni"]
-# Online-serving tests launch the server via subprocess.Popen and stop it with
-# SIGTERM. pytest-cov no longer instruments subprocesses implicitly, and a
-# SIGTERM-ed process discards its data, so without patch/sigterm a coverage run
-# over an online test measures only the pytest parent.
-patch = ["subprocess"]
-concurrency = ["multiprocessing", "thread"]
-sigterm = true
-
-[tool.coverage.report]
-show_missing = true
-skip_covered = true
-
 [tool.typos.default]
 extend-ignore-identifiers-re = [
    ".*_thw",
@@ -1,175 +0,0 @@
-# LTX-2.5
-
-> Text-to-video and first-frame image-to-video generation with synchronized audio
-
-## Overview
-
-vLLM-Omni supports the gated
-[`Lightricks/LTX-2.5-Diffusers`](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
-checkpoint for T2V and first-frame I2V. Every output includes synchronized
-48 kHz stereo audio.
-
-The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
-repository is not directly loadable with `--model`. It supplies the official
-upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
-both model licenses and authenticate before first use.
-
-## Choose a pipeline
-
-| Pipeline | Mode | Default size | Schedule |
-|---|---|---:|---:|
-| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
-| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
-| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
-| `LTX2DistilledTwoStagePipeline` | Distilled two-stage | 1920x1088 | 8 + 3 steps |
-
-Both two-stage pipelines generate at half resolution, apply the official x2
-latent upsampler, and run a three-step refinement stage. Select the class with
-`--model-class-name`; no `--task-type` flag is required. Supplying one initial
-image selects I2V, while omitting it selects T2V.
-
-## Prerequisites
-
-```bash
-hf auth login
-export MODEL=Lightricks/LTX-2.5-Diffusers
-```
-
-Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
-`ffprobe` are on `PATH`. I2V requires PyAV backed by an FFmpeg build with the
-`libx264` encoder.
-
-## Hardware
-
-| GPU | Status | Recommended scope |
-|---|---|---|
-| NVIDIA B300 | Verified | All four canonical pipelines |
-| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
-| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
-| NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 |
-
-The 1920x1088 two-stage examples require about 114 GB of peak GPU memory.
-80 GB GPUs do not have enough safety margin for the canonical two-stage
-configuration; use a larger GPU or reduce the output configuration.
-
-### H100 80 GB FP8
-
-Use the distilled one-stage pipeline with FP8 and cuDNN attention:
-
-```bash
-vllm serve Lightricks/LTX-2.5-Diffusers \
-  --omni \
-  --model-class-name LTX2DistilledOneStagePipeline \
-  --quantization fp8 \
-  --diffusion-attention-backend CUDNN_ATTN \
-  --host 0.0.0.0 \
-  --port 8000
-```
-
-Use the T2V or I2V request below after the server is ready.
-
-## Offline inference
-
-Choose values from the pipeline table. For example, the distilled two-stage
-path uses:
-
-```bash
-export PIPELINE=LTX2DistilledTwoStagePipeline
-export WIDTH=1920
-export HEIGHT=1088
-export STEPS=8
-```
-
-### Text-to-video
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-  --model "${MODEL}" \
-  --model-class-name "${PIPELINE}" \
-  --prompt "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot." \
-  --width "${WIDTH}" \
-  --height "${HEIGHT}" \
-  --num-frames 121 \
-  --num-inference-steps "${STEPS}" \
-  --frame-rate 24 \
-  --fps 24 \
-  --seed 42 \
-  --output ltx25-t2v.mp4
-```
-
-### First-frame image-to-video
-
-```bash
-python examples/offline_inference/image_to_video/image_to_video.py \
-  --model "${MODEL}" \
-  --model-class-name "${PIPELINE}" \
-  --image /absolute/path/to/first-frame.png \
-  --prompt "The red fox walks forward while the camera tracks alongside." \
-  --width "${WIDTH}" \
-  --height "${HEIGHT}" \
-  --num-frames 121 \
-  --num-inference-steps "${STEPS}" \
-  --frame-rate 24 \
-  --fps 24 \
-  --seed 42 \
-  --output ltx25-i2v.mp4
-```
-
-LTX-2.5 uses the official CRF-18 first-frame conditioning path by default.
-
-## Online serving
-
-Start one server for the selected pipeline:
-
-```bash
-vllm serve "${MODEL}" \
-  --omni \
-  --model-class-name "${PIPELINE}" \
-  --host 0.0.0.0 \
-  --port 8000 \
-  --stage-init-timeout 900
-```
-
-T2V request:
-
-```bash
-curl -sS --fail-with-body \
-  -X POST http://127.0.0.1:8000/v1/videos/sync \
-  -F 'prompt=A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot.' \
-  -o ltx25-online-t2v.mp4
-```
-
-For I2V, add exactly one first frame:
-
-```bash
-export FIRST_FRAME=/absolute/path/to/first-frame.png
-
-curl -sS --fail-with-body \
-  -X POST http://127.0.0.1:8000/v1/videos/sync \
-  -F "input_reference=@${FIRST_FRAME};type=image/png" \
-  -F 'prompt=The red fox walks forward while the camera tracks alongside.' \
-  -o ltx25-online-i2v.mp4
-```
-
-Restart the server after changing `PIPELINE` because each class selects a
-different weight layout and execution topology.
-
-## Constraints
-
- Online `width`, `height`, `num_frames`, `fps`, and `num_inference_steps`
-  are optional and use the selected pipeline defaults when omitted. `seed` is
-  optional and random when omitted.
- `num_frames` must be `8k+1` when overridden.
- One-stage width and height must be divisible by 32; two-stage final
-  dimensions must be divisible by 64.
- Full/SFT pipelines accept negative prompts. Distilled pipelines are
-  positive-only and reject negative prompts.
- I2V accepts exactly one initial image per prompt.
- The supported `--model` value is `Lightricks/LTX-2.5-Diffusers`.
-
-## References
-
- [LTX-2.5-Diffusers model card](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
- [Official LTX-2 repository](https://github.com/Lightricks/LTX-2)
- [Generic T2V example](../../examples/offline_inference/text_to_video/text_to_video.py)
- [Generic I2V example](../../examples/offline_inference/image_to_video/image_to_video.py)
@@ -8,12 +8,10 @@
 |---|---|---|
 | `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
 | `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
-| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
-| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
+| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
 | `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
 | `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
-| `LTX2DistilledOneStagePipeline` | LTX-2.3 merged-distilled one-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` |
-| `LTX2DistilledTwoStagePipeline` | LTX-2.3 merged-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
+| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |

 Repositories in the table are download units. A full pipeline repository
 contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
@@ -28,12 +26,9 @@ image selects I2V. Both one-stage repositories declare this class, so
 the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.

 `LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
-then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a
-merged distilled Transformer without upsampling, while
-`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support
-T2V and I2V; select their class explicitly. The deprecated
-`LTX2DistilledPipeline` name remains an alias for
-`LTX2DistilledTwoStagePipeline`.
+then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
+merged distilled Transformer in both stages. Both entries support T2V and I2V;
+select their class explicitly.

 ## API Migration

@@ -57,8 +52,8 @@ The consolidation also removes these registry names without aliases:
 | `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
 | `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
 | `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
-| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` |
-| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` |
+| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
+| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |

 Passing any second positional argument now raises `TypeError`. These changes
 affect direct Python callers and explicit `--model-class-name` overrides;
@@ -117,10 +112,10 @@ Start a two-stage checkpoint with its explicit class:

 ```bash
 vllm serve rootonchair/LTX-2-19b-distilled --omni \
-  --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
+  --model-class-name LTX2DistilledPipeline --stage-init-timeout 600
 # LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
 vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
-  --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
+  --model-class-name LTX2DistilledPipeline --stage-init-timeout 600
 # Ordinary LTX-2.3 two-stage
 vllm serve diffusers/LTX-2.3-Diffusers --omni \
  --model-class-name LTX2TwoStagePipeline \
@@ -1,195 +0,0 @@
-# MiniMax-H3 on a single NPU 950PR
-
-This recipe runs MiniMax-H3 with online INT8 quantization on one NPU 950PR
-NPU (128 GB HiBL 1.0 HBM). It covers the single-card T2VA configuration at
-1024x576. For the eight-card Atlas 800I A3 BF16 route at 768P, see
-[MiniMax-H3-NPU.md](MiniMax-H3-NPU.md).
-
-## Capacity requirements
-
-| Resource | Requirement |
-| --- | ---: |
-| NPU | 1x NPU 950PR |
-| NPU HBM | 128 GiB (131,072 MiB reported by `npu-smi`) |
-| Observed HBM high-water mark | 118,442 MiB (90.4% of capacity) |
-| Checkpoint storage | 135 GiB per partition, local disk strongly preferred |
-| Container shared memory | 8 GiB, or a writable local filesystem for spill |
-
-The observed high-water mark leaves about 12.3 GiB of headroom, so this card is
-sized for exactly one resident task partition at 1024x576. Do not expect a
-larger output shape or concurrency greater than one to fit.
-
-Because a single 950PR holds the whole partition in HBM, this route does not
-need layerwise offload, and the 200 GiB system-RAM floor carried by the 72 GiB
-GPU recipes does not apply. The validated run completed inside a container
-limited to 32 GiB of RAM. Passing `--enable-layerwise-offload` on this card is
-actively harmful: it stages weights in non-reclaimable host memory and the
-container's OOM killer terminates the server with exit code -9.
-
-
-## Environment
-
- Host architecture: x86_64
-  driver: 25.7.rc1.6 (hal 7.35.23)
-  firmware: 9.0.0.105.229
- CANN toolkit: 9.1.0 (`/usr/local//cann-9.1.0`)
- npu-smi: 25.7.rc1.6
- Python: 3.12.13
- PyTorch: 2.10.0+cpu
- torch_npu: 2.10.0.post2
- vLLM: 0.26.0
- vLLM-Omni: 0.26.1.dev103+g584d78c67.npu (commit `584d78c6`)
-
-Install vLLM-Omni from a checkout with MiniMax-H3 support:
-
-```bash
-uv venv
-source .venv/bin/activate
-uv pip install -e .
-```
-
-## Start a server
-
-```bash
-export MODEL=/path/to/MiniMax-H3/FL2VA
-export PORT=8000
-export VLLM_WORKER_MULTIPROC_METHOD=spawn
-export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400
-
-vllm serve "${MODEL}" \
-  --omni \
-  --host 0.0.0.0 \
-  --port "${PORT}" \
-  --trust-remote-code \
-  --num-gpus 1 \
-  --tensor-parallel-size 1 \
-  --usp 1 \
-  --ring 1 \
-  --text-encoder-tp-size 1 \
-  --vae-patch-parallel-size 1 \
-  --vae-parallel-mode tile \
-  --vae-use-tiling \
-  --init-timeout 14400 \
-  --stage-init-timeout 14400 \
-  --quantization int8
-```
-
-Do not pass `--diffusion-attention-backend CUDNN_ATTN`. cuDNN attention has no
-NPU implementation and fails at the first denoise step. Leave the backend
-unset: it resolves to `TORCH_SDPA`, or to `FLASH_ATTN` when MindIE-SD is
-installed (see [MiniMax-H3-NPU.md § Environment](MiniMax-H3-NPU.md#environment)).
-
-H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
-
-## Validated evidence
-
-Measured on one NPU 950PR with the server command above plus
-`--enable-diffusion-pipeline-profiler`, generating a 5 s 1024x576 T2VA clip from
-the request in [§ T2VA request example](#t2va-request-example). MiniMax-H3
-requested 60 denoise steps and executed 59 denoise updates, so per-step latency
-is `denoise / 59`. `ffprobe` confirms every returned MP4 is 1024x576, 24 fps,
-124 video frames, 5.175 s, with a 32 kHz audio track of 165,600 samples.
-
-Five requests were issued back to back with no idle time between them. The card
-is measurably faster on the first request than in steady state, so both regimes
-are reported.
-
-| Stage | First request, cold card | Thermal steady state |
-| --- | ---: | ---: |
-| End-to-end request | 462.13 s | 506.38 s |
-| Text encode | 0.55 s | 0.05 s |
-| Denoise (59 updates) | 429.29 s | 471.74 s |
-| Per denoise update | 7,276 ms | 7,996 ms |
-| VAE decode | 8.43 s | 8.14 s |
-| Worker-to-server handoff | 24 s | 27 s |
-| MP4 muxing | 0.57 s | 0.50 s |
-| Server-reported `denoise_step_latency_ms` | 7,702 ms | 8,440 ms |
-| Peak HBM (`npu-smi`) | 118,442 MiB | 118,442 MiB |
-| Average NPU power | 566.9 W | 556.9 W |
-| Peak NPU temperature | 104 C | 104 C |
-
-Sample counts differ per row. The cold column is a single request. In the steady
-column, text encode, denoise and VAE decode are means over the four subsequent
-requests — denoise spanned 470.37 to 473.21 s, a 0.60% spread — while
-end-to-end, handoff and muxing are means over the two of those four that
-survived the handoff timeout described above.
-
-**Report the steady-state column.** The cold card completes denoise 9.9% faster
-while drawing only 1.8% more average power, and both regimes reach the same
-104 C ceiling, so the first request is buying a short window of higher clocks
-before the die saturates. Any benchmark that issues a single request against an
-idle card will overstate throughput by about 10%.
-
-**Do not read `denoise_step_latency_ms` as a per-step time.** The server divides
-the entire stage wall time by the 60 *requested* steps, so it absorbs text
-encoding, VAE decode and the 837 MiB handoff. All three successful requests
-satisfy `denoise_step_latency_ms == e2e_stage_wall_time_ms / 60` exactly. The
-per-update figure in the table is `diffuse / 59` from the pipeline profiler.
-
-Denoise accounts for 93% of the steady-state request and the handoff accounts
-for 5.3%, so keeping the transfer in shared memory should bring end-to-end down
-to roughly 480 s.
-
-Peak HBM is sampled externally with `npu-smi` at a 3 s interval. It is a
-reserved high-water mark rather than live allocation: it reads 111,272 MiB after
-weight loading and before the first request, and does not fall back after
-requests complete. Re-measure for longer outputs, a different output shape, or
-concurrency greater than one.
-
-### Startup
-
-Weight loading took **31 min 50 s** on the validated host (first log line
-08:04:59, `Application startup complete` 08:36:49), with the checkpoint on a
-shared GlusterFS network mount. This is why the server command sets
-`--init-timeout 14400` and `--stage-init-timeout 14400`; the stock 600 s and
-1800 s timeouts both abort mid-load with `TimeoutError` and exit code 143.
-Reduce both timeouts when the partition is staged on local disk.
-
-## T2VA request example
-
-```bash
-export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
-
-curl -sS --max-time 1800 -X POST "${API_URL}" \
-  -F 'prompt=At night, three cats march into a bedroom playing tiny brass instruments, then abruptly file out, with synchronized room ambience.' \
-  -F 'width=1024' \
-  -F 'height=576' \
-  -F 'aspect_ratio=16:9' \
-  -F 'fps=24' \
-  -F 'num_inference_steps=60' \
-  -F 'flow_shift=12' \
-  -F 'seed=1101' \
-  -F 'extra_params={"task":"t2va","duration":5,"audio_flow_shift":3.0}' \
-  -o t2va.mp4
-```
-
-## Known limitations
-
- INT8 online quantization is validated for T2VA on this card. Use BF16 for
-  FL2VA and Ref2VA, or re-measure before relying on it.
- Single-card serving loads one task partition at a time. For Ref2VA, stop the
-  server and restart it against the `Ref2VA` directory.
- Sustained load drives the die to 104 C and costs about 10% of denoise
-  throughput relative to a cold card. Warm the server with at least one full
-  request before measuring, and treat steady state as the reportable number.
- Requests can fail during result handoff on hosts where the IPC spill path is
-  slow. See [§ Result handoff and shared
-  memory](#result-handoff-and-shared-memory).
- Loading the checkpoint from a network mount takes over 30 minutes and forces
-  very large init timeouts. Stage the partition on local disk when possible.
- The configuration measured here is 1024x576 at 60 steps, which differs from
-  the 1344x768 at 50 steps used by the GPU recipes in this directory. Denoise
-  cost scales superlinearly with token count, so the numbers are not directly
-  comparable across recipes.
- The image ships a `triton` package whose  backend is not built
-  (`No module named 'triton._C.libtriton.'`). vLLM logs this as an error
-  at startup and disables Triton; the diffusion path does not need it and the
-  run completes normally.
-
-## Additional resources
-
- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — eight-card Atlas 800I A3 BF16 guide
- [Int8 quantization](../../docs/user_guide/quantization/int8.md)
- [Supported models](../../docs/models/supported_models.md)
- [Video API](../../docs/serving/videos_api.md)
@@ -357,19 +357,12 @@ No restart is needed: `task=fl2va` routes to `FL2VA/transformer`, while

 ### Online FP8 quantization

-MiniMax H3 supports online FP8 quantization of both the DiT and the Qwen3-VL
-text decoder. The checkpoint remains BF16 on disk; vLLM-Omni creates FP8
-weights at runtime and uses dynamic activation scaling during inference. By
-default, `--quantization fp8` quantizes eligible attention and MLP linears in
-the text decoder, token refiner, and main DiT blocks, as well as the DiT
-condition and AdaLN projections. The Qwen vision tower, embeddings, norms,
-RoPE, both VAEs, and the model's FP32 patch, timestep, and output projections
-keep checkpoint precision.
-
-To select a component, use `--diffusion-quantization-config` with
-`{"transformer":{"method":"fp8"}}` for DiT-only FP8 or
-`{"text_encoder":{"method":"fp8"}}` for text-decoder-only FP8. The two
-entries can be combined. The shorthand below enables both components.
+MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
+remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
+uses dynamic activation scaling during inference. By default, attention and
+MLP linears in the token refiner and main DiT blocks, the condition
+projection, and all AdaLN projections use FP8. Patch, timestep, and final
+projections remain FP32; the text encoder and VAEs are unchanged.

 Add this option to an existing H3 server command:

@@ -377,28 +370,6 @@ Add this option to an existing H3 server command:
 --quantization fp8
 ```

-#### Single 96 GB GPU, no-offload capacity check
-
-Use the FL2VA-only partition for this capacity test. Loading the combined
-service would also load the Ref2VA DiT and would test a different memory
-budget. A no-offload capacity check should contain none of
-`--enable-cpu-offload`, `--enable-layerwise-offload`, or
-`--enable-distributed-layerwise-offload`. VAE tiling changes decode placement
-but does not offload model weights to the CPU.
-
-The run passes the capacity check when the server initializes, the request
-finishes without CUDA OOM or Xid errors, `peak_used_mib` remains below the
-card's reported `memory_total_mib`, and `ffprobe` reports H.264 video plus
-32 kHz stereo AAC audio. Report the measured headroom rather than assuming
-that every nominal 96 GB SKU exposes the same MiB total.
-
-As a capacity proxy only, the same five-second case on one B300 measured a
-92,946 MiB whole-device first-request peak and a 92,146 MiB worker peak. Its
-encode, diffuse, decode, and client wall times were 8.664 s, 134.504 s,
-6.016 s, and 151.583 s. These numbers suggest that a 96 GB card may fit, but
-they are not an RTX PRO 6000 validation: kernels, allocator behavior, and
-reported device capacity differ across GPUs.
-
 Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
 resolves the `transformer` component before constructing the DiT, so names do
 not start with `transformer.`. Entries are exact runtime linear prefixes; a
@@ -415,10 +386,10 @@ For example, keep the first main block's attention projections in BF16 with:
  '{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
 ```

-The structured option replaces `--quantization fp8`. Online FP8 can be used
-with H3 layerwise offload and with distributed layerwise offload's full-weight
-per-rank path (`--dlo-no-use-allgather`). The sharded DLO AllGather path is not
-supported for runtime-created FP8 weights.
+The structured option replaces `--quantization fp8`. Online FP8 is currently
+incompatible with H3 layerwise offload because the offload path produces a
+weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
+parallelism and VAE tiling instead.

 ## AMD ROCm (gfx942 / gfx950)

@@ -779,13 +750,6 @@ audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
 resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
 reduction.

-For direct human inspection, see the external
-[BF16 versus global-FP8 comparison](https://lishunyang12.github.io/vllm-omni-rankings/scripts/minimax_h3_global_fp8_vs_bf16/),
-which includes matched five-second T2VA, I2VA, and Ref2VA videos. The
-[comparison sources](https://github.com/lishunyang12/vllm-omni-rankings/tree/main/scripts/minimax_h3_global_fp8_vs_bf16)
-also record per-task fidelity metrics and provenance without storing generated
-media in this repository.
-
 ## TeaCache acceleration

 TeaCache reuses DiT block residuals across denoising steps when consecutive
@@ -870,8 +834,7 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
  H3 native `tile` mode only.
 - A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
  pure Ulysses.
- Runtime-created FP8 weights do not support the sharded DLO AllGather path;
-  use `--dlo-no-use-allgather` when combining online FP8 with DLO.
+- Online FP8 is not compatible with layerwise offload.
 - TeaCache and Cache-DiT cannot be enabled on the same server.
 - Image+audio Ref2VA accepts exactly one image and one audio reference.
 - Video Ref2VA accepts one or more video files, but not an additional standalone
@@ -55,7 +55,6 @@ recipes/
 | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
 | [`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) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
-| [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
 | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
@@ -1,15 +1,15 @@
 # Common dependencies for all platforms
-# transformers: required unified multimodal model classes are unavailable
-# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction
-# regression. transformers 5.9 also changed
-# _adjust_missing_and_unexpected_keys from
+# transformers: kept exactly aligned with the upstream vLLM constraint.
+# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
+# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
+# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
 # list+list to set|set, breaking models that declare
 # _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
 # with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
 # That breakage is now handled in code instead of by pinning:
 # MossTTSNanoForGeneration wraps weight loading in
 # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
-transformers >= 5.10.1, < 5.15
+transformers >= 5.5.3
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
@@ -5,7 +5,6 @@ import sys
 from pathlib import Path

 import pytest
-import yaml

 sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "common" / "scripts"))

@@ -39,6 +38,8 @@ def _render(changed_files: list[str]) -> str:


 def test_bootstrap_if_injected_by_step_key() -> None:
+    import yaml
+
    rendered = _render_bootstrap_pipeline(
        BOOTSTRAP_STEPS_TEMPLATE,
        decision=resolve_ci_decision([]),
@@ -138,65 +139,3 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
    assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [
        {"name": "swr-secret"},
    ]
-
-
-def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
-    rendered = _render_test_pipeline(doc, changed_files=changed_files)
-    labels: set[str] = set()
-
-    def walk(steps: list | None) -> None:
-        for step in steps or []:
-            if not isinstance(step, dict):
-                continue
-            if "label" in step:
-                labels.add(step["label"])
-            walk(step.get("steps"))
-
-    walk(rendered.get("steps"))
-    return labels
-
-
-# Synthetic coverage-style job: shared inputs that change what the split measures.
-_COVERAGE_SHARED_INPUTS_DOC = {
-    "steps": [
-        {
-            "label": "Coverage Pilot",
-            "source_file_dependencies": [
-                "tests/e2e/online_serving/test_example.py",
-                ".buildkite/common/scripts/run_cov_split.sh",
-                "pyproject.toml",
-            ],
-            "commands": [".buildkite/common/scripts/run_cov_split.sh --model-id example"],
-        },
-        {
-            "label": "Unrelated Model Test",
-            "source_file_dependencies": [
-                "tests/e2e/online_serving/test_other.py",
-            ],
-            "commands": ["pytest -sv tests/e2e/online_serving/test_other.py"],
-        },
-    ],
-}
-
-
-@pytest.mark.parametrize(
-    "changed_file",
-    [
-        ".buildkite/common/scripts/run_cov_split.sh",
-        "pyproject.toml",
-    ],
-)
-def test_coverage_shared_inputs_select_dependent_job(changed_file: str) -> None:
-    """Jobs that list coverage shared inputs must stay selected when those files change."""
-    labels = _surviving_labels(_COVERAGE_SHARED_INPUTS_DOC, [changed_file])
-    assert "Coverage Pilot" in labels
-    assert "Unrelated Model Test" not in labels
-
-
-def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None:
-    labels = _surviving_labels(
-        _COVERAGE_SHARED_INPUTS_DOC,
-        ["vllm_omni/entrypoints/openai/serving_chat.py"],
-    )
-    assert "Coverage Pilot" not in labels
-    assert "Unrelated Model Test" not in labels
@@ -14,8 +14,7 @@ import pytest
 from transformers import PretrainedConfig, Qwen3OmniMoeConfig

 from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
-from vllm_omni.config import config_factory as config_factory_module
-from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
+from vllm_omni.config.config_factory import StageConfigFactory
 from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
 from vllm_omni.config.omni_config import VllmOmniConfig
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
@@ -56,7 +55,6 @@ def clear_config_factory_caches():
    yield
    StageConfigFactory.get_hf_config.cache_clear()
    StageConfigFactory.try_infer_model_type.cache_clear()
-    _materialize_object_storage_configs.cache_clear()


 Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
@@ -2767,89 +2765,3 @@ class TestPipelineConfigResolvers:
            pass

        assert resolver(NotTheRightHfConfig()) is None
-
-
-class TestObjectStorageConfigResolution:
-    """Object-storage URIs must be materialized locally before any HF-style reads.
-
-    Regression coverage for review on the Run:AI PR: ``get_config`` and the
-    config.json/model_index.json fallbacks crashed with ``HFValidationError``
-    for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so
-    a bucket named after another pipeline could hijack pipeline selection.
-    """
-
-    @pytest.fixture
-    def fake_object_storage(self, monkeypatch, tmp_path):
-        """Replace ObjectStorageModel with a fake that "pulls" into a tmp dir.
-
-        Returns a recorder: calls to the returned callable write files into
-        the materialized directory, and ``recorder.pulls`` records every
-        ``pull_files`` invocation.
-        """
-        pulls: list[tuple[str, list[str] | None, list[str] | None]] = []
-        materialized_files: list[tuple[str, str]] = []
-
-        class FakeObjectStorageModel:
-            def __init__(self, url: str):
-                self.dir = str(tmp_path)
-
-            def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None):
-                pulls.append((model_path, allow_pattern, ignore_pattern))
-                for name, content in materialized_files:
-                    (tmp_path / name).write_text(content)
-
-        monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel)
-
-        class Recorder:
-            def set_files(self, files: list[tuple[str, str]]) -> None:
-                materialized_files.clear()
-                materialized_files.extend(files)
-
-        recorder = Recorder()
-        recorder.pulls = pulls
-        return recorder
-
-    def test_passthrough_for_non_uri(self, fake_object_storage):
-        assert _materialize_object_storage_configs("org/model") == "org/model"
-        assert _materialize_object_storage_configs("/local/model") == "/local/model"
-        assert fake_object_storage.pulls == []
-
-    def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage):
-        uri = "s3://bucket/model"
-
-        first = _materialize_object_storage_configs(uri)
-        second = _materialize_object_storage_configs(uri)
-
-        assert first == second
-        assert len(fake_object_storage.pulls) == 1
-        pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0]
-        assert pulled_uri == uri
-        assert "*.json" in allow_pattern
-        assert ignore_pattern is None
-
-    def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage):
-        """A misleading bucket name must not beat the real config.json content."""
-        fake_object_storage.set_files(
-            [
-                (
-                    "config.json",
-                    '{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}',
-                )
-            ]
-        )
-        uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
-
-        model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False)
-
-        assert model == "qwen3_omni_moe"
-        assert fake_object_storage.pulls  # materialization happened
-
-    def test_name_match_fallback_scans_basename_only(self, fake_object_storage):
-        # Empty config.json (CosyVoice3 style) forces the name-match fallback.
-        fake_object_storage.set_files([("config.json", "{}")])
-
-        deceptive = "s3://qwen3-tts-models/plain-checkpoint"
-        assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None
-
-        matching = "s3://any-bucket/my-cosyvoice3-model"
-        assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"
@@ -424,10 +424,10 @@
        "baseline": {
          "H100": {
            "median_audio_ttfp_ms": [
-              82.6344
+              82.1886
            ],
            "median_audio_rtf": [
-              0.1277
+              0.1241
            ]
          }
        }
@@ -453,13 +453,13 @@
        "baseline": {
          "H100": {
            "median_audio_ttfp_ms": [
-              170.3737
+              160.9168
            ],
            "median_audio_rtf": [
-              0.2076
+              0.1907
            ],
            "audio_throughput": [
-              37.0834
+              39.2332
            ]
          }
        }
@@ -485,13 +485,13 @@
        "baseline": {
          "H100": {
            "median_audio_ttfp_ms": [
-              101.0293
+              100.1265
            ],
            "median_audio_rtf": [
-              0.1498
+              0.1437
            ],
            "audio_throughput": [
-              10.1186
+              11.7605
            ]
          }
        }
@@ -60,18 +60,9 @@ def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
    return torch.einsum("bhqk,bkhd->bqhd", probs, value)


-def _gpu_flash_attn_usable() -> bool:
+def _cuda_flash_attn_usable() -> bool:
    if not torch.cuda.is_available():
        return False
-    if torch.version.hip is not None:
-        for module in ("aiter", "flash_attn"):
-            try:
-                imported = __import__(module, fromlist=["flash_attn_varlen_func"])
-                if getattr(imported, "flash_attn_varlen_func", None) is not None:
-                    return True
-            except ImportError:
-                pass
-        return False
    try:
        spec = find_spec("vllm.vllm_flash_attn")
        if spec is None or spec.origin is None:
@@ -236,11 +227,12 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
    assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)


-@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
+@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
@@ -5,7 +5,6 @@ from contextlib import nullcontext

 import pytest
 import torch
-from torch.nn.attention import SDPBackend

 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
@@ -13,48 +12,6 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]


-@pytest.mark.parametrize(
-    ("query_sequence_length", "key_sequence_length", "expected_backend"),
-    [
-        (1, 1, SDPBackend.MATH),
-        (8, 1, SDPBackend.MATH),
-        (1, 8, SDPBackend.MATH),
-        (8, 8, SDPBackend.CUDNN_ATTENTION),
-    ],
-)
-def test_cudnn_routes_single_token_attention_to_math(
-    monkeypatch, query_sequence_length, key_sequence_length, expected_backend
-):
-    selected_backends = []
-
-    def fake_sdpa_kernel(backends):
-        selected_backends.append(backends)
-        return nullcontext()
-
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
-        fake_sdpa_kernel,
-    )
-    monkeypatch.setattr(
-        torch.nn.functional,
-        "scaled_dot_product_attention",
-        lambda query, _key, _value, **_kwargs: query,
-    )
-    impl = CuDNNAttentionImpl(
-        num_heads=2,
-        head_size=4,
-        softmax_scale=0.5,
-    )
-    query = torch.randn(1, query_sequence_length, 2, 4)
-    key = torch.randn(1, key_sequence_length, 2, 4)
-    value = torch.randn_like(key)
-
-    output = impl.forward_cuda(query, key, value)
-
-    assert output.shape == query.shape
-    assert selected_backends == [[expected_backend]]
-
-
 def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
    observed = {}

@@ -265,15 +265,10 @@ def test_fa_vs_sdpa():
    print(f"Max absolute difference (valid region): {max_diff:.6f}")
    print(f"Mean absolute difference (valid region): {mean_diff:.6f}")

-    # AITER's ROCm BF16 kernel is deterministic, but differs from ROCm SDPA by
-    # about 1.3% relative error. An absolute-only bound misclassifies that
-    # expected rounding difference for larger output values (for example,
-    # -2.65625 vs -2.6875). Keep the tighter historical bounds elsewhere.
-    if current_omni_platform.is_rocm():
-        torch.testing.assert_close(output_fa_valid, output_sdpa_valid, rtol=1e-2, atol=1.5e-2)
-    else:
-        assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
-        assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
+    # Assert that outputs are close
+    # Using higher tolerance for bfloat16 and different implementations
+    assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
+    assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"

    print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")

@@ -36,6 +36,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

 SEPARATE_CFG_TRANSFORMERS = [
    DreamIdOmniModel,
+    LTX2VideoTransformer3DModel,
    HeliosTransformer3DModel,
    LongCatImageTransformer2DModel,
    Cosmos3VFMTransformer,
@@ -123,13 +124,6 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
    assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True


-def test_ltx2_cache_dit_uses_one_forward_per_denoise_step():
-    """LTX batches all guidance passes into one Transformer invocation."""
-    adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config
-
-    assert adapter_config.has_separate_cfg is False
-
-
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
@patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
 def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):
@@ -53,10 +53,14 @@ def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
 def _extract_single_image(outputs) -> Image.Image:
    first_output = outputs[0]
    assert first_output.final_output_type == "image"
-    if not isinstance(first_output, OmniRequestOutput):
-        raise ValueError(f"Expected an OmniRequestOutput, got {type(first_output).__name__}")
+    if not hasattr(first_output, "request_output") or not first_output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")

-    images = first_output.images
+    req_out = first_output.request_output
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images' key")
+
+    images = req_out.images
    if images is None or len(images) != 1:
        raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
    return images[0]
@@ -416,8 +416,8 @@ def test_sp_correctness_advanced(model_name: str):


@pytest.mark.skipif(
-    not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
-    reason="Only tested on CUDA, ROCM and XPU",
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
 )
@pytest.mark.diffusion
@pytest.mark.parallel
@@ -447,8 +447,8 @@ def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None:


@pytest.mark.skipif(
-    not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
-    reason="Only tested on CUDA, ROCM and XPU",
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
 )
@pytest.mark.diffusion
@pytest.mark.parallel
@@ -468,8 +468,8 @@ def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> No


@pytest.mark.skipif(
-    not (current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_xpu()),
-    reason="Only tested on CUDA, ROCM and XPU",
+    not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()),
+    reason="Only tested on CUDA and XPU",
 )
@pytest.mark.diffusion
@pytest.mark.parallel
@@ -72,7 +72,8 @@ def _to_float_array(data: Any) -> np.ndarray:


 def _extract_output_array(output: Any) -> np.ndarray:
-    payload = output.images[0]
+    req_out = output.request_output
+    payload = req_out.images[0]
    return _to_float_array(payload)


@@ -24,20 +24,11 @@ from safetensors.torch import save_file

 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer, OmniServerParams
-from vllm_omni.platforms import current_omni_platform

 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

 MODEL = "Tongyi-MAI/Z-Image-Turbo"

-# ROCm's default FLASH_ATTN diffusion backend resolves to AITER. On a cold CI
-# worker, JIT-building its MHA kernel can take longer than the server readiness
-# timeout. This test validates LoRA switching rather than attention-backend or
-# compiled-execution performance, so use the eager SDPA path on ROCm only.
-ROCM_SERVER_ARGS = (
-    ["--enforce-eager", "--diffusion-attention-backend", "TORCH_SDPA"] if current_omni_platform.is_rocm() else []
-)
-

 PROMPT = "a photo of a cat sitting on a laptop keyboard"
 SIZE = "256x256"
@@ -51,14 +42,7 @@ server_params = [
            server_args=[
                "--num-gpus",
                "1",
-                *ROCM_SERVER_ARGS,
            ],
-            # A cold-cache Z-Image load can take slightly more than the
-            # orchestrator's 900s default on CI (weights alone have been
-            # observed at ~690s and full orchestrator readiness at ~906s).
-            # Match the large-model offline runner's startup allowance without
-            # weakening timeout detection for every online-serving test.
-            init_timeout=1800,
        ),
        id="zimage_turbo_lora",
    ),
@@ -16,11 +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.host_weight_plan import (
-    HostWeightPlan,
-    HostWeightPlanResult,
-    TensorBinding,
-)
 from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model

@@ -115,85 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check():
        loader.load_weights(model)


-def test_stream_online_quant_weights_offloads_layers_after_processing():
-    from vllm.model_executor.model_loader.reload.layerwise import (
-        get_layerwise_info,
-    )
-
-    events: list[str] = []
-
-    class _OnlineQuantMethod:
-        uses_meta_device = True
-
-    class _TrackedLayer(nn.Linear):
-        def __init__(self, name: str):
-            super().__init__(2, 2, bias=False)
-            self.name = name
-            self.quant_method = _OnlineQuantMethod()
-            get_layerwise_info(self).load_numel_total = self.weight.numel()
-
-        def to(self, *args, **kwargs):
-            events.append(self.name)
-            return super().to(*args, **kwargs)
-
-    class _StreamingModel(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.first = _TrackedLayer("first")
-            self.second = _TrackedLayer("second")
-
-    model = _StreamingModel()
-    weights = iter(
-        [
-            ("first.weight", torch.zeros((2, 2))),
-            ("second.weight", torch.zeros((2, 2))),
-        ]
-    )
-    streamed = DiffusersPipelineLoader._stream_online_quant_weights_to_cpu(model, weights)
-
-    assert next(streamed)[0] == "first.weight"
-    get_layerwise_info(model.first).reset()
-    assert next(streamed)[0] == "second.weight"
-    assert events == ["first"]
-
-    get_layerwise_info(model.second).reset()
-    with pytest.raises(StopIteration):
-        next(streamed)
-    assert events == ["first", "second"]
-
-
-def test_process_weights_skips_completed_online_quant_layer(monkeypatch):
-    from unittest.mock import Mock
-
-    from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
-    from vllm.model_executor.model_loader.reload import layerwise
-
-    class _TrackedLayer(nn.Linear):
-        def __init__(self):
-            super().__init__(2, 2, bias=False)
-            self.to_calls: list[object] = []
-
-        def to(self, *args, **kwargs):
-            self.to_calls.append(args[0] if args else kwargs.get("device"))
-            return super().to(*args, **kwargs)
-
-    model = nn.Module()
-    model.layer = _TrackedLayer()
-    quant_method = Mock(spec=QuantizeMethodBase)
-    quant_method.uses_meta_device = True
-    model.layer.quant_method = quant_method
-    model.layer._already_called_process_weights_after_loading = True
-    finalize = Mock()
-    monkeypatch.setattr(layerwise, "finalize_layerwise_processing", finalize)
-
-    loader = _make_loader_with_weights([])
-    loader._process_weights_after_loading(model, torch.device("cuda"))
-
-    finalize.assert_called_once_with(model, model_config=None)
-    quant_method.process_weights_after_loading.assert_not_called()
-    assert model.layer.to_calls == []
-
-
 class _ConfigAwareModel(nn.Module):
    def __init__(self, *, od_config):
        super().__init__()
@@ -263,137 +179,6 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
    assert get_current_diffusion_config_or_none() is None


-def test_dlo_transfers_loader_plan_and_skips_ordinary_weight_loading(monkeypatch):
-    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
-
-    od_config = SimpleNamespace(
-        dtype=torch.float32,
-        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
-        quantization_config=None,
-        enable_distributed_layerwise_offload=True,
-        dlo_use_allgather=False,
-        model="unused",
-    )
-    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
-    model = nn.Module()
-    model.transformer = nn.Linear(2, 2, bias=False)
-    plan = HostWeightPlan(
-        backing_kind="checkpoint_mmap",
-        bindings={},
-    )
-    loaded_ordinary_weights = False
-
-    def load_weights(_model):
-        nonlocal loaded_ordinary_weights
-        loaded_ordinary_weights = True
-
-    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
-    loader.load_weights = load_weights  # type: ignore[method-assign]
-    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
-    monkeypatch.setattr(
-        loader_mod,
-        "build_checkpoint_mmap_plan",
-        lambda *_args, **_kwargs: HostWeightPlanResult(plan),
-    )
-
-    assert loader.load_model(load_device="cpu") is model
-    assert not loaded_ordinary_weights
-    assert loader.take_host_weight_plan() is plan
-    assert loader.take_host_weight_plan() is None
-
-
-def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
-    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
-
-    od_config = SimpleNamespace(
-        dtype=torch.float32,
-        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
-        quantization_config=None,
-        enable_distributed_layerwise_offload=True,
-        dlo_use_allgather=False,
-        model="unused",
-    )
-    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
-
-    class MixedSourceModel(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.transformer = nn.Linear(2, 2, bias=False)
-            self.text_encoder = nn.Linear(2, 2, bias=False)
-            self.weights_sources = (
-                DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="transformer."),
-                DiffusersPipelineLoader.ComponentSource("unused", None, None, prefix="text_encoder."),
-            )
-            self.loaded_weight_names: list[str] = []
-
-        def load_weights(self, weights):
-            self.loaded_weight_names = [name for name, _ in weights]
-            return set(self.loaded_weight_names)
-
-    model = MixedSourceModel()
-    plan = HostWeightPlan(
-        backing_kind="checkpoint_mmap",
-        bindings={
-            "transformer.weight": TensorBinding(
-                checkpoint_key="weight",
-                file_path="unused",
-            )
-        },
-        planned_source_prefixes=frozenset({"transformer."}),
-    )
-    requested_prefixes: list[str] = []
-
-    def get_weights(source, model=None):
-        del model
-        requested_prefixes.append(source.prefix)
-        yield source.prefix + "weight", torch.ones(2, 2)
-
-    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
-    loader._get_weights_iterator = get_weights  # type: ignore[method-assign]
-    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
-    monkeypatch.setattr(
-        loader_mod,
-        "build_checkpoint_mmap_plan",
-        lambda *_args, **_kwargs: HostWeightPlanResult(plan),
-    )
-
-    assert loader.load_model(load_device="cpu") is model
-    assert requested_prefixes == ["text_encoder."]
-    assert model.loaded_weight_names == ["text_encoder.weight"]
-    assert loader.take_host_weight_plan() is plan
-
-
-def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
-    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
-
-    od_config = SimpleNamespace(
-        dtype=torch.float32,
-        parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1),
-        quantization_config=None,
-        enable_distributed_layerwise_offload=True,
-        dlo_use_allgather=False,
-        model="unused",
-    )
-    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
-    model = nn.Module()
-    model.transformer = nn.Linear(2, 2, bias=False)
-    calls: list[str] = []
-
-    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
-    loader.load_weights = lambda _model: calls.append("load")  # type: ignore[method-assign]
-    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
-    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
-    monkeypatch.setattr(
-        loader_mod,
-        "build_checkpoint_mmap_plan",
-        lambda *_args, **_kwargs: HostWeightPlanResult(None, "not direct-compatible"),
-    )
-
-    assert loader.load_model(load_device="cpu") is model
-    assert calls == ["load", "process"]
-    assert loader.take_host_weight_plan() 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
@@ -1,72 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import contextlib
-
-import huggingface_hub
-import pytest
-
-from vllm_omni.diffusion.model_loader import hub_prefetch
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-def test_prefetch_subfolders_propagates_revision(monkeypatch):
-    calls = []
-
-    def fake_snapshot_download(**kwargs):
-        calls.append(kwargs)
-
-    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
-    monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
-
-    hub_prefetch.prefetch_subfolders(
-        "org/model",
-        ["transformer"],
-        revision="pinned-revision",
-    )
-
-    assert calls == [
-        {
-            "repo_id": "org/model",
-            "revision": "pinned-revision",
-            "allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"],
-        }
-    ]
-
-
-def test_from_pretrained_retry_preserves_revision(monkeypatch):
-    factory_calls = []
-    prefetch_calls = []
-
-    def fake_factory(model, **kwargs):
-        factory_calls.append((model, kwargs))
-        if len(factory_calls) == 1:
-            raise OSError("partial cache")
-        return "loaded"
-
-    def fake_prefetch(model, subfolders, **kwargs):
-        prefetch_calls.append((model, tuple(subfolders), kwargs))
-
-    monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch)
-    monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None)
-
-    result = hub_prefetch.from_pretrained_with_prefetch(
-        fake_factory,
-        "org/model",
-        subfolder="transformer",
-        prefetch_list=["transformer", "vae"],
-        revision="pinned-revision",
-        max_attempts=2,
-    )
-
-    assert result == "loaded"
-    assert len(factory_calls) == 2
-    assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls)
-    assert prefetch_calls == [
-        (
-            "org/model",
-            ("transformer", "vae"),
-            {"local_files_only": False, "revision": "pinned-revision"},
-        )
-    ]
@@ -7,17 +7,17 @@ from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
    Flux2Pipeline,
    _resolve_component_quant_config,
 )
-from vllm_omni.quantization import ComponentQuantizationConfig

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


 def test_resolve_component_quant_config_routes_components():
    text_encoder_config = object()
-    component_config = ComponentQuantizationConfig(
-        component_configs={"text_encoder": text_encoder_config},
-        default_config=None,
-    )
+    component_config = Mock()
+    component_config.resolve.side_effect = {
+        "text_encoder": text_encoder_config,
+        "transformer": None,
+    }.get

    assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
    assert _resolve_component_quant_config(component_config, "transformer") is None
@@ -1,881 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Focused LTX-2.5 pipeline correctness tests."""
-
-import json
-from dataclasses import replace
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.models.ltx2 import ltx2_components, ltx2_latents, ltx2_runtime
-from vllm_omni.diffusion.models.ltx2.ltx2_components import (
-    LTX25_DISTILLED_COMPONENT_PROFILE,
-    LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
-    LTX25_FULL_COMPONENT_PROFILE,
-    LTX25_TWO_STAGE_COMPONENT_PROFILE,
-    detect_ltx_model_version,
-    resolve_ltx_checkpoint_kind,
-    resolve_ltx_component_profile,
-)
-from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
-    LTXPhaseResult,
-    _first_frame_keyframes_mask,
-    _official_ltx_sigmas,
-)
-from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
-    LTX_GUIDANCE_EXECUTOR,
-    LTXGuidanceSpec,
-    LTXModalityGuidance,
-    velocity_from_x0,
-    x0_from_velocity,
-)
-from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
-from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
-    LTX2_DISTILLED_TWO_STAGE_RECIPE,
-    LTX2_ONE_STAGE_RECIPE,
-    LTX23_ONE_STAGE_RECIPE,
-    LTX25_DISTILLED_ONE_STAGE_RECIPE,
-    LTX25_DISTILLED_TWO_STAGE_RECIPE,
-    LTX25_FULL_RECIPE,
-    LTX25_TWO_STAGE_RECIPE,
-    LTX_DISTILLED_SIGMAS,
-    LTX_POSITIVE_ONLY_RECIPE,
-    resolve_ltx_pipeline_recipe,
-)
-from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
-from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
-    LTX2DistilledOneStagePipeline,
-    LTX2Pipeline,
-)
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
-    LTX2DistilledTwoStagePipeline,
-    LTX2TwoStagePipeline,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-@pytest.mark.parametrize("model_version", ["2", "2.3", "2.5"])
-def test_ltx_timestep_adaln_expands_per_token_like_official_without_sp(model_version):
-    pipe = object.__new__(LTX2Pipeline)
-    pipe.model_version = model_version
-    pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
-    ts = torch.tensor([750.0, 500.0])
-
-    kwargs = LTXRuntime._denoise_timestep_kwargs(
-        pipe,
-        ts,
-        SimpleNamespace(),
-        SimpleNamespace(),
-        video_token_count=4,
-        audio_token_count=3,
-    )
-
-    torch.testing.assert_close(kwargs["timestep"], ts[:, None].expand(-1, 4))
-    torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None].expand(-1, 3))
-    assert kwargs["sigma"] is ts
-    assert kwargs["audio_sigma"] is ts
-
-
-def test_ltx_ancestral_step_clears_sp_audio_padding_after_scheduler_update(monkeypatch):
-    state = LTXAVState(
-        video=torch.zeros(1, 2, 2),
-        audio=torch.zeros(1, 4, 2),
-    )
-    forward_ctx = SimpleNamespace(
-        original_audio_num_frames=3,
-        sampler="euler_ancestral",
-    )
-    denoise_ctx = SimpleNamespace(latents=None, audio_latents=None)
-    pipe = SimpleNamespace(
-        _predict_noise_for_step=lambda *_args: (torch.zeros_like(state.video), torch.zeros_like(state.audio))
-    )
-
-    def fake_step(_pipeline, actual_forward_ctx, _denoise_ctx, *_args):
-        assert actual_forward_ctx.sampler == "euler_ancestral"
-        updated_audio = torch.ones_like(state.audio)
-        updated_audio[:, 3:] = 99.0
-        return torch.ones_like(state.video), updated_audio
-
-    monkeypatch.setattr(ltx2_runtime, "step_denoised_latents", fake_step)
-
-    actual = LTXRuntime._denoise_step(
-        pipe,
-        0,
-        torch.tensor(1.0),
-        state,
-        forward_ctx,
-        denoise_ctx,
-    )
-
-    torch.testing.assert_close(actual.audio[:, :3], torch.ones_like(actual.audio[:, :3]))
-    torch.testing.assert_close(actual.audio[:, 3:], torch.zeros_like(actual.audio[:, 3:]))
-
-
-def test_ltx25_first_causal_latent_frame_is_marked_for_keyframe_embedding():
-    reference = torch.zeros(2, 12, 8)
-
-    mask = _first_frame_keyframes_mask(reference, latent_num_frames=3)
-
-    assert mask.shape == (2, 12, 1)
-    torch.testing.assert_close(mask[:, :4], torch.ones(2, 4, 1))
-    torch.testing.assert_close(mask[:, 4:], torch.zeros(2, 8, 1))
-
-
-def _make_ltx_request_pipe(cls):
-    pipe = object.__new__(cls)
-    torch.nn.Module.__init__(pipe)
-    pipe.device = torch.device("cpu")
-    pipe.tokenizer_max_length = 99
-    pipe.vae_spatial_compression_ratio = 32
-    pipe.vae_temporal_compression_ratio = 8
-    return pipe
-
-
-def _resolve_request_inputs_for_test(
-    pipe,
-    req,
-    *,
-    guidance_scale=4.0,
-    negative_prompt=None,
-):
-    return pipe._resolve_request_inputs(
-        req,
-        prompt=None,
-        negative_prompt=negative_prompt,
-        height=None,
-        width=None,
-        num_frames=None,
-        frame_rate=None,
-        num_inference_steps=None,
-        guidance_scale=guidance_scale,
-        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,
-    )
-
-
-def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
-    pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
-    od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
-
-    monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
-    monkeypatch.setattr(ltx2_components, "prefetch_subfolders", lambda *_args, **_kwargs: None)
-    monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", lambda *_args, **_kwargs: object())
-
-    with pytest.raises(ImportError) as exc_info:
-        ltx2_components.initialize_pipeline_components(pipe, od_config)
-
-    assert str(exc_info.value) == (
-        "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15."
-    )
-
-
-def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
-    revision = "pinned-revision"
-    calls = {"components": []}
-    profile = replace(
-        LTX25_DISTILLED_COMPONENT_PROFILE,
-        text_encoder_cls=object,
-        video_vae_cls=object,
-        vocoder_cls=object,
-        vocoder_fallback_cls=None,
-        scheduler_use_dynamic_shifting=False,
-    )
-    pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
-    od_config = SimpleNamespace(
-        model="org/converted-ltx25",
-        revision=revision,
-        dtype=torch.bfloat16,
-        quantization_config=None,
-    )
-
-    monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
-    monkeypatch.setattr(ltx2_components, "_install_connector_attention", lambda _connectors, **_kwargs: None)
-    monkeypatch.setattr(ltx2_components, "_place_aux_components", lambda _pipeline: None)
-
-    def fake_prefetch(model, subfolders, **kwargs):
-        calls["prefetch"] = (model, tuple(subfolders), kwargs)
-
-    def fake_tokenizer(*args, **kwargs):
-        calls["tokenizer"] = (args, kwargs)
-        return SimpleNamespace(model_max_length=1024)
-
-    def fake_component(component_cls, model, subfolder, **kwargs):
-        calls["components"].append((component_cls, model, subfolder, kwargs))
-        if subfolder == "vae":
-            return SimpleNamespace(spatial_compression_ratio=32, temporal_compression_ratio=8)
-        if subfolder == "audio_vae":
-            return SimpleNamespace(
-                mel_compression_ratio=4,
-                temporal_compression_ratio=4,
-                config=SimpleNamespace(sample_rate=16_000, mel_hop_length=160),
-            )
-        return object()
-
-    def fake_transformer_config(model, subfolder, local_files_only, *, revision):
-        calls["transformer_config"] = (model, subfolder, local_files_only, revision)
-        return {"component": "transformer"}
-
-    monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
-    monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
-    monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
-    monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
-    monkeypatch.setattr(
-        ltx2_components,
-        "create_transformer_from_config",
-        lambda *_args, **_kwargs: SimpleNamespace(config=SimpleNamespace(patch_size=1, patch_size_t=1)),
-    )
-
-    def fake_scheduler(*args, **kwargs):
-        calls["scheduler"] = (args, kwargs)
-        return SimpleNamespace(config={"use_dynamic_shifting": False, "shift_terminal": None})
-
-    monkeypatch.setattr(ltx2_components.FlowMatchEulerDiscreteScheduler, "from_pretrained", fake_scheduler)
-
-    ltx2_components.initialize_pipeline_components(pipeline, od_config)
-
-    assert pipeline.weights_sources[0].revision == revision
-    assert calls["prefetch"][2]["revision"] == revision
-    assert calls["tokenizer"][1]["revision"] == revision
-    assert {call[2] for call in calls["components"]} == {
-        "text_encoder",
-        "connectors",
-        "vae",
-        "audio_vae",
-        "vocoder",
-        "latent_upsampler",
-    }
-    assert all(call[3]["revision"] == revision for call in calls["components"])
-    assert calls["transformer_config"] == (
-        od_config.model,
-        profile.transformer_subfolder,
-        False,
-        revision,
-    )
-    assert calls["scheduler"][1]["revision"] == revision
-
-
-def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
-    (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
-    (tmp_path / "transformer").mkdir()
-    (tmp_path / "transformer" / "config.json").write_text(json.dumps({"ff_bias": False}))
-
-    assert detect_ltx_model_version(str(tmp_path)) == "2.3"
-
-
-def test_ltx25_checkpoint_selects_full_one_stage_profile(tmp_path, monkeypatch):
-    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
-
-    (tmp_path / "model_index.json").write_text(
-        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
-    )
-
-    def stub_components(pipe, od_config):
-        pipe.od_config = od_config
-        pipe.vae_spatial_compression_ratio = 32
-
-    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
-    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
-
-    pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
-
-    assert pipe.model_version == "2.5"
-    assert pipe.component_profile is LTX25_FULL_COMPONENT_PROFILE
-    assert pipe.pipeline_recipe is LTX25_FULL_RECIPE
-    assert pipe.preserve_sp_padded_audio_duration
-    assert pipe.reports_stage_durations
-
-
-def test_ltx25_full_recipe_matches_official_sft_defaults():
-    (phase,) = LTX25_FULL_RECIPE.phases
-
-    assert (LTX25_FULL_RECIPE.height, LTX25_FULL_RECIPE.width) == (544, 960)
-    assert LTX25_FULL_RECIPE.num_frames == 121
-    assert LTX25_FULL_RECIPE.frame_rate == 24.0
-    assert LTX25_FULL_RECIPE.num_inference_steps == 30
-    assert phase.guidance.video == LTXModalityGuidance(
-        cfg_scale=3.0,
-        stg_scale=1.0,
-        modality_scale=3.0,
-        rescale_scale=0.7,
-        stg_blocks=(28,),
-    )
-    assert phase.guidance.audio == LTXModalityGuidance(
-        cfg_scale=7.0,
-        stg_scale=1.0,
-        modality_scale=3.0,
-        rescale_scale=0.7,
-        stg_blocks=(28,),
-    )
-    assert phase.sigmas is None
-    assert phase.noise_scale == 1.0
-    assert phase.use_official_sigma_schedule
-    assert LTX25_FULL_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
-    assert LTX25_FULL_COMPONENT_PROFILE.scheduler_use_dynamic_shifting
-    assert LTX25_FULL_COMPONENT_PROFILE.scheduler_shift_terminal == 0.1
-
-
-def test_ltx25_full_sigma_schedule_uses_official_default_sequence_anchor():
-    scheduler = SimpleNamespace(
-        config={
-            "base_image_seq_len": 1024,
-            "max_image_seq_len": 4096,
-            "base_shift": 0.95,
-            "max_shift": 2.05,
-            "shift_terminal": 0.1,
-        }
-    )
-
-    sigmas = _official_ltx_sigmas(
-        scheduler,
-        steps=30,
-        device=torch.device("cpu"),
-    )
-
-    torch.testing.assert_close(
-        sigmas[[0, 1, -2, -1]],
-        torch.tensor([1.0, 0.9949570, 0.1, 0.0]),
-        rtol=1e-5,
-        atol=1e-6,
-    )
-
-
-def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypatch):
-    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
-
-    (tmp_path / "model_index.json").write_text(
-        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
-    )
-
-    def stub_components(pipe, od_config):
-        pipe.od_config = od_config
-        pipe.vae_spatial_compression_ratio = 32
-
-    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
-    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
-
-    pipe = LTX2DistilledTwoStagePipeline(
-        od_config=SimpleNamespace(
-            model=str(tmp_path),
-            enable_diffusion_pipeline_profiler=False,
-        )
-    )
-
-    assert pipe.component_profile is LTX25_DISTILLED_COMPONENT_PROFILE
-    assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
-
-
-def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-    (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "LTX2Pipeline", "model_version": "2.5"}))
-    for subfolder, marker in (("transformer", "distilled"), ("transformer_full", "full")):
-        path = tmp_path / subfolder
-        path.mkdir()
-        (path / "config.json").write_text(json.dumps({"marker": marker}))
-
-    config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="LTX2Pipeline")
-    config.enrich_config()
-
-    assert config.tf_model_config.get("marker") == "full"
-
-
-def test_ltx25_four_public_pipeline_semantics_are_disjoint():
-    assert resolve_ltx_component_profile("one_stage", "2.5") is LTX25_FULL_COMPONENT_PROFILE
-    assert resolve_ltx_pipeline_recipe("one_stage", "2.5") is LTX25_FULL_RECIPE
-    assert resolve_ltx_component_profile("two_stage", "2.5") is LTX25_TWO_STAGE_COMPONENT_PROFILE
-    assert resolve_ltx_pipeline_recipe("two_stage", "2.5") is LTX25_TWO_STAGE_RECIPE
-    assert resolve_ltx_component_profile("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
-    assert resolve_ltx_pipeline_recipe("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_RECIPE
-    assert resolve_ltx_component_profile("distilled_two_stage", "2.5") is LTX25_DISTILLED_COMPONENT_PROFILE
-    assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.5") is LTX25_DISTILLED_TWO_STAGE_RECIPE
-
-    assert LTX2Pipeline.pipeline_kind == "one_stage"
-    assert LTX2TwoStagePipeline.pipeline_kind == "two_stage"
-    assert LTX2DistilledOneStagePipeline.pipeline_kind == "distilled_one_stage"
-    assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
-
-
-@pytest.mark.parametrize(
-    ("pipeline_kind", "expected"),
-    [
-        ("one_stage", "regular"),
-        ("two_stage", "regular"),
-        ("distilled_one_stage", "distilled"),
-        ("distilled_two_stage", "distilled"),
-        ("dmd2", None),
-    ],
-)
-def test_ltx_checkpoint_kind_is_derived_from_pipeline_kind(pipeline_kind, expected):
-    assert resolve_ltx_checkpoint_kind(pipeline_kind) == expected
-
-
-def test_ltx_checkpoint_kind_rejects_unknown_pipeline():
-    with pytest.raises(ValueError, match="Unsupported LTX pipeline kind"):
-        resolve_ltx_checkpoint_kind("unknown")
-
-
-def test_ltx25_distilled_two_stage_recipe_matches_model_card_resolution():
-    stage1, stage2 = LTX25_DISTILLED_TWO_STAGE_RECIPE.phases
-
-    assert (LTX25_DISTILLED_TWO_STAGE_RECIPE.height, LTX25_DISTILLED_TWO_STAGE_RECIPE.width) == (1088, 1920)
-    assert stage1.spatial_downscale == 2
-    assert stage1.sigmas == LTX_DISTILLED_SIGMAS
-    assert stage1.sampler == "euler_ancestral"
-    assert stage2.sampler == "euler"
-    assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
-    assert stage2.input_transform == "spatial_upsample"
-
-    assert LTX25_DISTILLED_TWO_STAGE_RECIPE.negative_prompt == ""
-    assert not LTX25_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt
-
-
-def test_ltx25_distilled_one_stage_uses_native_low_resolution():
-    (phase,) = LTX25_DISTILLED_ONE_STAGE_RECIPE.phases
-
-    assert (LTX25_DISTILLED_ONE_STAGE_RECIPE.height, LTX25_DISTILLED_ONE_STAGE_RECIPE.width) == (544, 960)
-    assert phase.spatial_downscale == 1
-    assert phase.sigmas == LTX_DISTILLED_SIGMAS
-    assert phase.sampler == "euler_ancestral"
-    assert not LTX25_DISTILLED_ONE_STAGE_RECIPE.supports_cache_dit
-
-
-def test_ltx25_regular_two_stage_uses_full_transformer_then_distilled_lora():
-    stage1, stage2 = LTX25_TWO_STAGE_RECIPE.phases
-
-    assert (LTX25_TWO_STAGE_RECIPE.height, LTX25_TWO_STAGE_RECIPE.width) == (1088, 1920)
-    assert LTX25_TWO_STAGE_RECIPE.num_inference_steps == 30
-    assert stage1.spatial_downscale == 2
-    assert stage1.guidance == LTX25_FULL_RECIPE.request_guidance
-    assert stage2.adapter_slot == "ltx_distilled"
-    assert LTX25_TWO_STAGE_RECIPE.allow_request_phase_sigmas
-    assert LTX25_TWO_STAGE_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
-    assert LTX25_TWO_STAGE_COMPONENT_PROFILE.distilled_lora_filename is not None
-
-
-@pytest.mark.parametrize(
-    "recipe",
-    [
-        LTX2_ONE_STAGE_RECIPE,
-        LTX23_ONE_STAGE_RECIPE,
-        LTX_POSITIVE_ONLY_RECIPE,
-    ],
-)
-def test_ltx_one_stage_recipes_declare_cache_dit_supported(recipe):
-    assert recipe.supports_cache_dit
-
-
-@pytest.mark.parametrize(
-    "recipe",
-    [
-        LTX25_FULL_RECIPE,
-        LTX25_DISTILLED_ONE_STAGE_RECIPE,
-        LTX2_DISTILLED_TWO_STAGE_RECIPE,
-        LTX25_DISTILLED_TWO_STAGE_RECIPE,
-    ],
-)
-def test_ltx25_and_multistage_recipes_declare_cache_dit_unsupported(recipe):
-    assert not recipe.supports_cache_dit
-
-
-def test_ltx_positive_only_guidance_preserves_official_x0_roundtrip():
-    sample = torch.full((1, 1), 1e8, dtype=torch.float32)
-    velocity = torch.ones((1, 1), dtype=torch.bfloat16)
-    sigma = torch.tensor(0.5)
-
-    actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
-        sample,
-        {"cond": velocity},
-        sigma,
-        LTXModalityGuidance(),
-    )
-    expected = velocity_from_x0(sample, x0_from_velocity(sample, velocity, sigma), sigma)
-
-    assert actual.dtype == torch.float32
-    torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
-
-
-def test_ltx_velocity_from_x0_materializes_official_scalar_sigma():
-    item_calls = []
-
-    class TrackingTensor(torch.Tensor):
-        @staticmethod
-        def __new__(cls):
-            return torch.Tensor._make_subclass(cls, torch.tensor(0.725), False)
-
-        def item(self):
-            item_calls.append(True)
-            return super().item()
-
-    sample = torch.tensor([[1.0]], dtype=torch.bfloat16)
-    x0 = torch.tensor([[0.5]], dtype=torch.bfloat16)
-
-    velocity_from_x0(sample, x0, TrackingTensor())
-
-    assert item_calls == [True]
-
-
-def test_ltx_ancestral_positive_only_guidance_preserves_raw_velocity():
-    sample = torch.full((1, 1), 1e8, dtype=torch.float32)
-    velocity = torch.ones((1, 1), dtype=torch.bfloat16)
-
-    actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
-        sample,
-        {"cond": velocity},
-        torch.tensor(0.5),
-        LTXModalityGuidance(),
-        preserve_positive_velocity=True,
-    )
-
-    assert actual is velocity
-
-
-@pytest.mark.parametrize("model_version", ["2.5", "2.3"])
-def test_ltx_t2v_denoise_state_uses_official_bfloat16_dtype(model_version):
-    expected_dtype = torch.bfloat16
-    captured = {}
-
-    def prepare_latents(**kwargs):
-        captured["video_dtype"] = kwargs["dtype"]
-        return torch.empty(1, 1, 1, dtype=kwargs["dtype"])
-
-    def prepare_audio_latents(*_args, **kwargs):
-        captured["audio_dtype"] = kwargs["dtype"]
-        return torch.empty(1, 1, 1, dtype=kwargs["dtype"]), 1, 1
-
-    pipeline = SimpleNamespace(
-        model_version=model_version,
-        transformer=SimpleNamespace(config=SimpleNamespace(in_channels=128)),
-        prepare_latents=prepare_latents,
-        prepare_audio_latents=prepare_audio_latents,
-        audio_sampling_rate=48_000,
-        audio_hop_length=160,
-        audio_vae_temporal_compression_ratio=4,
-        audio_vae_mel_compression_ratio=4,
-        audio_vae=SimpleNamespace(config=SimpleNamespace(mel_bins=64, latent_channels=8)),
-        _resolve_audio_latent_length=lambda length, _latents: length,
-    )
-    request = SimpleNamespace(
-        num_videos_per_prompt=1,
-        height=64,
-        width=64,
-        num_frames=1,
-        frame_rate=24.0,
-        generator=None,
-        latents=None,
-        audio_latents=None,
-        audio_latents_normalized=False,
-    )
-    prompt = SimpleNamespace(
-        batch_size=1,
-        positive_connector_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
-        positive_connector_audio_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
-    )
-
-    LTXRuntime._prepare_video_latents_stage(
-        pipeline,
-        request,
-        prompt,
-        device=torch.device("cpu"),
-        noise_scale=1.0,
-    )
-    LTXRuntime._prepare_audio_latents_stage(
-        pipeline,
-        request,
-        prompt,
-        device=torch.device("cpu"),
-        noise_scale=1.0,
-    )
-
-    assert captured == {"video_dtype": expected_dtype, "audio_dtype": expected_dtype}
-
-
-def test_ltx_seeded_latents_match_official_packed_rng_layout():
-    pipeline = SimpleNamespace(
-        model_version="2.5",
-        vae_spatial_compression_ratio=2,
-        vae_temporal_compression_ratio=2,
-        transformer_spatial_patch_size=1,
-        transformer_temporal_patch_size=1,
-        audio_vae_mel_compression_ratio=2,
-        od_config=SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)),
-    )
-    device = torch.device("cpu")
-    actual_generator = torch.Generator(device=device).manual_seed(123)
-    expected_generator = torch.Generator(device=device).manual_seed(123)
-
-    actual_video = ltx2_latents.prepare_video_latents(
-        pipeline,
-        batch_size=1,
-        num_channels_latents=2,
-        height=4,
-        width=6,
-        num_frames=3,
-        dtype=torch.float32,
-        device=device,
-        generator=actual_generator,
-    )
-    actual_audio, original_length, padded_length = ltx2_latents.prepare_audio_latents(
-        pipeline,
-        batch_size=1,
-        num_channels_latents=2,
-        audio_latent_length=3,
-        num_mel_bins=4,
-        dtype=torch.float32,
-        device=device,
-        generator=actual_generator,
-    )
-
-    expected_video = torch.randn((1, 12, 2), generator=expected_generator)
-    expected_audio = torch.randn((1, 3, 4), generator=expected_generator)
-
-    torch.testing.assert_close(actual_video, expected_video, rtol=0.0, atol=0.0)
-    torch.testing.assert_close(actual_audio, expected_audio, rtol=0.0, atol=0.0)
-    assert (original_length, padded_length) == (3, 3)
-
-
-def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
-    request_inputs = LTXRequestInputs(
-        prompt="prompt",
-        negative_prompt="",
-        height=64,
-        width=64,
-        num_frames=1,
-        frame_rate=24.0,
-        num_inference_steps=8,
-        guidance=LTXGuidanceSpec.positive_only(),
-        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.25,
-        decode_noise_scale=0.5,
-        output_type="np",
-        max_sequence_length=16,
-    )
-    prompt_context = object()
-    phase_calls = []
-
-    def resolve_request_inputs(req, **kwargs):
-        return request_inputs
-
-    def run_phase(req, inputs, *, prompt_context=None, **kwargs):
-        phase_recipe = kwargs["phase_recipe"]
-        assert kwargs["image"] is source_image
-        phase_calls.append((phase_recipe, inputs, prompt_context))
-        if len(phase_calls) == 1:
-            assert prompt_context is None
-            assert inputs.decode_timestep == 0.25
-            assert inputs.decode_noise_scale == 0.5
-            context = prompt_context_sentinel
-            assert phase_recipe.name == "generate_lowres"
-            assert (inputs.height, inputs.width) == (32, 32)
-            assert inputs.num_inference_steps == 2
-            assert kwargs["sigmas"] == stage_1_sigmas
-            assert kwargs["noise_scale"] == 1.0
-            video = torch.ones(1, 128, 1, 1, 1)
-            audio = torch.full((1, 8, 1, 2), 2.0)
-            audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0)
-        else:
-            assert prompt_context is prompt_context_sentinel
-            assert phase_recipe.name == "refine"
-            assert (inputs.height, inputs.width) == (64, 64)
-            assert inputs.latents.shape == (1, 128, 1, 2, 2)
-            torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0))
-            assert inputs.audio_latents_normalized
-            assert inputs.guidance_scale == 1.0
-            assert inputs.num_inference_steps == 2
-            assert kwargs["sigmas"] == stage_2_sigmas
-            assert kwargs["noise_scale"] == stage_2_sigmas[0]
-            assert inputs.decode_timestep == 0.0
-            assert inputs.decode_noise_scale is None
-            context = prompt_context
-            video = torch.full((1, 128, 1, 2, 2), 3.0)
-            audio = torch.full((1, 8, 1, 2), 4.0)
-            audio_for_next_phase = None
-        return LTXPhaseResult(
-            forward_context=SimpleNamespace(prompt_context=context),
-            video=video,
-            audio=audio,
-            audio_for_next_phase=audio_for_next_phase,
-        )
-
-    def decode_phase(phase):
-        return DiffusionOutput(output=(phase.video, phase.audio))
-
-    class FakeUpsampler(torch.nn.Module):
-        dtype = torch.float32
-
-        def forward(self, latents):
-            return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
-
-    prompt_context_sentinel = prompt_context
-    stage_1_sigmas = [1.0, 0.5, 0.0]
-    stage_2_sigmas = [0.8, 0.2, 0.0]
-    pipeline = object.__new__(LTX2TwoStagePipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.pipeline_recipe = LTX25_DISTILLED_TWO_STAGE_RECIPE
-    pipeline.device = torch.device("cpu")
-    pipeline.vae_spatial_compression_ratio = 32
-    pipeline.vae_temporal_compression_ratio = 8
-    pipeline.latent_upsampler = FakeUpsampler()
-    object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
-    object.__setattr__(pipeline, "run_phase", run_phase)
-    object.__setattr__(pipeline, "decode_phase", decode_phase)
-
-    source_image = object()
-    output = pipeline.forward(
-        SimpleNamespace(sampling_params_list=[], prompts=[]),
-        image=source_image,
-        stage_1_sigmas=stage_1_sigmas,
-        stage_2_sigmas=stage_2_sigmas,
-    )
-
-    assert len(phase_calls) == 2
-    assert phase_calls[1][2] is prompt_context_sentinel
-    torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
-    torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
-
-
-class TestLTXRequestParsing:
-    def test_request_resolves_custom_sigmas_from_extra_args(self):
-        from vllm_omni.diffusion.request import OmniDiffusionRequest
-        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-        sigmas = [1.0, 0.5, 0.0]
-        req = DiffusionRequestBatch(
-            [
-                OmniDiffusionRequest(
-                    prompt="prompt",
-                    sampling_params=OmniDiffusionSamplingParams(extra_args={"sigmas": sigmas}),
-                    request_id="ltx-custom-sigmas",
-                )
-            ]
-        )
-
-        assert _make_ltx_request_pipe(LTX2Pipeline)._resolve_request_sigmas(req, None) == sigmas
-
-    def test_request_resolves_per_stage_sigmas_from_extra_args(self):
-        from vllm_omni.diffusion.request import OmniDiffusionRequest
-        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-        stage_1_sigmas = [1.0, 0.4, 0.0]
-        req = DiffusionRequestBatch(
-            [
-                OmniDiffusionRequest(
-                    prompt="prompt",
-                    sampling_params=OmniDiffusionSamplingParams(extra_args={"stage_1_sigmas": stage_1_sigmas}),
-                    request_id="ltx-custom-phase-sigmas",
-                )
-            ]
-        )
-
-        resolved = _make_ltx_request_pipe(LTX2TwoStagePipeline)._resolve_request_phase_sigmas(
-            req, [1.0, 0.0], [0.8, 0.0]
-        )
-
-        assert resolved == (stage_1_sigmas, [0.8, 0.0])
-
-    @pytest.mark.parametrize("value", [-1, 52, 1.5, True])
-    def test_request_rejects_invalid_image_crf(self, value):
-        from vllm_omni.diffusion.request import OmniDiffusionRequest
-        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-        req = DiffusionRequestBatch(
-            [
-                OmniDiffusionRequest(
-                    prompt="prompt",
-                    sampling_params=OmniDiffusionSamplingParams(extra_args={"image_crf": value}),
-                    request_id="ltx-invalid-image-crf",
-                )
-            ]
-        )
-
-        with pytest.raises(ValueError, match="image_crf"):
-            _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req)
-
-
-class TestLTXForwardStages:
-    def test_gemma_tokenization_matches_official_bos_and_left_padding(self):
-        pipe = _make_ltx_request_pipe(LTX2Pipeline)
-
-        class FakeTokenizer:
-            bos_token_id = 2
-            eos_token_id = 1
-            eos_token = "<eos>"
-            pad_token_id = 0
-            pad_token = "<pad>"
-            padding_side = "right"
-
-            def __init__(self):
-                self.calls = []
-
-            def __call__(self, text, **kwargs):
-                self.calls.append((text, kwargs))
-                ids = [7, 8, 9, self.eos_token_id] if text == "long" else [6, self.eos_token_id]
-                return SimpleNamespace(input_ids=ids)
-
-            def pad(self, encoded, *, padding, max_length, return_tensors, return_attention_mask):
-                assert padding == "max_length"
-                assert return_tensors == "pt"
-                assert return_attention_mask is True
-                rows = []
-                masks = []
-                for ids in encoded["input_ids"]:
-                    pad = max_length - len(ids)
-                    rows.append([self.pad_token_id] * pad + ids)
-                    masks.append([0] * pad + [1] * len(ids))
-                return SimpleNamespace(input_ids=torch.tensor(rows), attention_mask=torch.tensor(masks))
-
-        class FakeTextEncoder(torch.nn.Module):
-            dtype = torch.float32
-
-            def __init__(self):
-                super().__init__()
-                self.seen_input_ids = None
-                self.seen_attention_mask = None
-
-            def forward(self, *, input_ids, attention_mask, output_hidden_states):
-                assert output_hidden_states is True
-                self.seen_input_ids = input_ids.clone()
-                self.seen_attention_mask = attention_mask.clone()
-                return SimpleNamespace(hidden_states=(input_ids.unsqueeze(-1).float(),))
-
-        tokenizer = FakeTokenizer()
-        text_encoder = FakeTextEncoder()
-        object.__setattr__(pipe, "tokenizer", tokenizer)
-        object.__setattr__(pipe, "text_encoder", text_encoder)
-
-        pipe._get_gemma_prompt_embeds(
-            [" long ", "short"],
-            max_sequence_length=4,
-            device=torch.device("cpu"),
-            dtype=torch.float32,
-        )
-
-        torch.testing.assert_close(text_encoder.seen_input_ids, torch.tensor([[2, 7, 8, 9], [0, 2, 6, 1]]))
-        torch.testing.assert_close(text_encoder.seen_attention_mask, torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]]))
-        assert tokenizer.padding_side == "left"
-        assert [call[0] for call in tokenizer.calls] == ["long", "short"]
@@ -3,7 +3,6 @@

 """Unit tests for LTX image-to-video input and conditioning behavior."""

-import sys
 from types import SimpleNamespace

 import pytest
@@ -24,7 +23,7 @@ def _make_ltx23_request_pipe(cls):


 class TestLTXImageToVideoForwardStages:
-    def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self):
+    def test_i2v_pil_preprocessing_matches_official_center_crop(self):
        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images

        pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
@@ -36,136 +35,6 @@ class TestLTXImageToVideoForwardStages:

        torch.testing.assert_close(actual, expected)

-    def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch):
-        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
-
-        pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
-        pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
-        image = Image.fromarray(pixels.numpy())
-        monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array)
-
-        actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18)
-        expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
-
-        torch.testing.assert_close(actual, expected)
-
-    def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch):
-        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
-
-        captured_crfs = []
-        captured_prepare_kwargs = []
-
-        def fake_preprocess(image, *, height, width, crf, device, dtype):
-            del image, device
-            captured_crfs.append(crf)
-            return torch.zeros(1, 3, height, width, dtype=dtype)
-
-        def fake_prepare_latents(**kwargs):
-            captured_prepare_kwargs.append(kwargs)
-            return kwargs["image"], None
-
-        monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess)
-        pipe = object.__new__(LTX2Pipeline)
-        torch.nn.Module.__init__(pipe)
-        pipe.model_version = "2.5"
-        object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
-        object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
-        request_inputs = SimpleNamespace(
-            height=4,
-            width=4,
-            num_frames=3,
-            num_videos_per_prompt=1,
-            generator=None,
-            latents=None,
-        )
-        prompt_context = SimpleNamespace(
-            batch_size=1,
-            positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
-        )
-
-        pipe._prepare_video_latents_stage(
-            request_inputs,
-            prompt_context,
-            device=torch.device("cpu"),
-            noise_scale=0.0,
-            image=Image.new("RGB", (8, 4)),
-        )
-
-        assert captured_crfs == [18]
-        assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16
-        assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16
-
-    def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
-
-        pipe = object.__new__(LTX2Pipeline)
-        torch.nn.Module.__init__(pipe)
-        pipe.model_version = "2.5"
-
-        with pytest.raises(ValueError, match="requires PIL images"):
-            pipe._prepare_video_latents_stage(
-                SimpleNamespace(),
-                SimpleNamespace(),
-                device=torch.device("cpu"),
-                noise_scale=0.0,
-                image=torch.zeros(3, 4, 4),
-            )
-
-    def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
-
-        captured_prepare_kwargs = []
-
-        def fake_prepare_latents(**kwargs):
-            captured_prepare_kwargs.append(kwargs)
-            return kwargs["image"], None
-
-        pipe = object.__new__(LTX2Pipeline)
-        torch.nn.Module.__init__(pipe)
-        pipe.model_version = "2.5"
-        object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
-        object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
-        request_inputs = SimpleNamespace(
-            height=4,
-            width=4,
-            num_frames=3,
-            num_videos_per_prompt=1,
-            generator=None,
-            latents=None,
-            image_crf=0,
-        )
-        prompt_context = SimpleNamespace(
-            batch_size=1,
-            positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
-        )
-
-        pipe._prepare_video_latents_stage(
-            request_inputs,
-            prompt_context,
-            device=torch.device("cpu"),
-            noise_scale=0.0,
-            image=torch.zeros(3, 4, 4),
-        )
-
-        assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4)
-
-    def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
-
-        image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy()
-        monkeypatch.setitem(sys.modules, "av", None)
-
-        assert _apply_image_conditioning_crf(image, 18) is image
-
-    def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
-
-        monkeypatch.setitem(sys.modules, "av", None)
-
-        with pytest.raises(ImportError, match="PyAV with a libx264 encoder"):
-            _apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18)
-
    def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
        from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -229,7 +98,7 @@ class TestLTXImageToVideoForwardStages:
        with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
            pipe._resolve_request_image(mixed_req, None, request_inputs)

-    def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self):
+    def test_denoise_timestep_kwargs_masks_video_only(self):
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

        pipe = object.__new__(LTX2Pipeline)
@@ -320,18 +189,13 @@ class TestLTXImageToVideoConditioning:
        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
        torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))

-    @pytest.mark.parametrize(
-        ("model_version", "sampled_shape"),
-        [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
-    )
-    def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
+    def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

        pipe = object.__new__(LTX2Pipeline)
        torch.nn.Module.__init__(pipe)
-        pipe.model_version = model_version
        pipe.vae_spatial_compression_ratio = 1
        pipe.vae_temporal_compression_ratio = 1
        pipe.transformer_spatial_patch_size = 1
@@ -373,20 +237,14 @@ class TestLTXImageToVideoConditioning:

        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
        torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
-        assert sampled_shapes == [sampled_shape]
+        assert sampled_shapes == [(1, 3, 2)]

-    @pytest.mark.parametrize(
-        ("model_version", "sampled_shape"),
-        [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
-    )
-    def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
+    def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
-        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

        pipe = object.__new__(LTX2Pipeline)
        torch.nn.Module.__init__(pipe)
-        pipe.model_version = model_version
        pipe.vae_spatial_compression_ratio = 1
        pipe.vae_temporal_compression_ratio = 1
        pipe.transformer_spatial_patch_size = 1
@@ -408,7 +266,6 @@ class TestLTXImageToVideoConditioning:
            return torch.ones(shape, device=device, dtype=dtype)

        monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
-        monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)

        out, conditioning_mask = pipe.prepare_latents(
            image=torch.zeros(1, 3, 1, 1),
@@ -422,54 +279,10 @@ class TestLTXImageToVideoConditioning:
            device=torch.device("cpu"),
        )

-        assert sampled_shapes == [sampled_shape]
+        assert sampled_shapes == [(1, 3, 2)]
        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_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch):
-        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
-
-        pipe = object.__new__(LTX2Pipeline)
-        torch.nn.Module.__init__(pipe)
-        pipe.model_version = "2.5"
-        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),
-        )
-        object.__setattr__(
-            pipe,
-            "_encode_i2v_image_latents",
-            lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16),
-        )
-        captured = {}
-
-        def fake_create_noised_state(latents, noise_scale, generator=None):
-            del generator
-            captured["noise_scale"] = noise_scale
-            return latents
-
-        monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state)
-        pipe.prepare_latents(
-            image=torch.zeros(1, 3, 1, 1),
-            batch_size=1,
-            num_channels_latents=2,
-            height=1,
-            width=1,
-            num_frames=3,
-            noise_scale=0.909375,
-            dtype=torch.bfloat16,
-            device=torch.device("cpu"),
-            latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16),
-        )
-
-        assert captured["noise_scale"] == 0.909375
-
    def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

@@ -572,99 +385,6 @@ class TestLTXImageToVideoConditioning:
        torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
        assert pipeline.step_indices == [0, 1]

-    def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
-
-        device = torch.device("cpu")
-        sigmas = torch.tensor([1.0, 0.5, 0.0])
-        pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas))
-        mask = torch.tensor([[1.0, 0.0]])
-        adapter = LTXVideoAudioStepAdapter(
-            pipeline,
-            SimpleNamespace(sigmas=sigmas.clone()),
-            latent_num_frames=1,
-            latent_height=1,
-            latent_width=1,
-            image_conditioned=True,
-            sampler="euler_ancestral",
-            generator=torch.Generator(device=device).manual_seed(42),
-            conditioning_mask=mask,
-        )
-        video = torch.tensor([[[3.0], [2.0]]])
-        audio = torch.tensor([[[4.0]]])
-        video_velocity = torch.tensor([[[7.0], [0.25]]])
-        audio_velocity = torch.tensor([[[0.5]]])
-
-        reference_generator = torch.Generator(device=device).manual_seed(10042)
-        video_noise = torch.randn(video.shape, generator=reference_generator)
-        audio_noise = torch.randn(audio.shape, generator=reference_generator)
-
-        def official_step(sample, velocity, noise):
-            sigma = sigmas[0]
-            sigma_next = sigmas[1]
-            denoised = sample.float() - velocity.float() * sigma
-            sigma_down = sigma_next * (sigma_next / sigma)
-            ratio = sigma_down / sigma
-            deterministic = ratio * sample.float() + (1.0 - ratio) * denoised
-            alpha_next = 1.0 - sigma_next
-            alpha_down = 1.0 - sigma_down
-            coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
-            return (alpha_next / alpha_down) * deterministic + noise * coeff
-
-        expected_video = official_step(video, video_velocity, video_noise)
-        expected_video[:, :1] = video[:, :1]
-        expected_audio = official_step(audio, audio_velocity, audio_noise)
-
-        ((actual_video, actual_audio),) = adapter.step(
-            (video_velocity, audio_velocity),
-            (sigmas[0], sigmas[0]),
-            (video, audio),
-        )
-        torch.testing.assert_close(actual_video, expected_video)
-        torch.testing.assert_close(actual_audio, expected_audio)
-
-        terminal_video_velocity = torch.full_like(actual_video, 0.125)
-        terminal_audio_velocity = torch.full_like(actual_audio, 0.25)
-        ((terminal_video, terminal_audio),) = adapter.step(
-            (terminal_video_velocity, terminal_audio_velocity),
-            (sigmas[1], sigmas[1]),
-            (actual_video, actual_audio),
-        )
-        expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1]
-        expected_terminal_video[:, :1] = actual_video[:, :1]
-        expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1]
-        torch.testing.assert_close(terminal_video, expected_terminal_video)
-        torch.testing.assert_close(terminal_audio, expected_terminal_audio)
-
-    def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity
-
-        sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16)
-        velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16)
-        sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32)
-        generator = torch.Generator(device="cpu").manual_seed(10042)
-        reference_generator = torch.Generator(device="cpu").manual_seed(10042)
-
-        sigma, sigma_next = sigmas[:2]
-        denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float()
-        sigma_down = sigma_next * (sigma_next / sigma)
-        sigma_down_ratio = sigma_down / sigma
-        expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
-        alpha_next = 1.0 - sigma_next
-        alpha_down = 1.0 - sigma_down
-        coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
-        noise = torch.randn(
-            sample.shape,
-            generator=reference_generator,
-            dtype=sample.dtype,
-            device=sample.device,
-        )
-        expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype)
-
-        actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator)
-
-        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-
    def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
@@ -202,68 +202,3 @@ def test_prepare_audio_latents_rejects_incompatible_provided_length():
            device=torch.device("cpu"),
            latents=latents,
        )
-
-
-def test_create_noised_state_matches_official_fp32_lerp():
-    latents = torch.linspace(-2, 2, 4096, dtype=torch.bfloat16).reshape(1, 32, 128)
-    expected_generator = torch.Generator().manual_seed(42)
-    noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
-    expected = torch.lerp(latents.float(), noise.float(), torch.tensor(0.15, dtype=torch.bfloat16).float()).to(
-        latents.dtype
-    )
-
-    actual = latent_ops.create_noised_state(
-        latents,
-        torch.tensor(0.15, dtype=torch.bfloat16),
-        torch.Generator().manual_seed(42),
-    )
-
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-
-
-def test_create_conditioned_noised_state_matches_official_two_lerps():
-    latents = torch.linspace(-2, 2, 24, dtype=torch.bfloat16).reshape(1, 3, 8)
-    clean_latents = latents.clone()
-    clean_latents[:, 0] = 3.0
-    denoise_mask = torch.tensor([[[0.0], [1.0], [1.0]]])
-    expected_generator = torch.Generator().manual_seed(42)
-    noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
-    noised = torch.lerp(latents.float(), noise.float(), 0.909375)
-    expected = torch.lerp(clean_latents.float(), noised, denoise_mask).to(latents.dtype)
-
-    actual = latent_ops.create_conditioned_noised_state(
-        latents,
-        clean_latents,
-        denoise_mask,
-        0.909375,
-        torch.Generator().manual_seed(42),
-    )
-
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-
-
-def test_prepare_supplied_video_latents_uses_official_token_major_layout():
-    pipeline = _make_pipeline(LTX2Pipeline)
-    supplied = torch.arange(512, dtype=torch.float32).reshape(1, 32, 16)
-    actual_video = pipeline.prepare_latents(
-        batch_size=1,
-        num_channels_latents=4,
-        height=64,
-        width=64,
-        num_frames=9,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-        latents=supplied,
-    )
-
-    torch.testing.assert_close(actual_video, supplied, rtol=0, atol=0)
-    assert actual_video.stride()[1:] == (1, actual_video.shape[1])
-
-
-def test_clear_audio_padding_keeps_padding_outside_sampler_state():
-    updated = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [99.0, -99.0]]])
-
-    actual = latent_ops.clear_audio_padding(updated, 2)
-
-    torch.testing.assert_close(actual[:, :2], updated[:, :2])
-    torch.testing.assert_close(actual[:, 2:], torch.zeros_like(actual[:, 2:]))
@@ -28,15 +28,6 @@ def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
 class TestCFGParallelHelpers:
    """Test LTX-2.3 CFG helper math without loading model weights."""

-    def test_single_guidance_pass_preserves_official_token_layout(self):
-        from vllm_omni.diffusion.models.ltx2.ltx2_guidance import _repeat_batch
-
-        tokens = torch.randn(1, 16, 32).transpose(1, 2).contiguous().transpose(1, 2)
-        repeated = _repeat_batch(tokens, 1)
-
-        assert repeated is tokens
-        assert repeated.stride() == tokens.stride()
-
    def test_combine_cfg_noise_matches_x0_space_formula(self):
        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

@@ -20,7 +20,6 @@ from vllm_omni.diffusion.models.ltx2 import (
 from vllm_omni.diffusion.models.ltx2.ltx2_components import (
    LTX2_COMPONENT_PROFILE,
    LTX2_DISTILLED_COMPONENT_PROFILE,
-    LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
    LTX2_TWO_STAGE_COMPONENT_PROFILE,
    LTX23_COMPONENT_PROFILE,
    LTX23_DISTILLED_COMPONENT_PROFILE,
@@ -48,7 +47,6 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
 from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
-    LTX2_DISTILLED_ONE_STAGE_RECIPE,
    LTX2_DISTILLED_TWO_STAGE_RECIPE,
    LTX2_ONE_STAGE_RECIPE,
    LTX2_TWO_STAGE_RECIPE,
@@ -66,14 +64,12 @@ from vllm_omni.diffusion.models.ltx2.ltx2_request import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
-    LTX2DistilledOneStagePipeline,
    LTX2I2VDMD2Pipeline,
    LTX2Pipeline,
    LTX2T2VDMD2Pipeline,
 )
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
    LTX2DistilledPipeline,
-    LTX2DistilledTwoStagePipeline,
    LTX2TwoStagePipeline,
 )

@@ -129,21 +125,15 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():

    assert issubclass(LTX2Pipeline, LTXRuntime)
    assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
-    assert issubclass(LTX2DistilledOneStagePipeline, LTXRuntime)
-    assert issubclass(LTX2DistilledTwoStagePipeline, LTXRuntime)
    assert issubclass(LTX2DistilledPipeline, LTXRuntime)
-    assert LTX2DistilledPipeline is LTX2DistilledTwoStagePipeline
    assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
    assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
    assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
    assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
-    assert LTX2DistilledOneStagePipeline.component_profile is LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
-    assert LTX2DistilledOneStagePipeline.pipeline_recipe is LTX2_DISTILLED_ONE_STAGE_RECIPE
    assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
    for pipeline_cls, profile in (
        (LTX2Pipeline, LTX2_COMPONENT_PROFILE),
        (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
-        (LTX2DistilledOneStagePipeline, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE),
        (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
    ):
        assert pipeline_cls._dit_modules == list(profile.dit_modules)
@@ -264,7 +254,7 @@ def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_pa

    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)

-    with pytest.raises(ValueError, match="not qualified for this LTX recipe"):
+    with pytest.raises(ValueError, match="only one-stage LTX recipes"):
        LTX2TwoStagePipeline(
            od_config=SimpleNamespace(
                model=str(tmp_path),
@@ -719,48 +709,15 @@ def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch):

    assert len(installed) == 2
    assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed)
+    assert all(processor.has_learned_registers for processor in installed)
    assert all(kernel["role"] == "ltx2.connector" for kernel in kernels)
    assert all(kernel["role_category"] == "self" for kernel in kernels)
    assert all(kernel["skip_sequence_parallel"] for kernel in kernels)
    assert all(kernel["disable_kv_quant"] for kernel in kernels)
-    assert all(processor.has_learned_registers for processor in installed)
-    assert not any(processor.preserve_learned_register_mask for processor in installed)


-def test_ltx25_connector_attention_preserves_learned_register_mask(monkeypatch):
-    installed = []
-
-    class OmniAttention:
-        def __init__(self, **_kwargs):
-            pass
-
-    class Attention:
-        heads = 32
-        head_dim = 128
-        inner_kv_dim = 4096
-
-        def set_processor(self, processor):
-            installed.append(processor)
-
-    monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention)
-    connectors = SimpleNamespace(
-        video_connector=SimpleNamespace(
-            learnable_registers=object(),
-            transformer_blocks=[SimpleNamespace(attn1=Attention())],
-        ),
-        audio_connector=SimpleNamespace(
-            learnable_registers=object(),
-            transformer_blocks=[SimpleNamespace(attn1=Attention())],
-        ),
-    )
-
-    _install_connector_attention(connectors, preserve_learned_register_mask=True)
-
-    assert all(processor.has_learned_registers for processor in installed)
-    assert all(processor.preserve_learned_register_mask for processor in installed)
-
-
-def test_ltx_connector_attention_dispatches_through_omni_kernel():
+@pytest.mark.parametrize("has_learned_registers", [False, True])
+def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers):
    calls = []

    class Backend:
@@ -775,16 +732,12 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel():
            calls.append((query, key, value, metadata))
            return query

-    class UpcastingNorm(torch.nn.Module):
-        def forward(self, tensor):
-            return tensor.float()
-
    attention = SimpleNamespace(
        heads=2,
        head_dim=4,
        inner_kv_dim=8,
-        norm_q=UpcastingNorm(),
-        norm_k=UpcastingNorm(),
+        norm_q=torch.nn.Identity(),
+        norm_k=torch.nn.Identity(),
        to_q=torch.nn.Identity(),
        to_k=torch.nn.Identity(),
        to_v=torch.nn.Identity(),
@@ -793,29 +746,24 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel():
        rope_type="split",
        omni_attention=OmniAttention(),
    )
-    hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    hidden_states = torch.randn(2, 3, 8)
    additive_mask = torch.zeros(2, 1, 3, 3)
-    rotary = (
-        torch.randn(2, 2, 3, 2, dtype=torch.float32),
-        torch.randn(2, 2, 3, 2, dtype=torch.float32),
-    )

-    output = ltx2_components._LTXConnectorAttnProcessor()(
+    output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)(
        attention,
        hidden_states,
        attention_mask=additive_mask,
-        query_rotary_emb=rotary,
    )

    query, key, value, metadata = calls[0]
    assert query.shape == key.shape == value.shape == (2, 3, 2, 4)
-    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
-    assert metadata.attn_mask.shape == (2, 3)
-    assert metadata.attn_mask.dtype == torch.bool
-    assert metadata.attn_mask.all()
-    expected_rotary = tuple(component.to(torch.bfloat16) for component in rotary)
-    expected = ltx2_components.apply_split_rotary_emb(hidden_states, expected_rotary, head_dim=4)
-    torch.testing.assert_close(output, expected, rtol=0, atol=0)
+    if has_learned_registers:
+        assert metadata is None
+    else:
+        assert metadata.attn_mask.shape == (2, 3)
+        assert metadata.attn_mask.dtype == torch.bool
+        assert metadata.attn_mask.all()
+    torch.testing.assert_close(output, hidden_states)


 def test_ltx2_distilled_pipeline_shares_recipe_runtime():
@@ -1759,7 +1707,6 @@ class TestLTXForwardStages:
        assert seen["request_inputs"].output_type == "latent"
        assert seen["kwargs"] == {
            "request_sigmas": request_sigmas,
-            "request_phase_sigmas": None,
            "image": None,
        }

@@ -1877,16 +1824,6 @@ class TestRegistryIntegration:
            "pipeline_ltx2_two_stage",
            "LTX2TwoStagePipeline",
        )
-        assert _DIFFUSION_MODELS["LTX2DistilledOneStagePipeline"] == (
-            "ltx2",
-            "pipeline_ltx2",
-            "LTX2DistilledOneStagePipeline",
-        )
-        assert _DIFFUSION_MODELS["LTX2DistilledTwoStagePipeline"] == (
-            "ltx2",
-            "pipeline_ltx2_two_stage",
-            "LTX2DistilledTwoStagePipeline",
-        )
        assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
            "ltx2",
            "pipeline_ltx2_two_stage",
@@ -1911,8 +1848,6 @@ class TestRegistryIntegration:
        expected = [
            "LTX2Pipeline",
            "LTX2TwoStagePipeline",
-            "LTX2DistilledOneStagePipeline",
-            "LTX2DistilledTwoStagePipeline",
            "LTX2DistilledPipeline",
            "LTX2T2VDMD2Pipeline",
            "LTX2I2VDMD2Pipeline",
@@ -1926,8 +1861,6 @@ class TestRegistryIntegration:
        [
            "LTX2Pipeline",
            "LTX2TwoStagePipeline",
-            "LTX2DistilledOneStagePipeline",
-            "LTX2DistilledTwoStagePipeline",
            "LTX2DistilledPipeline",
            "LTX2T2VDMD2Pipeline",
            "LTX2I2VDMD2Pipeline",
@@ -7,54 +7,18 @@ import pytest
 import torch
 from torch import nn

+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
-from vllm_omni.diffusion.models.ltx2 import ltx2_transformer
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
-    LTX2FeedForward,
+    LTX2AudioVideoAttnProcessor,
    LTX2VideoTransformer3DModel,
-    LTX2VideoTransformerBlock,
+    _LTX2ParallelAttention,
    _make_rms_norm,
-    apply_interleaved_rotary_emb,
-    apply_keyframes_absolute_embedding,
-    apply_split_rotary_emb,
 )

 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


-def test_ltx_interleaved_rope_matches_official_model_dtype_arithmetic():
-    torch.manual_seed(0)
-    x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
-    cos = torch.randn(2, 3, 8, dtype=x.dtype)
-    sin = torch.randn(2, 3, 8, dtype=x.dtype)
-    x_real, x_imag = x.unflatten(2, (4, 2)).unbind(-1)
-    rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
-    expected = x * cos + rotated * sin
-
-    actual = apply_interleaved_rotary_emb(x, (cos, sin))
-
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-
-
-def test_ltx_split_rope_matches_official_model_dtype_arithmetic():
-    torch.manual_seed(1)
-    x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
-    cos = torch.randn(2, 3, 4, dtype=x.dtype)
-    sin = torch.randn(2, 3, 4, dtype=x.dtype)
-    split_x = x.reshape(2, 3, 2, 4)
-    first_x, second_x = split_x.chunk(2, dim=-2)
-    cos_u = cos.unsqueeze(-2)
-    sin_u = sin.unsqueeze(-2)
-    expected = split_x * cos_u
-    expected[..., :1, :].addcmul_(-sin_u, second_x)
-    expected[..., 1:, :].addcmul_(sin_u, first_x)
-    expected = expected.reshape_as(x)
-
-    actual = apply_split_rotary_emb(x, (cos, sin), head_dim=8)
-
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-
-
 def test_ltx_rms_norm_no_affine_matches_official_torch_module():
    norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)

@@ -72,53 +36,11 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
    assert "weight" in norm.state_dict()


-def test_ltx_attention_restores_projection_dtype_after_upcasting_qk_norms():
-    calls = []
-
-    class UpcastingNorm(nn.Module):
-        def forward(self, tensor):
-            return tensor.float()
-
-    class RecordingAttention:
-        def __call__(self, query, key, value, metadata):
-            calls.append((query, key, value, metadata))
-            return query
-
-    attention = SimpleNamespace(
-        heads=2,
-        head_dim=4,
-        to_qkv=None,
-        to_q=nn.Identity(),
-        to_k=nn.Identity(),
-        to_v=nn.Identity(),
-        norm_q=UpcastingNorm(),
-        norm_k=UpcastingNorm(),
-        to_gate_logits=None,
-        attn=RecordingAttention(),
-        to_out=(nn.Identity(), nn.Identity()),
-        rope_type="split",
-    )
-    hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
-    rotary = (
-        torch.ones(2, 2, 3, 2, dtype=torch.float32),
-        torch.zeros(2, 2, 3, 2, dtype=torch.float32),
-    )
-
-    processor = ltx2_transformer.LTX2AudioVideoAttnProcessor()
-    processor._slice_rope_for_tp = lambda rope, _: rope
-    output = processor(attention, hidden_states, query_rotary_emb=rotary)
-
-    query, key, value, metadata = calls[0]
-    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
-    assert metadata is None
-    torch.testing.assert_close(output, hidden_states)
-
-
-def test_ltx_transformer_has_fused_cfg_cache_dit_config():
+def test_ltx_transformer_has_separate_cfg_cache_dit_config():
    adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")

    assert isinstance(adapter_config, CacheDiTAdapterConfig)
-    assert not adapter_config.has_separate_cfg
+    assert adapter_config.has_separate_cfg


 def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
@@ -154,9 +76,8 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):


 def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
-    processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor
-    monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True))
-    processor = processor_cls()
+    monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
+    processor = LTX2AudioVideoAttnProcessor()
    hidden_states = torch.zeros(1, 3, 4)
    encoder_hidden_states = torch.zeros(1, 4, 4)
    key_padding_mask = torch.tensor([[True, True, True, False]])
@@ -192,7 +113,7 @@ def _fake_ltx_parallel_attention():

        return forward

-    layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
+    layer = object.__new__(_LTX2ParallelAttention)
    nn.Module.__init__(layer)
    layer.attn_backend = Backend
    layer.attention = SimpleNamespace(forward=run("native"))
@@ -202,281 +123,33 @@ def _fake_ltx_parallel_attention():


@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
-def test_ltx_masked_attention_uses_sdpa(query_length, key_length):
+def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
    layer, calls = _fake_ltx_parallel_attention()
    query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
    key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
    key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
-    metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
+    metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])

    layer._run_local_attention(query, key, value, metadata)

    assert calls == ["sdpa"]


-def test_ltx_unmasked_attention_keeps_native_backend():
+def test_ltx_masked_self_attention_uses_sdpa():
+    layer, calls = _fake_ltx_parallel_attention()
+    query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
+    key_padding_mask = torch.tensor([[True, True, True, False]])
+    metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
+
+    layer._run_local_attention(query, key, value, metadata)
+
+    assert calls == ["sdpa"]
+
+
+def test_ltx_unmasked_cross_attention_keeps_native_backend():
    layer, calls = _fake_ltx_parallel_attention()
    query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)

    layer._run_local_attention(query, key, value, None)

    assert calls == ["native"]
-
-
-def test_ltx25_feed_forward_bias_flag_controls_both_linears(monkeypatch):
-    class FakeParallelLinear(nn.Linear):
-        def __init__(self, dim_in, dim_out, *, bias=True, **kwargs):
-            super().__init__(dim_in, dim_out, bias=bias)
-
-    monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeParallelLinear)
-    monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeParallelLinear)
-
-    without_bias = LTX2FeedForward(8, inner_dim=16, bias=False)
-    with_bias = LTX2FeedForward(8, inner_dim=16, bias=True)
-
-    assert without_bias.net[0].proj.bias is None
-    assert without_bias.net[2].bias is None
-    assert with_bias.net[0].proj.bias is not None
-    assert with_bias.net[2].bias is not None
-
-
-def _build_tiny_ltx_transformer(monkeypatch, **overrides):
-    class FakeBlock(nn.Module):
-        def __init__(self, **kwargs):
-            super().__init__()
-            self.init_kwargs = kwargs
-
-        def forward(self, hidden_states, audio_hidden_states, *_args, **_kwargs):
-            return hidden_states, audio_hidden_states
-
-    monkeypatch.setattr(ltx2_transformer, "LTX2VideoTransformerBlock", FakeBlock)
-    kwargs = {
-        "in_channels": 4,
-        "out_channels": 4,
-        "num_attention_heads": 2,
-        "attention_head_dim": 4,
-        "cross_attention_dim": 8,
-        "vae_scale_factors": (1, 1, 1),
-        "pos_embed_max_pos": 4,
-        "base_height": 8,
-        "base_width": 8,
-        "audio_in_channels": 4,
-        "audio_out_channels": 4,
-        "audio_num_attention_heads": 2,
-        "audio_attention_head_dim": 4,
-        "audio_cross_attention_dim": 8,
-        "audio_scale_factor": 1,
-        "audio_pos_embed_max_pos": 4,
-        "audio_sampling_rate": 16,
-        "audio_hop_length": 4,
-        "num_layers": 1,
-        "caption_channels": 8,
-        "use_prompt_embeddings": False,
-        "cross_attn_mod": True,
-        "audio_cross_attn_mod": True,
-    }
-    kwargs.update(overrides)
-    return LTX2VideoTransformer3DModel(**kwargs)
-
-
-def test_ltx25_keyframe_absolute_embedding_matches_official():
-    hidden_states = torch.arange(32, dtype=torch.bfloat16).view(1, 4, 8)
-    keyframes_mask = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]])
-    embedding = torch.linspace(-0.5, 0.5, 8, dtype=torch.bfloat16).view(1, 8)
-
-    actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, embedding)
-    expected = hidden_states + (keyframes_mask > 0).to(torch.bfloat16) * embedding
-
-    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-    assert apply_keyframes_absolute_embedding(hidden_states, None, embedding) is hidden_states
-
-
-def test_ltx25_constructor_flags_and_native_weight_loader(monkeypatch):
-    legacy_model = _build_tiny_ltx_transformer(monkeypatch)
-
-    assert legacy_model.config.ff_bias is True
-    assert legacy_model.config.audio_ff_bias is True
-    assert legacy_model.config.use_prompt_adaln_single is True
-    assert legacy_model.config.use_keyframes_abs_pos_embedding is False
-    assert hasattr(legacy_model, "prompt_adaln")
-    assert hasattr(legacy_model, "audio_prompt_adaln")
-    assert not hasattr(legacy_model, "keyframes_abs_pos_embedding")
-    assert legacy_model.transformer_blocks[0].init_kwargs["ff_bias"] is True
-    assert legacy_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
-
-    ltx25_model = _build_tiny_ltx_transformer(
-        monkeypatch,
-        ff_bias=False,
-        audio_ff_bias=True,
-        use_prompt_adaln_single=False,
-        use_keyframes_abs_pos_embedding=True,
-    )
-
-    assert ltx25_model.keyframes_abs_pos_embedding.shape == (1, 8)
-    torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, torch.zeros(1, 8))
-    assert not hasattr(ltx25_model, "prompt_adaln")
-    assert not hasattr(ltx25_model, "audio_prompt_adaln")
-    assert ltx25_model.transformer_blocks[0].init_kwargs["ff_bias"] is False
-    assert ltx25_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
-
-    monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 1)
-    expected = torch.arange(8, dtype=torch.float32).view(1, 8)
-    loaded = ltx25_model.load_weights([("keyframes_abs_pos_embedding", expected)])
-
-    assert loaded == {"keyframes_abs_pos_embedding"}
-    torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, expected)
-
-
-def test_ltx_cross_attention_modulations_use_official_timesteps(monkeypatch):
-    class CaptureAda(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.calls = []
-
-        def forward(self, timestep, *, batch_size, hidden_dtype):
-            self.calls.append(timestep.detach().clone())
-            embedding = torch.zeros(timestep.numel(), 8, dtype=hidden_dtype)
-            return embedding, embedding
-
-    class CaptureRope(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.output_dtypes = []
-
-        def forward(self, _coords, *, device, out_dtype=None):
-            self.output_dtypes.append(out_dtype)
-            empty = torch.empty(0, device=device, dtype=out_dtype)
-            return empty, empty
-
-    model = _build_tiny_ltx_transformer(monkeypatch)
-    rope = CaptureRope()
-    model.rope = model.audio_rope = model.cross_attn_rope = model.cross_attn_audio_rope = rope
-    model.prompt_modulation = False
-    video_scale_shift = CaptureAda()
-    video_gate = CaptureAda()
-    audio_scale_shift = CaptureAda()
-    audio_gate = CaptureAda()
-    model.av_cross_attn_video_scale_shift = video_scale_shift
-    model.av_cross_attn_video_a2v_gate = video_gate
-    model.av_cross_attn_audio_scale_shift = audio_scale_shift
-    model.av_cross_attn_audio_v2a_gate = audio_gate
-
-    model(
-        hidden_states=torch.zeros(1, 2, 4),
-        audio_hidden_states=torch.zeros(1, 1, 4),
-        encoder_hidden_states=torch.zeros(1, 1, 8),
-        audio_encoder_hidden_states=torch.zeros(1, 1, 8),
-        timestep=torch.tensor([[0.0, 2.0]]),
-        audio_timestep=torch.tensor([[5.0]]),
-        sigma=torch.tensor([2.0]),
-        audio_sigma=torch.tensor([5.0]),
-        video_coords=torch.zeros(1, 3, 2, 2),
-        audio_coords=torch.zeros(1, 1, 1, 2),
-    )
-
-    torch.testing.assert_close(video_scale_shift.calls[0], torch.tensor([0.0, 2.0]))
-    torch.testing.assert_close(video_gate.calls[0], torch.tensor([5.0]))
-    torch.testing.assert_close(audio_scale_shift.calls[0], torch.tensor([5.0]))
-    torch.testing.assert_close(audio_gate.calls[0], torch.tensor([2.0]))
-    assert len(rope.output_dtypes) == 4
-    assert all(dtype is torch.float32 for dtype in rope.output_dtypes)
-
-
-def test_ltx25_static_prompt_table_is_used_without_prompt_adaln():
-    class CaptureAttention(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.encoder_hidden_states = None
-
-        def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
-            if encoder_hidden_states is not None:
-                self.encoder_hidden_states = encoder_hidden_states.detach().clone()
-            return torch.zeros_like(hidden_states)
-
-    class ZeroOutput(nn.Module):
-        def forward(self, hidden_states):
-            return torch.zeros_like(hidden_states)
-
-    block = object.__new__(LTX2VideoTransformerBlock)
-    nn.Module.__init__(block)
-    block.video_cross_attn_adaln = True
-    block.audio_cross_attn_adaln = True
-    block.cross_attn_adaln = True
-    block.scale_shift_table = nn.Parameter(torch.zeros(9, 2))
-    block.audio_scale_shift_table = nn.Parameter(torch.zeros(9, 2))
-    block.prompt_scale_shift_table = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
-    block.audio_prompt_scale_shift_table = nn.Parameter(torch.tensor([[0.5, 1.0], [1.5, 2.0]]))
-    block.norm1 = nn.Identity()
-    block.audio_norm1 = nn.Identity()
-    block.norm2 = nn.Identity()
-    block.audio_norm2 = nn.Identity()
-    block.norm3 = nn.Identity()
-    block.audio_norm3 = nn.Identity()
-    block.attn1 = CaptureAttention()
-    block.audio_attn1 = CaptureAttention()
-    block.attn2 = CaptureAttention()
-    block.audio_attn2 = CaptureAttention()
-    block.ff = ZeroOutput()
-    block.audio_ff = ZeroOutput()
-
-    hidden_states = torch.ones(1, 1, 2)
-    audio_hidden_states = torch.ones(1, 1, 2)
-    encoder_hidden_states = torch.ones(1, 1, 2)
-    audio_encoder_hidden_states = torch.ones(1, 1, 2)
-    block(
-        hidden_states,
-        audio_hidden_states,
-        encoder_hidden_states,
-        audio_encoder_hidden_states,
-        temb=torch.zeros(1, 1, 18),
-        temb_audio=torch.zeros(1, 1, 18),
-        temb_ca_scale_shift=torch.empty(0),
-        temb_ca_audio_scale_shift=torch.empty(0),
-        temb_ca_gate=torch.empty(0),
-        temb_ca_audio_gate=torch.empty(0),
-        temb_prompt=None,
-        temb_prompt_audio=None,
-        use_a2v_cross_attention=False,
-        use_v2a_cross_attention=False,
-    )
-
-    torch.testing.assert_close(block.attn2.encoder_hidden_states, torch.tensor([[[5.0, 7.0]]]))
-    torch.testing.assert_close(block.audio_attn2.encoder_hidden_states, torch.tensor([[[3.0, 4.0]]]))
-
-
-def test_ltx_gated_attention_uses_tp_sharded_projection(monkeypatch):
-    column_calls = []
-
-    class FakeColumnParallelLinear(nn.Linear):
-        def __init__(self, input_size, output_size, *, bias=True, **kwargs):
-            column_calls.append((output_size, kwargs))
-            super().__init__(input_size, output_size, bias=bias)
-
-    class FakeRowParallelLinear(nn.Linear):
-        def __init__(self, input_size, output_size, *, bias=True, **_kwargs):
-            super().__init__(input_size, output_size, bias=bias)
-
-    class FakeParallelAttention(nn.Module):
-        def __init__(self, **_kwargs):
-            super().__init__()
-
-    monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumnParallelLinear)
-    monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRowParallelLinear)
-    monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeParallelAttention)
-    monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
-
-    attention = ltx2_transformer.LTX2Attention(
-        query_dim=8,
-        heads=4,
-        kv_heads=4,
-        dim_head=2,
-        qk_norm="rms_norm",
-        apply_gated_attention=True,
-        pack_qkv=False,
-    )
-
-    gate_output_size, gate_kwargs = column_calls[-1]
-    assert gate_output_size == 4
-    assert gate_kwargs["gather_output"] is False
-    assert gate_kwargs["return_bias"] is False
-    assert attention.query_num_heads == 2
@@ -2,10 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 import json
-import sys
 from multiprocessing.reduction import ForkingPickler
 from types import SimpleNamespace
-from unittest.mock import Mock, patch
+from unittest.mock import Mock

 import numpy as np
 import pytest
@@ -363,18 +362,8 @@ def test_pipeline_loads_task_selected_dits_and_shared_components_once(
    if expected_partition == "ref2va":
        assert pipeline._transformer_for_task("ref2va") is pipeline.transformer
    assert [source.model_or_path for source in pipeline.weights_sources] == [
-        *(str(tmp_path / partition_name) for partition_name in source_partitions),
-        str(component_path),
+        str(tmp_path / partition_name) for partition_name in source_partitions
    ]
-    assert [source.subfolder for source in pipeline.weights_sources] == [
-        *("transformer" for _ in source_partitions),
-        "text_encoder",
-    ]
-    expected_prefixes = ["transformer."]
-    if expected_dits == 2:
-        expected_prefixes.append("transformers_ref.")
-    expected_prefixes.append("text_encoder.")
-    assert [source.prefix for source in pipeline.weights_sources] == expected_prefixes
    expected_patterns = (
        pipeline_module.MINIMAX_H3_DOWNLOAD_PATTERNS
        if expected_partition == "combined"
@@ -946,66 +935,6 @@ def test_text_encoder_stub_constructs_without_group_or_weights():
    assert list(encoder.named_parameters()) == []


-def test_global_quant_config_is_shared_by_dit_and_encoder():
-    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
-
-    global_config = build_quant_config("fp8")
-
-    assert resolve_component_quant_config(global_config, "transformer") is global_config
-    assert resolve_component_quant_config(global_config, "text_encoder") is global_config
-
-
-def test_minimax_h3_fp8_component_selection():
-    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
-
-    dit_only = build_quant_config({"transformer": {"method": "fp8"}})
-    encoder_only = build_quant_config({"text_encoder": {"method": "fp8"}})
-
-    assert resolve_component_quant_config(dit_only, "transformer").get_name() == "fp8"
-    assert resolve_component_quant_config(dit_only, "text_encoder") is None
-    assert resolve_component_quant_config(encoder_only, "transformer") is None
-    assert resolve_component_quant_config(encoder_only, "text_encoder").get_name() == "fp8"
-
-
-def test_text_encoder_rejects_serialized_fp8():
-    from vllm_omni.diffusion.models.minimax_h3.encoder import _validate_encoder_quant_config
-
-    config = SimpleNamespace(
-        get_name=lambda: "fp8",
-        is_checkpoint_fp8_serialized=True,
-    )
-    with pytest.raises(ValueError, match="online FP8 only"):
-        _validate_encoder_quant_config(config)
-
-
-def test_text_encoder_linear_delegates_quantization_to_vllm_factory():
-    from vllm.model_executor.layers.linear import UnquantizedLinearMethod
-
-    import vllm_omni.diffusion.models.minimax_h3.encoder as encoder_module
-
-    group = SimpleNamespace(rank_in_group=0, world_size=1)
-    method = UnquantizedLinearMethod()
-    quant_config = Mock()
-    quant_config.get_quant_method.return_value = method
-    prefix = "text_encoder.text_model.layers.0.mlp.gate_up_proj"
-
-    with (
-        patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
-        patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
-    ):
-        linear = encoder_module.MiniMaxH3Qwen3VLMergedColumnParallelLinear(
-            group,
-            input_size=4,
-            intermediate_size=8,
-            dtype=torch.bfloat16,
-            quant_config=quant_config,
-            prefix=prefix,
-        )
-
-    assert linear.quant_method is method
-    quant_config.get_quant_method.assert_called_once_with(linear, prefix=prefix)
-
-
 def test_no_offload_keeps_text_encoder_resident():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

@@ -1048,21 +977,15 @@ def test_model_offload_uses_hooked_text_encoder_call():
    pipeline.text_encoder.load_to_device.assert_not_called()


-@pytest.mark.parametrize(
-    "offload_flag",
-    ["enable_layerwise_offload", "enable_distributed_layerwise_offload"],
-)
-def test_layerwise_offload_releases_text_encoder(offload_flag):
+def test_layerwise_offload_releases_text_encoder():
    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline

    pipeline = object.__new__(MiniMaxH3Pipeline)
    torch.nn.Module.__init__(pipeline)
    pipeline.od_config = SimpleNamespace(
        enable_cpu_offload=False,
-        enable_layerwise_offload=False,
-        enable_distributed_layerwise_offload=False,
+        enable_layerwise_offload=True,
    )
-    setattr(pipeline.od_config, offload_flag, True)
    pipeline.text_encoder = Mock()
    expected = torch.ones(2, 3)
    pipeline.text_encoder.encode_ids.return_value = expected
@@ -1505,219 +1428,34 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti
    assert transcode_calls[0][1]["target_frame_count"] == 209


-def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch):
+def test_ref2va_qwen_sampling_decodes_prepared_video_once(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module

-    decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3)
-    decode_calls = []
+    decoded = np.arange(25 * 2 * 2 * 3, dtype=np.uint8).reshape(25, 2, 2, 3)
+    load_calls = []
    monkeypatch.setattr(
        reference_video_module,
        "_probe_video",
-        lambda _path: {
-            "frame_count": 25,
-            "width": 2,
-            "height": 2,
-        },
+        lambda _path: {"frame_count": len(decoded)},
    )
    monkeypatch.setattr(
        reference_video_module,
-        "_decode_video_frames_ffmpeg",
-        lambda path, **kwargs: decode_calls.append((path, kwargs)) or decoded,
+        "load_video_frames",
+        lambda path: load_calls.append(path) or decoded,
    )

    sampled = reference_video_module.sample_reference_video_frames("prepared.mp4")

-    assert decode_calls == [
-        (
-            "prepared.mp4",
-            {
-                "frame_count": 25,
-                "indices": [0, 12, 24],
-                "width": 2,
-                "height": 2,
-            },
-        )
-    ]
-    for frame, expected in zip(sampled["frames"], decoded, strict=True):
+    assert load_calls == ["prepared.mp4"]
+    for frame, expected in zip(
+        sampled["frames"],
+        (decoded[0], decoded[12], decoded[24]),
+        strict=True,
+    ):
        np.testing.assert_array_equal(frame, expected)
    assert sampled["block_timestamps"] == [0.25, 1.0]


-@pytest.mark.parametrize(
-    ("indices", "frame_count"),
-    [([0, 12], 25), (None, 2)],
-)
-def test_ref2va_ffmpeg_decode_reads_exact_rgb_frames(monkeypatch, indices, frame_count):
-    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
-
-    expected = np.arange(2 * 2 * 3 * 3, dtype=np.uint8).reshape(2, 2, 3, 3)
-    observed = []
-
-    def fake_run(command, **kwargs):
-        observed.append((command, kwargs))
-        return SimpleNamespace(stdout=expected.tobytes())
-
-    monkeypatch.setattr(reference_video_module.subprocess, "run", fake_run)
-
-    actual = reference_video_module._decode_video_frames_ffmpeg(
-        "prepared.mp4",
-        frame_count=frame_count,
-        indices=indices,
-        width=3,
-        height=2,
-    )
-
-    np.testing.assert_array_equal(actual, expected)
-    command, kwargs = observed[0]
-    assert command[0] == "ffmpeg"
-    assert command[command.index("-frames:v") + 1] == "2"
-    if indices is None:
-        assert "-vf" not in command
-    else:
-        assert command[command.index("-vf") + 1] == "select=eq(n\\,0)+eq(n\\,12)"
-    assert command[-1] == "pipe:1"
-    assert kwargs == {"check": True, "capture_output": True}
-
-
-def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
-
-    expected = np.zeros((2, 2, 3, 3), dtype=np.uint8)
-    decode_calls = []
-    monkeypatch.setitem(sys.modules, "decord", None)
-    monkeypatch.setattr(
-        reference_video_module,
-        "_probe_video",
-        lambda _path: {
-            "frame_count": 2,
-            "width": 3,
-            "height": 2,
-        },
-    )
-    monkeypatch.setattr(
-        reference_video_module,
-        "_decode_video_frames_ffmpeg",
-        lambda path, **kwargs: decode_calls.append((path, kwargs)) or expected,
-    )
-
-    actual = reference_video_module.load_video_frames("prepared.mp4")
-
-    assert actual is expected
-    assert decode_calls == [
-        (
-            "prepared.mp4",
-            {
-                "frame_count": 2,
-                "width": 3,
-                "height": 2,
-            },
-        )
-    ]
-
-
-@pytest.mark.parametrize(
-    ("kwargs", "message"),
-    [
-        ({"frame_count": 0, "width": 3, "height": 2}, "video has no frames"),
-        ({"frame_count": 1, "width": 0, "height": 2}, "invalid dimensions"),
-        (
-            {"frame_count": 1, "indices": [], "width": 3, "height": 2},
-            "invalid frame indices",
-        ),
-    ],
-)
-def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message):
-    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
-    from vllm_omni.errors import OmniClientError
-
-    with pytest.raises(OmniClientError, match=message):
-        reference_video_module._decode_video_frames_ffmpeg("prepared.mp4", **kwargs)
-
-
-def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
-
-    calls = []
-    probe = {
-        "streams": [
-            {
-                "codec_type": "video",
-                "width": 4,
-                "height": 2,
-                "r_frame_rate": "24/1",
-                "nb_frames": "3",
-                "duration": "0.125",
-                "sample_aspect_ratio": "1:1",
-            }
-        ],
-        "format": {"duration": "0.125", "size": "123"},
-    }
-    monkeypatch.setattr(
-        reference_video_module.subprocess,
-        "run",
-        lambda command, **kwargs: calls.append((command, kwargs)) or SimpleNamespace(stdout=json.dumps(probe)),
-    )
-
-    metadata = reference_video_module._probe_video("prepared.mp4")
-
-    assert metadata["frame_count"] == 3
-    assert len(calls) == 1
-    assert "-count_frames" not in calls[0][0]
-
-
-def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
-
-    calls = []
-    stream = {
-        "codec_type": "video",
-        "width": 4,
-        "height": 2,
-        "r_frame_rate": "24/1",
-        "duration": "0.125",
-        "sample_aspect_ratio": "1:1",
-    }
-    first_probe = {
-        "streams": [stream],
-        "format": {"duration": "0.125", "size": "123"},
-    }
-    second_probe = {
-        "streams": [{**stream, "nb_read_frames": "3"}],
-        "format": {"duration": "0.125", "size": "123"},
-    }
-    monkeypatch.setattr(
-        reference_video_module.subprocess,
-        "run",
-        lambda command, **kwargs: calls.append((command, kwargs))
-        or SimpleNamespace(stdout=json.dumps(first_probe if len(calls) == 1 else second_probe)),
-    )
-
-    metadata = reference_video_module._probe_video("prepared.mp4")
-
-    assert metadata["frame_count"] == 3
-    assert len(calls) == 2
-    assert "-count_frames" in calls[1][0]
-
-
-def test_ref2va_ffmpeg_decode_rejects_incomplete_output(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
-    from vllm_omni.errors import OmniClientError
-
-    monkeypatch.setattr(
-        reference_video_module.subprocess,
-        "run",
-        lambda *_args, **_kwargs: SimpleNamespace(stdout=b"short"),
-    )
-
-    with pytest.raises(OmniClientError, match="decoded 5 bytes.*expected 18"):
-        reference_video_module._decode_video_frames_ffmpeg(
-            "prepared.mp4",
-            frame_count=1,
-            width=3,
-            height=2,
-        )
-
-
 def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path):
    from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module

@@ -2064,25 +1802,21 @@ def _one_layer_text_encoder():
    group = _SingleRankEncoderGroup(0)
    layer = nn.Module()
    layer.self_attn = nn.Module()
+    layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
+        group,
+        hidden_size=_ENCODER_HIDDEN,
+        num_heads=_ENCODER_NUM_HEADS,
+        num_kv_heads=_ENCODER_NUM_KV_HEADS,
+        head_dim=_ENCODER_HEAD_DIM,
+        dtype=torch.float32,
+    )
    layer.mlp = nn.Module()
-    with (
-        patch("vllm.model_executor.parameter.get_tensor_model_parallel_rank", return_value=0),
-        patch("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", return_value=1),
-    ):
-        layer.self_attn.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear(
-            group,
-            hidden_size=_ENCODER_HIDDEN,
-            num_heads=_ENCODER_NUM_HEADS,
-            num_kv_heads=_ENCODER_NUM_KV_HEADS,
-            head_dim=_ENCODER_HEAD_DIM,
-            dtype=torch.float32,
-        )
-        layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
-            group,
-            input_size=_ENCODER_HIDDEN,
-            intermediate_size=_ENCODER_INTERMEDIATE,
-            dtype=torch.float32,
-        )
+    layer.mlp.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear(
+        group,
+        input_size=_ENCODER_HIDDEN,
+        intermediate_size=_ENCODER_INTERMEDIATE,
+        dtype=torch.float32,
+    )
    layer.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(_ENCODER_HIDDEN)

    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
@@ -2115,7 +1849,6 @@ def _write_text_encoder_checkpoint(path, *, omit=()):
    safetensors.torch.save_file(tensors, str(path / "model.safetensors"))
    index = {"weight_map": dict.fromkeys(tensors, "model.safetensors")}
    (path / "model.safetensors.index.json").write_text(json.dumps(index))
-    return tensors


@pytest.mark.parametrize(
@@ -2137,10 +1870,10 @@ def test_encoder_load_reports_missing_fused_source_shards(
    tmp_path, omitted_sources, expected_detail, unreported_target
 ):
    encoder = _one_layer_text_encoder()
-    weights = _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)
+    _write_text_encoder_checkpoint(tmp_path, omit=omitted_sources)

    with pytest.raises(RuntimeError) as excinfo:
-        encoder.load_weights(weights.items())
+        encoder._load_weights(str(tmp_path))

    message = str(excinfo.value)
    assert expected_detail in message
@@ -2149,17 +1882,17 @@ def test_encoder_load_reports_missing_fused_source_shards(

 def test_encoder_load_rejects_unloaded_plain_param(tmp_path):
    encoder = _one_layer_text_encoder()
-    weights = _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))
+    _write_text_encoder_checkpoint(tmp_path, omit=("input_layernorm.weight",))

    with pytest.raises(RuntimeError, match=r"weights not loaded.*input_layernorm"):
-        encoder.load_weights(weights.items())
+        encoder._load_weights(str(tmp_path))


 def test_encoder_load_places_every_source_in_its_own_rows(tmp_path):
    encoder = _one_layer_text_encoder()
-    weights = _write_text_encoder_checkpoint(tmp_path)
+    _write_text_encoder_checkpoint(tmp_path)

-    encoder.load_weights(weights.items())
+    encoder._load_weights(str(tmp_path))

    params = dict(encoder.named_parameters())
    q_rows = _ENCODER_NUM_HEADS * _ENCODER_HEAD_DIM
@@ -36,13 +36,7 @@ def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp():
    assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"]
    assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3
    assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg
-    assert set(MiniMaxH3DiTModel._sp_plan) == {
-        "sp_prepare",
-        "local_sp_prepare",
-        "sp_gather",
-    }
-    assert set(MiniMaxH3DiTModel._sp_plan["sp_prepare"]) == {0, 1, 2}
-    assert set(MiniMaxH3DiTModel._sp_plan["local_sp_prepare"]) == {2}
+    assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"}

    model = object.__new__(MiniMaxH3DiTModel)
    nn.Module.__init__(model)
@@ -95,279 +89,6 @@ def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
    torch.testing.assert_close(actual[..., 96:], x[..., 96:], atol=0, rtol=0)


-def test_h3_rope_table_materializes_local_rows_in_fused_kernel_layout():
-    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
-        _build_rope_table,
-    )
-
-    local_freqs_half = torch.randn(2, 48)
-    local_freqs = torch.cat((local_freqs_half, local_freqs_half), dim=-1)
-
-    actual = _build_rope_table(local_freqs)
-    expected = torch.cat(
-        (torch.cos(local_freqs_half), torch.sin(local_freqs_half)),
-        dim=-1,
-    ).to(torch.bfloat16)
-
-    assert actual.shape == (2, 96)
-    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
-
-
-def test_strict_sp_local_span_uses_rank_owned_rows(monkeypatch):
-    from vllm_omni.diffusion import forward_context
-    from vllm_omni.diffusion.distributed import parallel_state
-    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
-
-    monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
-    monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: "strict")
-    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
-    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
-    monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 8)
-    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
-    monkeypatch.setattr(parallel_state, "get_allgather_parallel_world_size", lambda: 1)
-
-    assert h3._sequence_parallel_local_span(16, hooks_applied=True) == (6, 2)
-
-
-@pytest.mark.parametrize(
-    ("hooks_applied", "mode", "seq_len", "ulysses", "ring", "allgather"),
-    [
-        (False, "strict", 16, 8, 1, 1),
-        (True, "advanced_uaa", 16, 8, 1, 1),
-        (True, "strict", 17, 8, 1, 1),
-        (True, "strict", 16, 4, 2, 1),
-        (True, "strict", 16, 1, 1, 8),
-    ],
-)
-def test_local_span_falls_back_when_local_embedding_is_unsafe(
-    monkeypatch,
-    hooks_applied,
-    mode,
-    seq_len,
-    ulysses,
-    ring,
-    allgather,
-):
-    from vllm_omni.diffusion import forward_context
-    from vllm_omni.diffusion.distributed import parallel_state
-    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
-
-    monkeypatch.setattr(forward_context, "is_forward_context_available", lambda: True)
-    monkeypatch.setattr(forward_context, "get_ulysses_mode", lambda **kwargs: mode)
-    monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 8)
-    monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 3)
-    monkeypatch.setattr(
-        parallel_state,
-        "get_ulysses_parallel_world_size",
-        lambda: ulysses,
-    )
-    monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: ring)
-    monkeypatch.setattr(
-        parallel_state,
-        "get_allgather_parallel_world_size",
-        lambda: allgather,
-    )
-
-    assert h3._sequence_parallel_local_span(seq_len, hooks_applied=hooks_applied) == (
-        0,
-        seq_len,
-    )
-
-
-def test_local_embedding_spans_reassemble_multimodal_rows():
-    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
-
-    class _IdentityProjection(nn.Module):
-        def forward(self, x):
-            return x, None
-
-    class _IdentityRefiner(nn.Module):
-        def forward(self, x, **kwargs):
-            return x
-
-    model = object.__new__(h3.MiniMaxH3DiTModel)
-    nn.Module.__init__(model)
-    model.hidden_size = 2
-    model.video_patch_proj = _IdentityProjection()
-    model.audio_patch_proj = _IdentityProjection()
-    model.condition_proj = _IdentityProjection()
-    model.token_refiner = _IdentityRefiner()
-    model.time_embedder = nn.Identity()
-
-    seq_len = 16
-    x = torch.arange(seq_len * 2, dtype=torch.float32).reshape(1, seq_len, 2)
-    audio_x = x + 100
-    text = torch.tensor(
-        [[200.0, 201.0], [202.0, 203.0], [204.0, 205.0], [206.0, 207.0]],
-        dtype=torch.bfloat16,
-    )
-    common = {
-        "x": x,
-        "audio_x": audio_x,
-        "text_embeddings_selected": text,
-        "unique_timesteps": torch.tensor([0.2, 0.3]),
-        "img_pos": torch.tensor([1, 2, 6, 7, 11, 12]),
-        "audio_pos": torch.tensor([3, 4, 8, 9, 13, 14]),
-        "text_pos": torch.tensor([0, 5, 10, 15]),
-        "refiner_cu_seqlens": torch.tensor([0, 4, 4], dtype=torch.int32),
-        "refiner_max_seqlen": 4,
-        "seq_len": seq_len,
-        "device": torch.device("cpu"),
-        "local_span": (0, seq_len),
-    }
-
-    full, full_t_emb = model._embed(**common)
-    local_spans = []
-    for rank in range(8):
-        local_span = (rank * (seq_len // 8), seq_len // 8)
-        local, local_t_emb = model._embed(
-            **{
-                **common,
-                "local_span": local_span,
-            }
-        )
-        assert local.shape == (2, 2)
-        torch.testing.assert_close(local_t_emb, full_t_emb)
-        local_spans.append(local)
-
-    torch.testing.assert_close(torch.cat(local_spans), full, atol=0, rtol=0)
-
-
-@pytest.mark.parametrize(
-    "local_projection",
-    [False, True],
-    ids=["fallback", "strict-local"],
-)
-def test_sp_output_boundary_preserves_logits_and_minimizes_gather(
-    monkeypatch,
-    local_projection,
-):
-    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
-
-    class _StaticRope(nn.Module):
-        input_rows = None
-
-        def forward(self, position_ids):
-            self.input_rows = position_ids.shape[1]
-            return torch.zeros(position_ids.shape[1], 6)
-
-    class _Prepare(nn.Module):
-        def __init__(self, shard):
-            super().__init__()
-            self.shard = shard
-
-        def forward(self, hidden, rope, combined):
-            local_len = hidden.shape[0] // 2 if self.shard else hidden.shape[0]
-            return hidden[:local_len], rope[:local_len], combined[:local_len]
-
-    class _Gather(nn.Module):
-        def __init__(self, remote_rows, replace):
-            super().__init__()
-            self.remote_rows = remote_rows
-            self.replace = replace
-            self.input_width = None
-
-        def forward(self, tensor):
-            self.input_width = tensor.shape[-1]
-            if self.replace:
-                return self.remote_rows
-            return torch.cat((tensor, self.remote_rows), dim=0)
-
-    class _FinalLayer(nn.Module):
-        def forward(self, hidden, *, t_emb, inverse_indices):
-            del t_emb
-            inverse = inverse_indices.to(torch.float32).unsqueeze(-1)
-            return (
-                hidden[:, :2].to(torch.float32) + inverse,
-                hidden[:, 2:3].to(torch.float32) - inverse,
-            )
-
-    seq_len = 8
-    full_hidden = torch.arange(seq_len * 4, dtype=torch.bfloat16).reshape(seq_len, 4)
-    inverse_indices = torch.tensor([0, 1, 0, 1, 1, 0, 1, 0])
-    final_layer = _FinalLayer()
-    remote_video, remote_audio = final_layer(
-        full_hidden[seq_len // 2 :],
-        t_emb=torch.empty(0),
-        inverse_indices=inverse_indices[seq_len // 2 :],
-    )
-    if local_projection:
-        remote_rows = torch.cat((remote_video, remote_audio), dim=-1)
-    else:
-        remote_rows = full_hidden
-    gather = _Gather(remote_rows, replace=not local_projection)
-
-    model = object.__new__(h3.MiniMaxH3DiTModel)
-    nn.Module.__init__(model)
-    model.arch = h3.MiniMaxH3DiTArchConfig(
-        latents_dim=2,
-        audio_latents_dim=1,
-        patch_size=(1, 1, 1),
-    )
-    rope = _StaticRope()
-    model.rope = rope
-    model.sp_prepare = _Prepare(shard=True)
-    model.local_sp_prepare = _Prepare(shard=False)
-    model.sp_gather = gather
-    model.blocks = nn.ModuleList()
-    model.final_layer = final_layer
-    monkeypatch.setattr(
-        model,
-        "_embed",
-        lambda **kwargs: (
-            full_hidden[: seq_len // 2] if local_projection else full_hidden,
-            torch.tensor([[0.5]], dtype=torch.float32),
-        ),
-    )
-
-    def _local_span(*args, **kwargs):
-        del args, kwargs
-        return (0, seq_len // 2) if local_projection else (0, seq_len)
-
-    monkeypatch.setattr(
-        h3,
-        "_sequence_parallel_local_span",
-        _local_span,
-    )
-
-    img_pos = torch.tensor([0, 2, 4, 6])
-    audio_pos = torch.tensor([1, 3, 5, 7])
-    video, audio = model(
-        x=torch.zeros(1, seq_len, 2),
-        audio_x=torch.zeros(1, seq_len, 1),
-        img_position_ids=torch.zeros(1, seq_len, 3, dtype=torch.long),
-        unique_timesteps=torch.tensor([0.5]),
-        inverse_indices=inverse_indices,
-        update_mask=torch.ones(img_pos.numel()),
-        token_tags=torch.zeros(seq_len, dtype=torch.long),
-        prompt_embeds=torch.empty(0, 2),
-        img_pos_info={"position_ids": img_pos},
-        audio_pos_info={"position_ids": audio_pos},
-        text_pos_info={"position_ids": torch.empty(0, dtype=torch.long)},
-        img_pos_for_infer_output_info={"position_ids": img_pos},
-        packed_seq_params={
-            "cu_seqlens_q": torch.tensor([0, seq_len, seq_len], dtype=torch.int32),
-            "max_seqlen_q": seq_len,
-        },
-        refiner_packed_seq_params={
-            "cu_seqlens_q": torch.tensor([0, 0, 0], dtype=torch.int32),
-            "max_seqlen_q": 0,
-        },
-    )
-
-    expected_video, expected_audio = final_layer(
-        full_hidden,
-        t_emb=torch.empty(0),
-        inverse_indices=inverse_indices,
-    )
-    torch.testing.assert_close(video, expected_video.index_select(0, img_pos))
-    torch.testing.assert_close(audio, expected_audio.index_select(0, audio_pos))
-    expected_gather_width = 3 if local_projection else full_hidden.shape[-1]
-    expected_rope_rows = seq_len // 2 if local_projection else seq_len
-    assert gather.input_width == expected_gather_width
-    assert rope.input_rows == expected_rope_rows
-
-
@pytest.mark.parametrize(
    ("tp_size", "message"),
    [
@@ -191,10 +191,7 @@ def test_model_load_weights_transforms_before_calling_vllm_loader():
    ]


-def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch):
-    from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
-        get_direct_mmap_adapter,
-    )
+def test_unquantized_qkv_declares_equivalent_mmap_transform(monkeypatch):
    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3

    monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
@@ -212,26 +209,10 @@ def test_loader_adapter_declares_equivalent_direct_mmap_transform(monkeypatch):
        quant_config=None,
        prefix="blocks.0.attn",
    )
-    transformer = object.__new__(h3.MiniMaxH3DiTModel)
-    nn.Module.__init__(transformer)
-    transformer.blocks = nn.ModuleList([nn.Module()])
-    transformer.blocks[0].attn = attention
-    pipeline = nn.Module()
-    pipeline.transformer = transformer
-
-    adapter = get_direct_mmap_adapter(pipeline)
-    assert adapter is not None
-    policy = adapter.policy_for(
-        "transformer.blocks.0.attn.qkv_proj.weight",
-        attention.qkv_proj.weight,
-    )
-    assert policy is not None
-    assert policy.allow_custom_loader
-    assert policy.transform is not None
+    transform = attention.qkv_proj.weight.mmap_weight_transform
    checkpoint_weight = torch.arange(6, dtype=torch.float32).reshape(6, 1)

-    assert policy.transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]
-    assert not hasattr(attention.qkv_proj.weight, "mmap_weight_transform")
+    assert transform(checkpoint_weight)[:, 0].tolist() == [0, 3, 1, 4, 2, 5]


 def test_pipeline_resolves_transformer_component_quant_config():
@@ -254,18 +235,3 @@ def test_pipeline_resolves_transformer_component_quant_config():
    assert transformer_config.ignored_layers == ignored_layers
    assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
    assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
-
-
-def test_pipeline_strips_prequantized_text_encoder_config():
-    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import _resolve_minimax_h3_text_encoder_quant_config
-    from vllm_omni.quantization import ComponentQuantizationConfig
-
-    prequantized = Mock()
-    prequantized.get_name.return_value = "modelopt"
-    online_fp8 = Mock()
-    online_fp8.get_name.return_value = "fp8"
-    component_config = ComponentQuantizationConfig({"text_encoder": prequantized})
-
-    assert _resolve_minimax_h3_text_encoder_quant_config(online_fp8) is online_fp8
-    assert _resolve_minimax_h3_text_encoder_quant_config(prequantized) is None
-    assert _resolve_minimax_h3_text_encoder_quant_config(component_config) is None
@@ -142,6 +142,12 @@ def _generate_joint_output(omni, config: QualityTestConfig):

    first = outputs[0]
    worker_peak_memory_mb = float(getattr(first, "peak_memory_mb", 0.0) or 0.0)
+    if hasattr(first, "request_output") and isinstance(first.request_output, list):
+        first = first.request_output[0]
+        worker_peak_memory_mb = max(
+            worker_peak_memory_mb,
+            float(getattr(first, "peak_memory_mb", 0.0) or 0.0),
+        )
    if not hasattr(first, "images") or not first.images:
        raise ValueError("Could not extract video frames from H3 output")
    frames = first.images[0]
@@ -1,104 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-from vllm_omni.platforms import current_omni_platform
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
-
-
-@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
-@pytest.mark.parametrize("num_q_heads,num_kv_heads,seq_len", [(8, 2, 16), (64, 8, 1)])
-def test_qwen3_vl_npu_rope_matches_bnsd_reference(
-    num_q_heads: int,
-    num_kv_heads: int,
-    seq_len: int,
-) -> None:
-    from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu
-
-    device = torch.device("npu")
-    dtype = torch.bfloat16
-    q = torch.randn(1, num_q_heads, seq_len, 128, device=device, dtype=dtype)
-    k = torch.randn(1, num_kv_heads, seq_len, 128, device=device, dtype=dtype)
-    cos = torch.randn(1, seq_len, 128, device=device, dtype=dtype)
-    sin = torch.randn(1, seq_len, 128, device=device, dtype=dtype)
-
-    q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin)
-    cos_bnsd = cos.unsqueeze(1)
-    sin_bnsd = sin.unsqueeze(1)
-
-    def rotate_half(x: torch.Tensor) -> torch.Tensor:
-        x1, x2 = x.chunk(2, dim=-1)
-        return torch.cat((-x2, x1), dim=-1)
-
-    torch.testing.assert_close(
-        q_out,
-        q * cos_bnsd + rotate_half(q) * sin_bnsd,
-        atol=2e-2,
-        rtol=2e-2,
-    )
-    torch.testing.assert_close(
-        k_out,
-        k * cos_bnsd + rotate_half(k) * sin_bnsd,
-        atol=2e-2,
-        rtol=2e-2,
-    )
-
-
-@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
-def test_qwen3_vl_npu_rope_preserves_batched_mrope_positions() -> None:
-    """Keep per-sample Qwen3-VL M-RoPE frequencies on the BNSD NPU path."""
-    from types import SimpleNamespace
-
-    from vllm_omni.diffusion.models.minimax_h3.encoder import (
-        MiniMaxH3Qwen3VLTextRotaryEmbedding,
-    )
-    from vllm_omni.platforms.npu.models.minimax_h3 import _apply_rotary_pos_emb_npu
-
-    device = torch.device("npu")
-    dtype = torch.bfloat16
-    batch_size, seq_len, head_dim = 2, 16, 128
-    config = SimpleNamespace(
-        hidden_size=head_dim,
-        num_attention_heads=1,
-        rope_parameters={
-            "rope_theta": 1_000_000.0,
-            "mrope_section": [24, 20, 20],
-        },
-    )
-    rope = MiniMaxH3Qwen3VLTextRotaryEmbedding(config).to(device)
-    positions = torch.arange(seq_len, device=device)
-    position_ids = torch.stack(
-        (
-            torch.stack((positions, positions + 7)),
-            torch.stack((positions * 2, positions * 2 + 3)),
-            torch.stack((positions * 3, positions * 3 + 5)),
-        )
-    )
-    q = torch.randn(batch_size, 8, seq_len, head_dim, device=device, dtype=dtype)
-    k = torch.randn(batch_size, 2, seq_len, head_dim, device=device, dtype=dtype)
-    cos, sin = rope(q, position_ids)
-
-    q_out, k_out = _apply_rotary_pos_emb_npu(q, k, cos, sin)
-    cos_bnsd = cos.unsqueeze(1)
-    sin_bnsd = sin.unsqueeze(1)
-
-    def rotate_half(x: torch.Tensor) -> torch.Tensor:
-        x1, x2 = x.chunk(2, dim=-1)
-        return torch.cat((-x2, x1), dim=-1)
-
-    assert not torch.equal(cos[0], cos[1])
-    torch.testing.assert_close(
-        q_out,
-        q * cos_bnsd + rotate_half(q) * sin_bnsd,
-        atol=2e-2,
-        rtol=2e-2,
-    )
-    torch.testing.assert_close(
-        k_out,
-        k * cos_bnsd + rotate_half(k) * sin_bnsd,
-        atol=2e-2,
-        rtol=2e-2,
-    )
@@ -1,46 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-import torch.nn.functional as F
-
-from vllm_omni.platforms import current_omni_platform
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
-
-
-@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
-@pytest.mark.parametrize("shape", [(1, 16, 256), (2, 1, 256)])
-def test_qwen3_vl_npu_swiglu_matches_split_reference(
-    shape: tuple[int, int, int],
-) -> None:
-    from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed
-
-    packed = torch.randn(*shape, device="npu", dtype=torch.bfloat16)
-    gate, up = packed.chunk(2, dim=-1)
-
-    torch.testing.assert_close(
-        npu_swiglu_from_packed(packed),
-        F.silu(gate) * up,
-        atol=2e-2,
-        rtol=2e-2,
-    )
-
-
-@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU")
-def test_qwen3_vl_npu_swiglu_matches_two_projection_reference() -> None:
-    from vllm_omni.platforms.npu.models.minimax_h3 import npu_swiglu_from_packed
-
-    x = torch.randn(1, 16, 128, device="npu", dtype=torch.bfloat16)
-    weight = torch.randn(256, 128, device="npu", dtype=torch.bfloat16)
-    packed = F.linear(x, weight)
-    gate = F.linear(x, weight[:128])
-    up = F.linear(x, weight[128:])
-
-    torch.testing.assert_close(
-        npu_swiglu_from_packed(packed),
-        F.silu(gate) * up,
-        atol=2e-2,
-        rtol=2e-2,
-    )
@@ -8,28 +8,15 @@ import torch
 from PIL import Image
 from torch import nn

-from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
+from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
    Wan22I2VPipeline,
-    get_wan22_i2v_post_process_func,
    get_wan22_i2v_pre_process_func,
 )
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]


-def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
-    video = torch.zeros(1, 4, 1, 2, 2)
-
-    output = get_wan22_i2v_post_process_func(SimpleNamespace())(
-        video,
-        sampling_params=SimpleNamespace(output_type="latent"),
-    )
-
-    assert output is video
-
-
 def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
    pipeline = object.__new__(Wan22I2VPipeline)
    nn.Module.__init__(pipeline)
@@ -44,29 +31,6 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
    return pipeline


-def _make_i2v_sampling(**overrides):
-    values = {
-        "height": 16,
-        "width": 16,
-        "num_frames": 5,
-        "num_inference_steps": 1,
-        "guidance_scale_provided": True,
-        "guidance_scale": 1.0,
-        "guidance_scale_2": None,
-        "guidance_scale_2_provided": False,
-        "boundary_ratio": None,
-        "generator": None,
-        "seed": None,
-        "num_outputs_per_prompt": 1,
-        "max_sequence_length": 8,
-        "latents": None,
-        "output_type": "latent",
-        "extra_args": {},
-    }
-    values.update(overrides)
-    return SimpleNamespace(**values)
-
-
 def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
    request = SimpleNamespace(
@@ -89,88 +53,6 @@ def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
        preprocess(missing_image)


-def _make_i2v_preprocess_request(last_image):
-    return SimpleNamespace(
-        prompt={
-            "prompt": "p",
-            "multi_modal_data": {
-                "image": Image.new("RGB", (320, 160), "red"),
-                "last_image": last_image,
-            },
-        },
-        sampling_params=SimpleNamespace(height=16, width=16),
-    )
-
-
-def test_i2v_preprocess_treats_empty_last_image_list_as_absent() -> None:
-    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
-
-    result = preprocess(_make_i2v_preprocess_request([]))
-
-    assert result.prompt["multi_modal_data"]["last_image"] is None
-    assert result.batch_compatibility_key == ("wan22_i2v_last_image", False)
-
-
-def test_i2v_preprocess_unwraps_single_last_image_list() -> None:
-    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
-    last_image = Image.new("RGB", (16, 16), "blue")
-
-    result = preprocess(_make_i2v_preprocess_request([last_image]))
-
-    assert result.prompt["multi_modal_data"]["last_image"] is last_image
-    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
-
-
-@pytest.mark.parametrize(
-    "last_image",
-    [
-        Image.new("RGB", (16, 16), "blue"),
-        torch.zeros(3, 16, 16),
-    ],
-)
-def test_i2v_preprocess_preserves_supported_last_image(last_image) -> None:
-    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
-
-    result = preprocess(_make_i2v_preprocess_request(last_image))
-
-    assert result.prompt["multi_modal_data"]["last_image"] is last_image
-    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
-
-
-def test_i2v_preprocess_loads_last_image_path(tmp_path) -> None:
-    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
-    path = tmp_path / "last.png"
-    Image.new("RGB", (16, 16), "blue").save(path)
-
-    result = preprocess(_make_i2v_preprocess_request(str(path)))
-
-    last_image = result.prompt["multi_modal_data"]["last_image"]
-    assert isinstance(last_image, Image.Image)
-    assert last_image.mode == "RGB"
-    assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
-
-
-def test_i2v_preprocess_rejects_multiple_last_images() -> None:
-    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
-
-    with pytest.raises(ValueError, match="at most one last_image"):
-        preprocess(
-            _make_i2v_preprocess_request(
-                [
-                    Image.new("RGB", (16, 16), "blue"),
-                    Image.new("RGB", (16, 16), "green"),
-                ]
-            )
-        )
-
-
-def test_i2v_preprocess_rejects_unsupported_last_image_type() -> None:
-    preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
-
-    with pytest.raises(TypeError, match="Unsupported last_image format"):
-        preprocess(_make_i2v_preprocess_request({"not": "an image"}))
-
-
 def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None:
    pipeline = _make_i2v_pipeline(expand_timesteps=True)
    latents = torch.zeros(1, 4, 2, 4, 4)
@@ -244,169 +126,3 @@ def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> N
    assert first_frame_mask.shape == (1, 1, 2, 2, 2)
    assert first_frame_mask[:, :, 0].sum() == 0
    assert first_frame_mask[:, :, 1].sum() == 4
-
-
-def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
-    pipeline = _make_i2v_pipeline(expand_timesteps=True)
-    generators = [
-        torch.Generator(device="cpu").manual_seed(1),
-        torch.Generator(device="cpu").manual_seed(2),
-    ]
-
-    latents, condition, first_frame_mask = pipeline.prepare_latents(
-        image=torch.zeros(2, 3, 16, 16),
-        batch_size=2,
-        num_channels_latents=4,
-        height=16,
-        width=16,
-        num_frames=5,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-        generator=generators,
-    )
-
-    assert latents.shape == (2, 4, 2, 2, 2)
-    assert condition.shape == (2, 4, 1, 2, 2)
-    assert first_frame_mask.shape == (2, 1, 2, 2, 2)
-
-
-def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
-        SimpleNamespace(is_available=lambda: False),
-    )
-    pipeline = _make_i2v_pipeline(expand_timesteps=True)
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 5.0
-    pipeline.boundary_ratio = 0.875
-    pipeline.has_image_encoder = False
-    pipeline._guidance_scale = None
-    pipeline._guidance_scale_2 = None
-    pipeline._num_timesteps = None
-    pipeline._current_timestep = None
-    pipeline.check_inputs = lambda **kwargs: None
-    prepare_call = {}
-
-    def _fake_encode_prompt(**kwargs):
-        batch_size = len(kwargs["prompt"])
-        n = batch_size * kwargs["num_videos_per_prompt"]
-        return torch.zeros(n, 2, 3), None
-
-    def _fake_prepare_latents(**kwargs):
-        prepare_call.update(kwargs)
-        batch_size = kwargs["batch_size"]
-        return (
-            kwargs["latents"],
-            torch.zeros(batch_size, 4, 1, 2, 2),
-            torch.ones(batch_size, 1, 2, 2, 2),
-        )
-
-    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
-    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
-    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
-
-    gen_a = torch.Generator(device="cpu").manual_seed(1)
-    gen_b = torch.Generator(device="cpu").manual_seed(2)
-    latents_a = torch.zeros(2, 4, 2, 2, 2)
-    latents_b = torch.ones(2, 4, 2, 2, 2)
-    image_a = torch.zeros(1, 3, 16, 16)
-    image_b = torch.ones(1, 3, 16, 16)
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={"prompt": "first", "multi_modal_data": {"image": image_a}},
-                sampling_params=_make_i2v_sampling(
-                    generator=gen_a,
-                    latents=latents_a,
-                    num_outputs_per_prompt=2,
-                ),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={"prompt": "second", "multi_modal_data": {"image": image_b}},
-                sampling_params=_make_i2v_sampling(
-                    generator=gen_b,
-                    latents=latents_b,
-                    num_outputs_per_prompt=2,
-                ),
-            ),
-        ]
-    )
-
-    outputs = pipeline.forward(batch)
-
-    assert prepare_call["batch_size"] == 4
-    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
-    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
-    torch.testing.assert_close(
-        prepare_call["image"],
-        torch.cat([image_a, image_a, image_b, image_b]),
-    )
-    assert len(outputs) == 2
-    torch.testing.assert_close(outputs[0].output, latents_a)
-    torch.testing.assert_close(outputs[1].output, latents_b)
-
-
-def test_i2v_forward_rejects_mismatched_tensor_condition_shapes(monkeypatch) -> None:
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
-        SimpleNamespace(is_available=lambda: False),
-    )
-    pipeline = _make_i2v_pipeline(expand_timesteps=True)
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=5.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 5.0
-    pipeline.boundary_ratio = 0.875
-    pipeline.has_image_encoder = False
-    pipeline._guidance_scale = None
-    pipeline._guidance_scale_2 = None
-    pipeline._num_timesteps = None
-    pipeline._current_timestep = None
-    pipeline.check_inputs = lambda **kwargs: None
-    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={"prompt": "first", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 16)}},
-                sampling_params=_make_i2v_sampling(),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={"prompt": "second", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 8)}},
-                sampling_params=_make_i2v_sampling(),
-            ),
-        ]
-    )
-
-    with pytest.raises(ValueError, match="image condition"):
-        pipeline.forward(batch)
-
-
-def test_i2v_forward_rejects_mixed_last_image_presence() -> None:
-    pipeline = _make_i2v_pipeline(expand_timesteps=True)
-    image = torch.zeros(1, 3, 16, 16)
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={
-                    "prompt": "first",
-                    "multi_modal_data": {"image": image, "last_image": image},
-                },
-                sampling_params=_make_i2v_sampling(),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={"prompt": "second", "multi_modal_data": {"image": image}},
-                sampling_params=_make_i2v_sampling(),
-            ),
-        ]
-    )
-
-    with pytest.raises(ValueError, match="mix of provided and missing last_image"):
-        pipeline.forward(batch)
@@ -93,29 +93,6 @@ def _make_pipeline() -> Wan22Pipeline:
    return pipeline


-def _make_sampling(**overrides):
-    values = {
-        "height": None,
-        "width": None,
-        "num_frames": 1,
-        "num_inference_steps": 2,
-        "guidance_scale_provided": True,
-        "guidance_scale": 1.0,
-        "guidance_scale_2": None,
-        "guidance_scale_2_provided": False,
-        "boundary_ratio": None,
-        "generator": None,
-        "seed": None,
-        "num_outputs_per_prompt": 1,
-        "max_sequence_length": 32,
-        "latents": None,
-        "output_type": "latent",
-        "extra_args": {},
-    }
-    values.update(overrides)
-    return SimpleNamespace(**values)
-
-
@pytest.mark.parametrize(
    ("sampling_params_kwargs", "expected_low", "expected_high"),
    [
@@ -152,10 +129,9 @@ def test_forward_delegates_denoising_to_diffuse(
    )
    batch = DiffusionRequestBatch(requests=[mock_req])

-    outputs = pipeline.forward(batch)
+    output = pipeline.forward(batch)

-    assert len(outputs) == 1
-    assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8)))
+    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"] == expected_low
@@ -166,134 +142,6 @@ def test_forward_delegates_denoising_to_diffuse(
    assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]


-def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
-    pipeline = _make_pipeline()
-    encode_call = {}
-    prepare_call = {}
-
-    def _fake_encode_prompt(**kwargs):
-        encode_call.update(kwargs)
-        batch_size = len(kwargs["prompt"])
-        n = batch_size * kwargs["num_videos_per_prompt"]
-        return torch.arange(n, dtype=torch.float32).view(n, 1, 1), torch.zeros(n, 1, 1)
-
-    def _fake_prepare_latents(**kwargs):
-        prepare_call.update(kwargs)
-        return kwargs["latents"]
-
-    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
-    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
-    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
-
-    gen_a = torch.Generator(device="cpu").manual_seed(1)
-    gen_b = torch.Generator(device="cpu").manual_seed(2)
-    latents_a = torch.zeros(2, 4, 1, 2, 2)
-    latents_b = torch.ones(2, 4, 1, 2, 2)
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={"prompt": "first", "negative_prompt": "bad first"},
-                sampling_params=_make_sampling(
-                    generator=gen_a,
-                    latents=latents_a,
-                    num_outputs_per_prompt=2,
-                ),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={"prompt": "second", "negative_prompt": "bad second"},
-                sampling_params=_make_sampling(
-                    generator=gen_b,
-                    latents=latents_b,
-                    num_outputs_per_prompt=2,
-                ),
-            ),
-        ]
-    )
-
-    outputs = pipeline.forward(batch)
-
-    assert encode_call["prompt"] == ["first", "second"]
-    assert encode_call["negative_prompt"] == ["bad first", "bad second"]
-    assert prepare_call["batch_size"] == 4
-    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
-    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
-    assert len(outputs) == 2
-    torch.testing.assert_close(outputs[0].output, latents_a)
-    torch.testing.assert_close(outputs[1].output, latents_b)
-
-
-def test_forward_batches_precomputed_prompt_embeddings() -> None:
-    pipeline = _make_pipeline()
-    diffuse_call = {}
-    pipeline.encode_prompt = lambda **kwargs: pytest.fail("text encoder must not run")  # type: ignore[method-assign]
-    pipeline.prepare_latents = lambda **kwargs: torch.zeros(kwargs["batch_size"], 4, 1, 2, 2)  # type: ignore[method-assign]
-
-    def _fake_diffuse(**kwargs):
-        diffuse_call.update(kwargs)
-        return kwargs["latents"]
-
-    pipeline.diffuse = _fake_diffuse  # type: ignore[method-assign]
-    embeds_a = torch.zeros(3, 4)
-    embeds_b = torch.ones(3, 4)
-    negative_a = torch.full((3, 4), 2.0)
-    negative_b = torch.full((3, 4), 3.0)
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={"prompt_embeds": embeds_a, "negative_prompt_embeds": negative_a},
-                sampling_params=_make_sampling(guidance_scale=4.0),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={"prompt_embeds": embeds_b, "negative_prompt_embeds": negative_b},
-                sampling_params=_make_sampling(guidance_scale=4.0),
-            ),
-        ]
-    )
-
-    outputs = pipeline.forward(batch)
-
-    torch.testing.assert_close(diffuse_call["prompt_embeds"], torch.stack([embeds_a, embeds_b]))
-    torch.testing.assert_close(diffuse_call["negative_prompt_embeds"], torch.stack([negative_a, negative_b]))
-    assert len(outputs) == 2
-
-
-def test_prepare_latents_with_request_generators_matches_single_generation() -> None:
-    pipeline = _make_pipeline()
-    kwargs = {
-        "num_channels_latents": 4,
-        "height": 16,
-        "width": 16,
-        "num_frames": 5,
-        "dtype": torch.float32,
-        "device": torch.device("cpu"),
-    }
-
-    batched = Wan22Pipeline.prepare_latents(
-        pipeline,
-        batch_size=2,
-        generator=[torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)],
-        **kwargs,
-    )
-    singles = torch.cat(
-        [
-            Wan22Pipeline.prepare_latents(
-                pipeline,
-                batch_size=1,
-                generator=torch.Generator().manual_seed(seed),
-                **kwargs,
-            )
-            for seed in (1, 2)
-        ]
-    )
-
-    torch.testing.assert_close(batched, singles)
-    assert not torch.equal(batched[0], batched[1])
-
-
 def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
    pipeline = _make_pipeline()
    latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32)
@@ -10,7 +10,6 @@ import torch
 import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
    create_transformer_from_config,
-    get_wan22_post_process_func,
    load_transformer_config,
    retrieve_latents,
 )
@@ -27,17 +26,6 @@ class _LatentDist:
        return torch.tensor([2.0])


-def test_wan22_postprocess_honors_request_output_type() -> None:
-    video = torch.zeros(1, 4, 1, 2, 2)
-
-    output = get_wan22_post_process_func(SimpleNamespace())(
-        video,
-        sampling_params=SimpleNamespace(output_type="latent"),
-    )
-
-    assert output is video
-
-
 def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None:
    generator = torch.Generator(device="cpu")

@@ -1,307 +1,14 @@
-from types import SimpleNamespace
 from unittest.mock import MagicMock, patch

-import numpy as np
-import PIL.Image
 import pytest
 import torch
 import torch.nn as nn

-from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
-    Wan22S2VPipeline,
-    _make_clip_generators,
-)
 from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


-def _make_s2v_sampling(**overrides):
-    values = {
-        "height": 16,
-        "width": 16,
-        "num_frames": 8,
-        "num_inference_steps": 1,
-        "guidance_scale_provided": True,
-        "guidance_scale": 1.0,
-        "generator": None,
-        "seed": None,
-        "num_outputs_per_prompt": 1,
-        "max_sequence_length": 8,
-        "latents": None,
-        "output_type": "latent",
-        "extra_args": {},
-    }
-    values.update(overrides)
-    return SimpleNamespace(**values)
-
-
-def _make_s2v_validation_pipeline() -> Wan22S2VPipeline:
-    pipeline = object.__new__(Wan22S2VPipeline)
-    nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
-    pipeline.transformer = SimpleNamespace(dtype=torch.float32)
-    pipeline.vae_scale_factor_spatial = 8
-    pipeline.resolution_divisor = 16
-    pipeline.motion_frames = 7
-    pipeline.drop_first_motion = True
-    pipeline._DEFAULT_INFER_FRAMES = 8
-    pipeline._guidance_scale = None
-    pipeline.check_inputs = lambda *args, **kwargs: None
-    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
-    return pipeline
-
-
-def test_s2v_predict_noise_keeps_batch_dimension() -> None:
-    pipeline = object.__new__(Wan22S2VPipeline)
-    nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
-    transformer = MagicMock()
-    transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)])
-    transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),)
-    pipeline.transformer = transformer
-
-    result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2))
-
-    assert result.shape == (1, 16, 2, 2, 2)
-
-
-def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None:
-    explicit_generator = torch.Generator(device="cpu").manual_seed(999)
-
-    generators = _make_clip_generators(
-        seeds=[1234, 5678, None, None],
-        request_generators=[
-            torch.Generator(device="cpu").manual_seed(1234),
-            torch.Generator(device="cpu").manual_seed(5678),
-            explicit_generator,
-            None,
-        ],
-        clip_index=2,
-        device=torch.device("cpu"),
-    )
-
-    assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2]
-    assert generators[2] is explicit_generator
-
-
-def test_s2v_forward_rejects_different_num_repeat() -> None:
-    pipeline = _make_s2v_validation_pipeline()
-    pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16)  # type: ignore[method-assign]
-    image = PIL.Image.new("RGB", (16, 16))
-    audio = np.zeros(16000, dtype=np.float32)
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={
-                    "prompt": "first",
-                    "multi_modal_data": {"image": image, "audio": audio},
-                    "additional_information": {"num_repeat": 1},
-                },
-                sampling_params=_make_s2v_sampling(),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={
-                    "prompt": "second",
-                    "multi_modal_data": {"image": image, "audio": audio},
-                    "additional_information": {"num_repeat": 2},
-                },
-                sampling_params=_make_s2v_sampling(),
-            ),
-        ]
-    )
-
-    with pytest.raises(ValueError, match="num_repeat"):
-        pipeline.forward(batch)
-
-
-def test_s2v_forward_rejects_different_audio_time_lengths() -> None:
-    pipeline = _make_s2v_validation_pipeline()
-    pipeline.encode_audio = lambda audio, **kwargs: (  # type: ignore[method-assign]
-        torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24),
-        1,
-        16 if audio[0] == 0 else 24,
-    )
-    image = PIL.Image.new("RGB", (16, 16))
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={
-                    "prompt": "first",
-                    "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
-                },
-                sampling_params=_make_s2v_sampling(),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={
-                    "prompt": "second",
-                    "multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)},
-                },
-                sampling_params=_make_s2v_sampling(),
-            ),
-        ]
-    )
-
-    with pytest.raises(ValueError, match="audio embedding time lengths"):
-        pipeline.forward(batch)
-
-
-def test_s2v_forward_rejects_different_init_first_frame() -> None:
-    pipeline = _make_s2v_validation_pipeline()
-    image = PIL.Image.new("RGB", (16, 16))
-    audio = np.zeros(16000, dtype=np.float32)
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={
-                    "prompt": "first",
-                    "multi_modal_data": {"image": image, "audio": audio},
-                    "additional_information": {"init_first_frame": True},
-                },
-                sampling_params=_make_s2v_sampling(),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={
-                    "prompt": "second",
-                    "multi_modal_data": {"image": image, "audio": audio},
-                    "additional_information": {"init_first_frame": False},
-                },
-                sampling_params=_make_s2v_sampling(),
-            ),
-        ]
-    )
-
-    with pytest.raises(ValueError, match="init_first_frame"):
-        pipeline.forward(batch)
-
-
-def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
-    pipeline = _make_s2v_validation_pipeline()
-    pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16)  # type: ignore[method-assign]
-    image = PIL.Image.new("RGB", (16, 16))
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={
-                    "prompt": "first",
-                    "multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
-                },
-                sampling_params=_make_s2v_sampling(),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={
-                    "prompt": "second",
-                    "multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)},
-                },
-                sampling_params=_make_s2v_sampling(),
-            ),
-        ]
-    )
-
-    with pytest.raises(ValueError, match="raw audio shapes and sample rates"):
-        pipeline.forward(batch)
-
-
-def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
-    pipeline = object.__new__(Wan22S2VPipeline)
-    nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
-    pipeline.transformer = MagicMock()
-    pipeline.transformer.dtype = torch.float32
-    pipeline.transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.float32)])
-    pipeline.transformer.casual_audio_encoder = None
-    pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
-    pipeline.vae = MagicMock()
-    pipeline.vae.dtype = torch.float32
-    pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
-    pipeline.od_config = SimpleNamespace(
-        enable_cpu_offload=False,
-        parallel_config=SimpleNamespace(use_hsdp=False),
-    )
-    pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
-    pipeline.vae_scale_factor_spatial = 8
-    pipeline.resolution_divisor = 16
-    pipeline.motion_frames = 7
-    pipeline.drop_first_motion = True
-    pipeline._DEFAULT_INFER_FRAMES = 8
-    pipeline._guidance_scale = None
-    pipeline._num_timesteps = None
-    pipeline.check_inputs = lambda *args, **kwargs: None
-    pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
-    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
-    pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
-    pipeline.prepare_motion_latents = MagicMock(
-        side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
-    )
-    pipeline._denormalize_latents = lambda latents: latents
-    pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])
-
-    image = PIL.Image.new("RGB", (16, 16))
-    audio_a = np.zeros(16000, dtype=np.float32)
-    audio_b = np.ones(16000, dtype=np.float32)
-    latents_a = torch.zeros(2, 16, 2, 2, 2)
-    latents_b = torch.ones(2, 16, 2, 2, 2)
-    generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
-    generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)]
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={
-                    "prompt": "first",
-                    "negative_prompt": "negative first",
-                    "multi_modal_data": {"image": image, "audio": audio_a},
-                },
-                sampling_params=_make_s2v_sampling(
-                    num_outputs_per_prompt=2,
-                    generator=generators_a,
-                    latents=latents_a,
-                ),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={
-                    "prompt": "second",
-                    "negative_prompt": "negative second",
-                    "multi_modal_data": {"image": image, "audio": audio_b},
-                },
-                sampling_params=_make_s2v_sampling(
-                    num_outputs_per_prompt=2,
-                    generator=generators_b,
-                    latents=latents_b,
-                ),
-            ),
-        ]
-    )
-
-    with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
-        platform.is_available.return_value = False
-        outputs = pipeline.forward(batch)
-
-    assert len(outputs) == 2
-    assert outputs[0].output[0].shape[0] == 2
-    assert outputs[1].output[0].shape[0] == 2
-    assert outputs[0].output[1].shape == (16000,)
-    assert outputs[1].output[1].shape == (16000,)
-    np.testing.assert_array_equal(outputs[0].output[1], audio_a)
-    np.testing.assert_array_equal(outputs[1].output[1], audio_b)
-    assert np.shares_memory(outputs[0].output[1], audio_a)
-    assert np.shares_memory(outputs[1].output[1], audio_b)
-    torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b]))
-    assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b
-    assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"]
-    assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"]
-
-
 def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks():
    model = object.__new__(WanS2VTransformer3DModel)
    nn.Module.__init__(model)
@@ -553,37 +260,31 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
    pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean

    # -- Build request --
-    sampling_params = _make_s2v_sampling(
-        height=704,
-        width=1024,
-        num_frames=80,
-        num_inference_steps=5,
-        guidance_scale=4.5,
-        generator=torch.Generator(device="cpu").manual_seed(42),
-        max_sequence_length=512,
-    )
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+    sampling_params = MagicMock()
+    sampling_params.height = 704
+    sampling_params.width = 1024
+    sampling_params.num_frames = 80
+    sampling_params.num_inference_steps = 5
+    sampling_params.guidance_scale = 4.5
+    sampling_params.guidance_scale_provided = True
+    sampling_params.generator = None
+    sampling_params.seed = 42

    ref_image = PIL.Image.new("RGB", (1024, 704))
    audio_data = np.zeros(16000, dtype=np.float32)

-    req = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="test",
-                prompt={
-                    "prompt": "test prompt",
-                    "negative_prompt": "bad quality",
-                    "multi_modal_data": {"image": ref_image, "audio": audio_data},
-                    "additional_information": {
-                        "audio_path": audio_data,
-                        "pose_video": None,
-                        "init_first_frame": False,
-                    },
-                },
-                sampling_params=sampling_params,
-            )
-        ]
-    )
+    req = MagicMock(spec=OmniDiffusionRequest)
+    req.prompts = [
+        {
+            "prompt": "test prompt",
+            "negative_prompt": "bad quality",
+            "multi_modal_data": {"image": ref_image, "audio": audio_data},
+            "additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False},
+        }
+    ]
+    req.sampling_params = sampling_params

    # -- Mock methods that use complex internal state --
    pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80))
@@ -594,7 +295,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
    pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False)

    # Mock predict_noise_maybe_with_cfg from CFGParallelMixin
-    pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128))
+    pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128))
    pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline)

    # Mock platform methods
@@ -605,7 +306,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
        with patch(
            "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000)
        ):
-            result = Wan22S2VPipeline.forward(pipeline, req=req)[0]
+            result = Wan22S2VPipeline.forward(pipeline, req=req)

    # -- Assertions --
    # Text encoder was called
@@ -8,14 +8,13 @@ import torch
 from PIL import Image
 from torch import nn

-from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
+from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
    Wan22VACEPipeline,
    create_vace_transformer_from_config,
    get_wan22_vace_pre_process_func,
 )
 from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]

@@ -33,27 +32,6 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
    return pipeline


-def _make_vace_sampling(**overrides):
-    values = {
-        "height": 16,
-        "width": 16,
-        "num_frames": 5,
-        "num_inference_steps": 1,
-        "guidance_scale_provided": True,
-        "guidance_scale": 1.0,
-        "boundary_ratio": None,
-        "generator": None,
-        "seed": None,
-        "num_outputs_per_prompt": 1,
-        "max_sequence_length": 8,
-        "latents": None,
-        "output_type": "latent",
-        "extra_args": {},
-    }
-    values.update(overrides)
-    return SimpleNamespace(**values)
-
-
 def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
    preprocess = get_wan22_vace_pre_process_func(SimpleNamespace())
    ref = Image.new("RGB", (320, 160), "green")
@@ -160,189 +138,3 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
    assert calls[0]["positive_kwargs"]["vace_context"] is vace_context
    assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75
    torch.testing.assert_close(result, torch.ones_like(latents))
-
-
-def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
-        SimpleNamespace(is_available=lambda: False),
-    )
-    pipeline = _make_vace_pipeline()
-    pipeline.transformer.vace_patch_embedding = None
-    pipeline.transformer_2 = None
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 3.0
-    pipeline.boundary_ratio = None
-    pipeline._guidance_scale = None
-    pipeline._num_timesteps = None
-    pipeline._current_timestep = None
-    pipeline.check_inputs = lambda **kwargs: None
-    prepare_call = {}
-
-    def _fake_encode_prompt(**kwargs):
-        batch_size = len(kwargs["prompt"])
-        n = batch_size * kwargs["num_videos_per_prompt"]
-        return torch.zeros(n, 2, 3), None
-
-    def _fake_prepare_latents(**kwargs):
-        prepare_call.update(kwargs)
-        return kwargs["latents"]
-
-    pipeline.encode_prompt = _fake_encode_prompt  # type: ignore[method-assign]
-    pipeline.prepare_latents = _fake_prepare_latents  # type: ignore[method-assign]
-    pipeline.diffuse = lambda **kwargs: kwargs["latents"]  # type: ignore[method-assign]
-    gen_a = torch.Generator(device="cpu").manual_seed(1)
-    gen_b = torch.Generator(device="cpu").manual_seed(2)
-    latents_a = torch.zeros(2, 4, 2, 2, 2)
-    latents_b = torch.ones(2, 4, 2, 2, 2)
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={"prompt": "first"},
-                sampling_params=_make_vace_sampling(
-                    generator=gen_a,
-                    latents=latents_a,
-                    num_outputs_per_prompt=2,
-                ),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={"prompt": "second"},
-                sampling_params=_make_vace_sampling(
-                    generator=gen_b,
-                    latents=latents_b,
-                    num_outputs_per_prompt=2,
-                ),
-            ),
-        ]
-    )
-
-    outputs = pipeline.forward(batch)
-
-    assert prepare_call["batch_size"] == 4
-    assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
-    torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
-    assert len(outputs) == 2
-    torch.testing.assert_close(outputs[0].output, latents_a)
-    torch.testing.assert_close(outputs[1].output, latents_b)
-
-
-def test_vace_forward_rejects_mismatched_condition_shapes(monkeypatch) -> None:
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
-        SimpleNamespace(is_available=lambda: False),
-    )
-    pipeline = _make_vace_pipeline()
-    pipeline.transformer.vace_patch_embedding = object()
-    pipeline.transformer_2 = None
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 3.0
-    pipeline.boundary_ratio = None
-    pipeline._guidance_scale = None
-    pipeline._num_timesteps = None
-    pipeline._current_timestep = None
-    pipeline.check_inputs = lambda **kwargs: None
-    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
-
-    def _fake_preprocess(video, **kwargs):
-        del kwargs
-        mask = torch.ones_like(video)
-        return video, mask, [[]]
-
-    pipeline.preprocess_conditions = _fake_preprocess  # type: ignore[method-assign]
-    pipeline.prepare_video_latents = (  # type: ignore[method-assign]
-        lambda video, mask, refs, generator, device: torch.zeros(1, 8, video.shape[2], 2, 2)
-    )
-    pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, mask.shape[2], 2, 2)  # type: ignore[method-assign]
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={
-                    "prompt": "first",
-                    "additional_information": {"source_video": torch.zeros(1, 3, 5, 16, 16)},
-                },
-                sampling_params=_make_vace_sampling(),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={
-                    "prompt": "second",
-                    "additional_information": {"source_video": torch.zeros(1, 3, 3, 16, 16)},
-                },
-                sampling_params=_make_vace_sampling(),
-            ),
-        ]
-    )
-
-    with pytest.raises(ValueError, match="VACE condition"):
-        pipeline.forward(batch)
-
-
-def test_vace_forward_rejects_mismatched_reference_image_counts(monkeypatch) -> None:
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
-        SimpleNamespace(is_available=lambda: False),
-    )
-    pipeline = _make_vace_pipeline()
-    pipeline.transformer.vace_patch_embedding = object()
-    pipeline.transformer_2 = None
-    pipeline.scheduler = StubScheduler([9])
-    pipeline.od_config = SimpleNamespace(flow_shift=3.0)
-    pipeline._sample_solver = "unipc"
-    pipeline._flow_shift = 3.0
-    pipeline.boundary_ratio = None
-    pipeline._guidance_scale = None
-    pipeline._num_timesteps = None
-    pipeline._current_timestep = None
-    pipeline.check_inputs = lambda **kwargs: None
-    pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None)  # type: ignore[method-assign]
-
-    def _fake_preprocess(video, mask, reference_images, **kwargs):
-        del kwargs
-        return video, mask, [reference_images]
-
-    pipeline.preprocess_conditions = _fake_preprocess  # type: ignore[method-assign]
-    pipeline.prepare_video_latents = (  # type: ignore[method-assign]
-        lambda video, mask, refs, generator, device: torch.zeros(1, 8, 5, 2, 2)
-    )
-    pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, 5, 2, 2)  # type: ignore[method-assign]
-    video = torch.zeros(1, 3, 5, 16, 16)
-    mask = torch.ones_like(video)
-    reference = torch.zeros(3, 16, 16)
-    batch = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="a",
-                prompt={
-                    "prompt": "first",
-                    "additional_information": {
-                        "source_video": video,
-                        "mask": mask,
-                        "reference_images": [reference],
-                    },
-                },
-                sampling_params=_make_vace_sampling(),
-            ),
-            SimpleNamespace(
-                request_id="b",
-                prompt={
-                    "prompt": "second",
-                    "additional_information": {
-                        "source_video": video,
-                        "mask": mask,
-                        "reference_images": [reference, reference],
-                    },
-                },
-                sampling_params=_make_vace_sampling(),
-            ),
-        ]
-    )
-
-    with pytest.raises(ValueError, match="same number of reference images"):
-        pipeline.forward(batch)
@@ -18,11 +18,6 @@ from torch import nn
 from torch.distributed.tensor import DeviceMesh, DTensor, Replicate

 import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dist_backend_module
-from vllm_omni.diffusion.model_loader.host_weight_plan import (
-    HostWeightPlan,
-    TensorBinding,
-    build_checkpoint_mmap_plan,
-)
 from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
 from vllm_omni.diffusion.offloader.block_discovery import (
    get_blocks_attr_names,
@@ -38,6 +33,9 @@ from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
 from vllm_omni.diffusion.offloader.offload_plan import (
    OffloadPlan,
    get_offload_plan,
+    has_online_quantization,
+    should_select_dlo_mmap,
+    supports_mmap_loading,
 )
 from vllm_omni.platforms import current_omni_platform

@@ -150,41 +148,6 @@ class TestDistributedLayerwiseOffloadHook:
        assert shard.numel() == 4
        assert torch.equal(shard, _make_values(10.0))

-    def test_prefetch_preserves_transposed_weight_stride(
-        self,
-        dist_group,
-        patched_offload_runtime,
-    ):
-        """Online-FP8 Cutlass weights must retain their transposed layout."""
-
-        class StridedBlock(nn.Module):
-            def __init__(self, start: float):
-                super().__init__()
-                base = torch.arange(start, start + 12).reshape(3, 4)
-                self.weight = nn.Parameter(base.t(), requires_grad=False)
-
-        current_block = StridedBlock(1.0)
-        next_block = StridedBlock(20.0)
-        expected = next_block.weight.detach().clone()
-        expected_stride = next_block.weight.stride()
-        assert expected_stride == (1, 4)
-
-        hook = DistributedLayerwiseOffloadHook(
-            next_block=next_block,
-            device=torch.device("cpu"),
-            dp_group=None,
-            dp_size=1,
-            rank=0,
-            copy_stream=DummyStream(),
-            comm_stream=DummyStream(),
-            pin_memory=False,
-        )
-        hook.initialize_hook(current_block)
-        hook.prefetch_layer(slot=0, non_blocking=False)
-
-        assert next_block.weight.stride() == expected_stride
-        assert torch.equal(next_block.weight, expected)
-
    def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
        """DTensor wrapper should be preserved through prefetch/offload cycle."""
        current_block = TinyBlock(_make_values(1.0))
@@ -362,9 +325,8 @@ class TestDistributedLayerwiseOffloadHook:
        block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
        next_block = nn.Module()
        next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
-
-        def transform(tensor):
-            return tensor.flip(0)
+        next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
+        next_block.weight.mmap_weight_transform_pending = True

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
@@ -373,17 +335,17 @@ class TestDistributedLayerwiseOffloadHook:
            dp_size=2,
            rank=0,
            pin_memory=False,
-            tensor_transforms={id(next_block.weight): transform},
        )
        hook.initialize_hook(block)

        assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))

-    def test_transform_is_not_applied_without_a_loader_plan(self, dist_group, patched_offload_runtime):
+    def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime):
        block = nn.Module()
        block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
        next_block = nn.Module()
        next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
+        next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
@@ -402,9 +364,8 @@ class TestDistributedLayerwiseOffloadHook:
        next_block = nn.Linear(2, 2, bias=False)
        next_block.weight.data.copy_(torch.arange(4, dtype=torch.float32).view(2, 2))
        source_storage = next_block.weight.untyped_storage().data_ptr()
-
-        def transform(tensor):
-            return tensor.t()
+        next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
+        next_block.weight.mmap_weight_transform_pending = True

        hook = DistributedLayerwiseOffloadHook(
            next_block=next_block,
@@ -414,7 +375,6 @@ class TestDistributedLayerwiseOffloadHook:
            rank=0,
            pin_memory=False,
            rank_local_mmap=True,
-            tensor_transforms={id(next_block.weight): transform},
        )
        hook.initialize_hook(current_block)

@@ -431,9 +391,8 @@ class TestDistributedLayerwiseOffloadHook:

        assert torch.equal(
            next_block.weight,
-            torch.tensor([[0.0, 2.0], [1.0, 3.0]]),
+            torch.tensor([[2.0, 3.0], [0.0, 1.0]]),
        )
-        assert next_block.weight.stride() == (1, 2)
        assert torch.equal(source, torch.arange(4, dtype=torch.float32).view(2, 2))

    def test_rank_local_mmap_staging_is_bounded_by_largest_block(self, patched_offload_runtime):
@@ -509,24 +468,6 @@ class TestPinnedResidentLayerGroup:

        group.offload()

-    def test_load_preserves_noncontiguous_weight_stride(self, patched_offload_runtime):
-        block = nn.Module()
-        weight = torch.arange(6, dtype=torch.float32).reshape(2, 3).t()
-        expected = weight.detach().clone()
-        expected_stride = weight.stride()
-        block.weight = nn.Parameter(weight)
-        group = PinnedResidentLayerGroup(
-            [block],
-            device=torch.device("cpu"),
-            copy_stream=DummyStream(),
-            pin_memory=False,
-        )
-
-        group.load()
-        assert torch.equal(block.weight, expected)
-        assert block.weight.stride() == expected_stride
-        group.offload()
-
    def test_rank_local_mmap_resident_layers_retain_sources(self, patched_offload_runtime):
        blocks = [nn.Linear(2, 2, bias=False) for _ in range(3)]
        expected = []
@@ -563,6 +504,24 @@ class TestPinnedResidentLayerGroup:
        group.offload()
        assert all(block.weight.numel() == 0 for block in blocks)

+    def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
+        module = nn.Linear(2, 2, bias=False)
+
+        def transform(tensor):
+            return tensor.flip(0)
+
+        module.weight.mmap_weight_transform = transform
+
+        backend = object.__new__(DistributedLayerwiseOffloadBackend)
+        backend._remember_mmap_param_attrs(module)
+        module.to_empty(device="meta")
+
+        replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
+        backend._attach_mmap_param_attrs("weight", replacement, module.weight)
+
+        assert replacement.mmap_weight_transform is transform
+        assert replacement.mmap_weight_transform_pending is True
+
    def test_shared_allgather_output_is_narrowed_to_current_block(
        self, dist_group, patched_offload_runtime, monkeypatch
    ):
@@ -704,6 +663,7 @@ class TestMmapWeightLoading:
            OffloadConfig(
                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
                pin_cpu_memory=False,
+                model_path=str(tmp_path),
            ),
            torch.device("cpu"),
        )
@@ -711,81 +671,13 @@ class TestMmapWeightLoading:
            dits=[pipeline.transformer],
            dit_names=["transformer"],
        )
-        plan = HostWeightPlan(
-            backing_kind="checkpoint_mmap",
-            bindings={
-                name: TensorBinding(
-                    checkpoint_key=name,
-                    file_path=str(weight_file),
-                )
-                for name in weights
-            },
-        )

-        backend._load_weights_via_mmap(pipeline, modules, plan)
+        backend._load_weights_via_mmap(pipeline, modules)

        assert pipeline.transformer.post_load_calls == 1
        assert pipeline.transformer.time_embedder.weight.dtype == torch.float32
        assert pipeline.transformer.blocks[0].weight.dtype == torch.bfloat16

-    def test_allgather_flag_with_group_size_one_uses_rank_local_mmap(
-        self,
-        tmp_path,
-        patched_offload_runtime,
-    ):
-        class Transformer(nn.Module):
-            _layerwise_offload_blocks_attrs = ["blocks"]
-
-            def __init__(self):
-                super().__init__()
-                self.blocks = nn.ModuleList(
-                    [
-                        nn.Linear(2, 2, bias=False),
-                        nn.Linear(2, 2, bias=False),
-                    ]
-                )
-
-        class Pipeline(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.transformer = Transformer()
-
-            @staticmethod
-            def _remap_ckpt_key(key):
-                return key
-
-        pipeline = Pipeline()
-        weights = {name: torch.ones_like(param) for name, param in pipeline.named_parameters()}
-        checkpoint_file = tmp_path / "model.safetensors"
-        save_file(weights, str(checkpoint_file))
-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=(),
-            model_path=str(tmp_path),
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
-        )
-        assert result.plan is not None
-        backend = DistributedLayerwiseOffloadBackend(
-            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
-                dp_size=1,
-                dlo_use_allgather=True,
-            ),
-            torch.device("cpu"),
-            host_weight_plan=result.plan,
-        )
-
-        backend.enable(pipeline)
-
-        assert backend._using_rank_local_mmap
-        assert backend.dp_group is None
-        assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
-        backend.disable()
-

 class TestGetBlocksFromDit:
    def test_get_blocks_from_dit_single_block_attr(self):
@@ -1005,6 +897,46 @@ class TestOffloadPlan:
        with pytest.raises(Exception):
            plan.block_attrs = {"x": ("y",)}  # type: ignore

+    def test_explicit_component_sources_can_opt_in_to_mmap(self):
+        class Pipeline(nn.Module):
+            _supports_mmap_loading = True
+            weights_sources = [object()]
+
+        assert supports_mmap_loading(Pipeline())
+
+    def test_component_sources_do_not_implicitly_enable_mmap(self):
+        class Pipeline(nn.Module):
+            weights_sources = [object()]
+
+        assert not supports_mmap_loading(Pipeline())
+
+    def test_online_quantization_detection(self):
+        model = nn.Sequential(nn.Linear(1, 1))
+        assert not has_online_quantization(model)
+
+        model[0].quant_method = SimpleNamespace(uses_meta_device=True)
+        assert has_online_quantization(model)
+
+    @pytest.mark.parametrize(
+        ("kwargs", "expected"),
+        [
+            ({"dlo_use_allgather": True, "tensor_parallel_size": 2}, True),
+            ({"dlo_use_allgather": False, "tensor_parallel_size": 1}, True),
+            ({"dlo_use_allgather": False, "tensor_parallel_size": 2}, False),
+            ({"dlo_use_allgather": False, "tensor_parallel_size": 1, "use_hsdp": True}, False),
+            ({"dlo_use_allgather": False, "tensor_parallel_size": 1, "has_online_quant": True}, False),
+        ],
+    )
+    def test_dlo_mmap_policy(self, kwargs, expected):
+        defaults = {
+            "mmap_supported": True,
+            "dlo_use_allgather": False,
+            "has_online_quant": False,
+            "tensor_parallel_size": 1,
+            "use_hsdp": False,
+        }
+        assert should_select_dlo_mmap(**(defaults | kwargs)) is expected
+
    def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime):
        class TokenRefiner(nn.Module):
            def __init__(self):
@@ -1051,23 +983,20 @@ class TestOffloadPlan:


 class TestMmapValidation:
-    """Tests for loader preflight and backend plan realization."""
+    """Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""

    def test_component_source_prefix_builds_model_namespace(self, tmp_path):
        source_root = tmp_path / "partition"
        checkpoint_dir = source_root / "transformer"
        checkpoint_dir.mkdir(parents=True)
        checkpoint_file = checkpoint_dir / "model.safetensors"
-        save_file(
-            {"blocks.0.weight": torch.arange(4, dtype=torch.float32).reshape(1, 4)},
-            str(checkpoint_file),
-        )
+        save_file({"blocks.0.weight": torch.arange(4, dtype=torch.float32)}, str(checkpoint_file))

        class Pipeline(nn.Module):
+            _supports_mmap_loading = True
+
            def __init__(self):
                super().__init__()
-                self.transformer = nn.Module()
-                self.transformer.blocks = nn.ModuleList([nn.Linear(4, 1, bias=False)])
                self.weights_sources = [
                    SimpleNamespace(
                        model_or_path=str(source_root),
@@ -1077,51 +1006,29 @@ class TestMmapValidation:
                    )
                ]

-        pipeline = Pipeline()
-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=tuple(pipeline.weights_sources),
-            model_path=None,
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
-        )
-
-        assert result.fallback_reason is None
-        assert result.plan is not None
-        assert result.plan.bindings["transformer.blocks.0.weight"] == TensorBinding(
-            checkpoint_key="blocks.0.weight",
-            file_path=str(checkpoint_file),
-        )
-        assert result.plan.planned_source_prefixes == frozenset({"transformer."})
-
-    def test_non_dedicated_component_source_falls_back_before_discovery(self):
-        pipeline = nn.Module()
-        pipeline.transformer = nn.Linear(2, 2, bias=False)
-        sources = (
-            SimpleNamespace(
-                model_or_path="unused",
-                subfolder=None,
-                revision=None,
-                prefix="",
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
            ),
+            torch.device("cpu"),
        )

-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=sources,
-            model_path=None,
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
+        model_map = backend._build_mmap_model_map(Pipeline())
+
+        assert model_map == {
+            "transformer.blocks.0.weight": (
+                "blocks.0.weight",
+                str(checkpoint_file),
+            )
+        }
+
+    def test_explicit_mmap_opt_in_rejects_missing_model_tensors(self, tmp_path, patched_offload_runtime, monkeypatch):
+        monkeypatch.setattr(
+            "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
+            lambda: 1,
        )

-        assert result.plan is None
-        assert "dedicated component weight source" in result.fallback_reason
-
-    def test_preflight_falls_back_before_mutation_for_missing_model_tensors(self, tmp_path):
        source_root = tmp_path / "partition"
        checkpoint_dir = source_root / "transformer"
        checkpoint_dir.mkdir(parents=True)
@@ -1138,6 +1045,8 @@ class TestMmapValidation:
                self.blocks = nn.ModuleList([nn.Linear(2, 2, bias=False), nn.Linear(2, 2, bias=False)])

        class Pipeline(nn.Module):
+            _supports_mmap_loading = True
+
            def __init__(self):
                super().__init__()
                self.transformer = Transformer()
@@ -1151,23 +1060,33 @@ class TestMmapValidation:
                ]

        pipeline = Pipeline()
-        original_weight = pipeline.transformer.blocks[0].weight
-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=tuple(pipeline.weights_sources),
-            model_path=None,
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+            ),
+            torch.device("cpu"),
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+
+        with pytest.raises(RuntimeError, match="Explicit mmap loading left 1"):
+            backend._load_weights_via_mmap(pipeline, modules)
+
+    def test_tp_rejected_when_tp_world_size_gt_1(self, tmp_path, patched_offload_runtime, monkeypatch):
+        """DLO+AllGather should reject when TP world size > 1."""
+        import json
+        from types import SimpleNamespace
+
+        from safetensors.torch import save_file
+
+        from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+
+        # Mock TP world size = 2
+        monkeypatch.setattr(
+            "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
+            lambda: 2,
        )

-        assert result.plan is None
-        assert "1 required DiT tensors" in result.fallback_reason
-        assert pipeline.transformer.blocks[0].weight is original_weight
-        assert not original_weight.is_meta
-
-    def test_tp_falls_back_before_checkpoint_discovery(self):
        class SimpleModel(nn.Module):
            _layerwise_offload_blocks_attrs = ["blocks"]

@@ -1185,118 +1104,41 @@ class TestMmapValidation:
                return key

        pipeline = SimplePipeline()
-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=(),
-            model_path="unused",
-            tensor_parallel_size=2,
-            use_hsdp=False,
-            online_quantization=False,
+        weights = {name: torch.ones(p.shape, dtype=p.dtype) for name, p in pipeline.named_parameters() if not p.is_meta}
+        save_file(weights, str(tmp_path / "model.safetensors"))
+        (tmp_path / "model.safetensors.index.json").write_text(
+            json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
        )

-        assert result.plan is None
-        assert result.fallback_reason == "TP=2 requires the ordinary loader"
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                model_path=str(tmp_path),
+            ),
+            torch.device("cpu"),
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])

-    @pytest.mark.parametrize(
-        ("use_hsdp", "online_quantization", "expected_reason"),
-        [
-            (True, False, "HSDP requires the ordinary loader"),
-            (False, True, "online quantization requires the ordinary loader"),
-        ],
-    )
-    def test_runtime_layouts_that_need_loader_callbacks_fall_back(
-        self,
-        use_hsdp,
-        online_quantization,
-        expected_reason,
-    ):
-        pipeline = nn.Module()
-        pipeline.transformer = nn.Linear(2, 2, bias=False)
+        with pytest.raises(ValueError, match="Tensor Parallel"):
+            backend._load_weights_via_mmap(pipeline, modules)

-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=(),
-            model_path="unused",
-            tensor_parallel_size=1,
-            use_hsdp=use_hsdp,
-            online_quantization=online_quantization,
+    def test_tp_allowed_when_tp_world_size_eq_1(self, tmp_path, patched_offload_runtime, monkeypatch):
+        """DLO+AllGather should work when TP world size = 1, even if params
+        have custom weight_loader attributes (e.g. QKVParallelLinear at TP=1)."""
+        import json
+        from types import SimpleNamespace
+
+        from safetensors.torch import save_file
+
+        from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+
+        # Mock TP world size = 1 (default)
+        monkeypatch.setattr(
+            "vllm.distributed.parallel_state.get_tensor_model_parallel_world_size",
+            lambda: 1,
        )

-        assert result.plan is None
-        assert result.fallback_reason == expected_reason
-
-    @pytest.mark.parametrize(
-        ("checkpoint_tensor", "expected_reason"),
-        [
-            (torch.ones((3, 2), dtype=torch.float32), "shape mismatch"),
-            (torch.ones((2, 2), dtype=torch.bfloat16), "dtype mismatch"),
-        ],
-    )
-    def test_source_metadata_mismatch_falls_back(self, tmp_path, checkpoint_tensor, expected_reason):
-        checkpoint_file = tmp_path / "model.safetensors"
-        save_file({"transformer.weight": checkpoint_tensor}, str(checkpoint_file))
-
-        class Pipeline(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.transformer = nn.Linear(2, 2, bias=False)
-
-            @staticmethod
-            def _remap_ckpt_key(key):
-                return key
-
-        pipeline = Pipeline()
-
-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=(),
-            model_path=str(tmp_path),
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
-        )
-
-        assert result.plan is None
-        assert expected_reason in result.fallback_reason
-
-    def test_non_persistent_buffers_do_not_need_checkpoint_bindings(self, tmp_path):
-        class Transformer(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.weight = nn.Parameter(torch.ones(2, 2))
-                self.register_buffer("derived", torch.ones(2), persistent=False)
-
-        checkpoint_file = tmp_path / "model.safetensors"
-        save_file({"transformer.weight": torch.ones(2, 2)}, str(checkpoint_file))
-
-        class Pipeline(nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.transformer = Transformer()
-
-            @staticmethod
-            def _remap_ckpt_key(key):
-                return key
-
-        pipeline = Pipeline()
-
-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=(),
-            model_path=str(tmp_path),
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
-        )
-
-        assert result.plan is not None
-        assert set(result.plan.bindings) == {"transformer.weight"}
-
-    def test_unadapted_custom_weight_loader_falls_back_at_tp1(self, tmp_path):
        def custom_weight_loader(param, weight):
            param.data.copy_(weight)

@@ -1326,67 +1168,27 @@ class TestMmapValidation:
            json.dumps({"weight_map": {k: "model.safetensors" for k in weights}})
        )

-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=(),
-            model_path=str(tmp_path),
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
-        )
-
-        assert result.plan is None
-        assert "custom weight_loader" in result.fallback_reason
-
-    def test_cosmos3_adapter_preserves_existing_tp1_direct_mmap_contract(self, tmp_path):
-        from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
-            Cosmos3OmniDiffusersPipeline,
-        )
-        from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
-            Cosmos3VFMTransformer,
-        )
-
-        transformer = object.__new__(Cosmos3VFMTransformer)
-        nn.Module.__init__(transformer)
-        transformer.proj_in = nn.Linear(2, 2, bias=False)
-        transformer.proj_in.weight.weight_loader = lambda *_args: None
-        pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
-        nn.Module.__init__(pipeline)
-        pipeline.transformer = transformer
-
-        source_root = tmp_path / "partition"
-        checkpoint_dir = source_root / "transformer"
-        checkpoint_dir.mkdir(parents=True)
-        checkpoint_file = checkpoint_dir / "model.safetensors"
-        save_file({"proj_in.weight": torch.ones(2, 2)}, str(checkpoint_file))
-        sources = (
-            SimpleNamespace(
-                model_or_path=str(source_root),
-                subfolder="transformer",
-                revision=None,
-                prefix="transformer.",
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                model_path=str(tmp_path),
            ),
+            torch.device("cpu"),
        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])

-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", transformer),),
-            sources=sources,
-            model_path=None,
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
-        )
-
-        assert result.plan is not None
-        assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
-            checkpoint_key="proj_in.weight",
-            file_path=str(checkpoint_file),
-        )
+        # Should NOT raise — TP=1 is allowed even with custom weight_loader
+        backend._load_weights_via_mmap(pipeline, modules)

    def test_validate_loaded_weights_called(self, tmp_path, patched_offload_runtime):
        """validate_loaded_weights should be called after mmap loading."""
+        import json
+        from types import SimpleNamespace
+
+        from safetensors.torch import save_file
+
+        from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy

        class ModelWithValidate(nn.Module):
            _layerwise_offload_blocks_attrs = ["blocks"]
@@ -1428,64 +1230,19 @@ class TestMmapValidation:

        backend = DistributedLayerwiseOffloadBackend(
            OffloadConfig(
-                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
-                pin_cpu_memory=False,
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, pin_cpu_memory=False, model_path=str(tmp_path)
            ),
            torch.device("cpu"),
        )
        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
-        result = build_checkpoint_mmap_plan(
-            pipeline,
-            dit_modules=(("transformer", pipeline.transformer),),
-            sources=(),
-            model_path=str(tmp_path),
-            tensor_parallel_size=1,
-            use_hsdp=False,
-            online_quantization=False,
-        )

-        assert result.plan is not None
-        backend._load_weights_via_mmap(pipeline, modules, result.plan)
+        backend._load_weights_via_mmap(pipeline, modules)
        assert pipeline.transformer.validate_called, "validate_loaded_weights should be called"


 class TestConfigValidation:
    """Tests for configuration validation in OffloadConfig / execute_request."""

-    def test_loader_plan_cannot_be_silently_dropped_on_unsupported_platform(self, monkeypatch):
-        import vllm_omni.diffusion.offloader as offloader_module
-
-        od_config = SimpleNamespace(
-            enable_cpu_offload=False,
-            enable_layerwise_offload=False,
-            enable_distributed_layerwise_offload=True,
-            dlo_use_allgather=False,
-            dlo_resident_layers=0,
-            pin_cpu_memory=False,
-            model="unused",
-            parallel_config=SimpleNamespace(
-                use_hsdp=False,
-                data_parallel_size=1,
-                sequence_parallel_size=1,
-            ),
-        )
-        plan = HostWeightPlan(
-            backing_kind="checkpoint_mmap",
-            bindings={},
-        )
-        monkeypatch.setattr(
-            offloader_module.current_omni_platform,
-            "supports_cpu_offload",
-            lambda: False,
-        )
-
-        with pytest.raises(RuntimeError, match="skipped ordinary weight materialization"):
-            offloader_module.get_offload_backend(
-                od_config,
-                device=torch.device("cpu"),
-                host_weight_plan=plan,
-            )
-
    def test_hsdp_with_allgather_rejected(self):
        """HSDP + DLO + AllGather should raise ValueError (double sharding)."""
        from vllm_omni.diffusion.offloader.base import OffloadConfig
@@ -79,11 +79,11 @@ def dist_group():


@pytest.fixture
-def patched_offload_runtime(monkeypatch):
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
-    monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
+def patched_offload_runtime(mocker):
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
+    mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)


 class TinyBlock(nn.Module):
@@ -128,33 +128,6 @@ class TestLayerwiseOffloadHook:
        assert current_block.weight.to_local().shape == torch.Size([4])
        assert not hook.is_materialized

-    def test_prefetch_preserves_transposed_weight_stride(self, patched_offload_runtime):
-        """Online-FP8 Cutlass weights must retain their transposed layout."""
-
-        class StridedBlock(nn.Module):
-            def __init__(self, start: float):
-                super().__init__()
-                base = torch.arange(start, start + 12).reshape(3, 4)
-                self.weight = nn.Parameter(base.t(), requires_grad=False)
-
-        current_block = StridedBlock(1.0)
-        next_block = StridedBlock(20.0)
-        expected = next_block.weight.detach().clone()
-        expected_stride = next_block.weight.stride()
-        assert expected_stride == (1, 4)
-
-        hook = LayerwiseOffloadHook(
-            next_block=next_block,
-            device=torch.device("cpu"),
-            stream=DummyStream(),
-            pin_memory=False,
-        )
-        hook.initialize_hook(current_block)
-        hook.prefetch_layer(non_blocking=False)
-
-        assert next_block.weight.stride() == expected_stride
-        assert torch.equal(next_block.weight, expected)
-

 class _DummyBlock(nn.Module):
    def __init__(self):
@@ -16,11 +16,10 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
 )
 from vllm_omni.quantization.component_config import (
    ComponentQuantizationConfig,
-    resolve_component_quant_config,
 )
 from vllm_omni.quantization.inc_config import OmniINCConfig

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


 # ---------------------------------------------------------------------------
@@ -348,24 +347,6 @@ class TestTalkerVisualRouting:
 # ===================================================================


-class TestResolveComponentQuantConfig:
-    def test_global_config_applies_to_each_quantization_aware_component(self):
-        config = _MockQuantConfig("fp8")
-
-        assert resolve_component_quant_config(config, "transformer") is config
-        assert resolve_component_quant_config(config, "text_encoder") is config
-
-    def test_component_map_is_the_only_scope_override(self):
-        transformer = _MockQuantConfig("fp8")
-        config = ComponentQuantizationConfig(
-            component_configs={"transformer": transformer},
-            default_config=None,
-        )
-
-        assert resolve_component_quant_config(config, "transformer") is transformer
-        assert resolve_component_quant_config(config, "text_encoder") is None
-
-
 class TestComponentResolve:
    def test_longest_prefix_match(self):
        a = _MockQuantConfig("a")
@@ -36,23 +36,15 @@ class _ToyKernel:
 # --- helper to build a toy module simulating online-FP8 post-load state ---


-def _make_fp8_toy_module(
-    out_features: int = 16,
-    in_features: int = 32,
-    *,
-    quant_method_cls=None,
-):
+def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32):
    """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
    ``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``.
    """
    from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod

-    if quant_method_cls is None:
-        quant_method_cls = Fp8LinearMethod
-
    # Create the quant-method instance without calling __init__ (avoids
    # heavyweight upstream config dependencies).  isinstance() still works.
-    qm = object.__new__(quant_method_cls)
+    qm = object.__new__(Fp8LinearMethod)
    qm.fp8_linear = _ToyKernel()

    # randn does not support float8 on CPU; zeros does and the test only
@@ -112,28 +104,6 @@ def test_rewrites_non_contiguous_weight():
    assert module.weight.shape == (16, 32)


-def test_rewrites_per_tensor_online_fp8_weight():
-    """Regression test for vLLM's new online FP8 frontend (#45463)."""
-    from vllm.model_executor.layers.quantization.online.fp8 import (
-        Fp8PerTensorOnlineLinearMethod,
-    )
-
-    module = _make_fp8_toy_module(
-        16,
-        32,
-        quant_method_cls=Fp8PerTensorOnlineLinearMethod,
-    )
-
-    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
-
-    assert n_rewritten_layers == 1
-    assert module.weight.is_contiguous()
-    assert module.weight.shape == (16, 32)
-    w, *_ = module.quant_method.fp8_linear._get_layer_params(module)
-    assert w.shape == (32, 16)
-    assert w.stride() == (1, 32)
-
-
 def test_patched_get_layer_params_returns_transposed_view():
    module = _make_fp8_toy_module(16, 32)
    prepare_fp8_layers_for_fsdp(module)
@@ -3,8 +3,6 @@
 """Unit tests for WorkerProc async output logic."""

 import queue
-import threading
-import time
 from unittest.mock import MagicMock

 import pytest
@@ -31,59 +29,9 @@ def _make_worker_proc(step_execution=False):
    proc._running = True
    proc._async_output_queue = queue.Queue() if not step_execution else None
    proc._async_output_thread = None
-    proc._async_output_done = threading.Condition()
-    proc._async_output_pending = 0
-    proc._result_mq_lock = threading.Lock()
    return proc


-class _OverlapDetectingQueue:
-    """Fails if two threads are inside enqueue() at the same time.
-
-    Mirrors vLLM's MessageQueue, whose acquire_write() reads current_idx,
-    writes that block, then advances it. Overlapping writers can target the
-    same ring block and silently drop a message.
-    """
-
-    def __init__(self):
-        self._active = 0
-        self._guard = threading.Lock()
-        self.overlaps = 0
-        self.messages = []
-
-    def enqueue(self, msg, timeout=None):
-        with self._guard:
-            self._active += 1
-            if self._active > 1:
-                self.overlaps += 1
-        time.sleep(0.001)  # widen the window a real ring-buffer write would have
-        with self._guard:
-            self.messages.append(msg)
-            self._active -= 1
-
-
-class TestResultQueueSingleWriter:
-    """Worker must serialize writes to the single-writer result queue."""
-
-    def test_concurrent_enqueues_do_not_overlap(self):
-        proc = _make_worker_proc(step_execution=False)
-        detector = _OverlapDetectingQueue()
-        proc.result_mq = detector
-
-        def writer(tag):
-            for i in range(20):
-                proc._enqueue_result(f"{tag}-{i}")
-
-        threads = [threading.Thread(target=writer, args=(tag,)) for tag in ("compute", "output")]
-        for t in threads:
-            t.start()
-        for t in threads:
-            t.join()
-
-        assert detector.overlaps == 0, f"{detector.overlaps} overlapping writes to result_mq"
-        assert len(detector.messages) == 40
-
-
 class TestGenerateAsyncOutputId:
    """Test _generate_async_output_id uniqueness."""

@@ -244,87 +192,6 @@ class TestAsyncOutputLoopLogic:
        assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls"


-class TestDrainAsyncOutputs:
-    """Async outputs must finish packing before device memory is released."""
-
-    def test_drain_returns_immediately_when_nothing_in_flight(self):
-        proc = _make_worker_proc(step_execution=False)
-
-        assert proc.drain_async_outputs(timeout=0.01) is True
-
-    def test_return_result_keeps_output_pending_until_loop_completes(self, mocker):
-        """_return_result marks work pending; the bg loop clears it."""
-        proc = _make_worker_proc(step_execution=False)
-
-        mock_platform = mocker.MagicMock()
-        mock_platform.record_device_event.return_value = MagicMock()
-        mocker.patch(
-            "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform",
-            mock_platform,
-        )
-        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm")
-        mock_torch = mocker.MagicMock()
-        mock_torch.accelerator.current_accelerator.return_value.type = "cpu"
-        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch)
-
-        proc._return_result(DiffusionOutput(output="data"), rpc_id="1")
-
-        # Packing has not run yet, so a sleep here would free tensors in flight.
-        assert proc._async_output_pending == 1
-        assert proc.drain_async_outputs(timeout=0.05) is False
-
-        thread = threading.Thread(target=proc._async_output_loop, daemon=True)
-        thread.start()
-        try:
-            assert proc.drain_async_outputs(timeout=5.0) is True
-        finally:
-            proc._async_output_queue.put(None)
-            thread.join(timeout=2.0)
-
-    def test_sleep_rpc_drains_before_releasing_memory(self, mocker):
-        proc = _make_worker_proc(step_execution=False)
-        proc.worker = MagicMock()
-
-        call_order: list[str] = []
-        mocker.patch.object(
-            proc,
-            "drain_async_outputs",
-            side_effect=lambda *a, **kw: call_order.append("drain") or True,
-        )
-        proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append("sleep")
-
-        proc._execute_rpc(
-            {
-                "method": "handle_sleep_task",
-                "args": (),
-                "kwargs": {},
-                "output_rank": 0,
-                "exec_all_ranks": False,
-                "collect_rank_status": False,
-            }
-        )
-
-        assert call_order == ["drain", "sleep"]
-
-    def test_non_sleep_rpc_does_not_drain(self, mocker):
-        proc = _make_worker_proc(step_execution=False)
-        proc.worker = MagicMock()
-        drain = mocker.patch.object(proc, "drain_async_outputs", return_value=True)
-
-        proc._execute_rpc(
-            {
-                "method": "execute_model",
-                "args": (),
-                "kwargs": {},
-                "output_rank": 0,
-                "exec_all_ranks": False,
-                "collect_rank_status": False,
-            }
-        )
-
-        drain.assert_not_called()
-
-
 class TestWorkerProcShutdown:
    def test_shutdown_stops_async_thread_and_releases_queues(self):
        import threading
@@ -123,7 +123,6 @@ class _ConcurrencyTrackingExecutor:
 def _make_request(tag: str):
    sampling_params = dict(_SAMPLING_KEY_DEFAULTS)
    sampling_params["num_inference_steps"] = 1
-    sampling_params["extra_args"] = {}
    return SimpleNamespace(
        request_id=tag,
        prompt=f"prompt_{tag}",
@@ -305,49 +305,6 @@ def test_batch_runner_output_round_trips_nested_results_through_shm() -> None:
    assert output["req-error"].result.error == "boom"


-def test_per_request_views_of_a_batch_tensor_do_not_ship_whole_storage() -> None:
-    """Regression: per-request slices used to pickle the whole batch storage.
-
-    Each view is under the packing threshold, but pickle writes the entire
-    shared storage per view, so a batch's message grew by the full batch
-    tensor once per request and overflowed the ring buffer.
-    """
-    import pickle
-
-    num_requests = 16
-    # Each row stays under the threshold; the shared storage is well over it.
-    row_numel = _SHM_TENSOR_THRESHOLD // (2 * torch.empty((), dtype=torch.float32).element_size())
-    batch = torch.arange(num_requests * row_numel, dtype=torch.float32).reshape(num_requests, row_numel)
-    per_request = batch[0]
-    assert per_request.nelement() * per_request.element_size() <= _SHM_TENSOR_THRESHOLD
-    assert per_request.untyped_storage().nbytes() > _SHM_TENSOR_THRESHOLD
-
-    output = BatchRunnerOutput.from_list(
-        [
-            RunnerOutput(
-                request_id=f"req-{i}",
-                finished=True,
-                result=DiffusionOutput(trajectory_latents=batch[i]),
-            )
-            for i in range(num_requests)
-        ]
-    )
-    unpacked_size = len(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL))
-
-    pack_diffusion_output_shm(output)
-    packed_size = len(pickle.dumps(output, protocol=pickle.HIGHEST_PROTOCOL))
-
-    assert output.runner_outputs[0].result.trajectory_latents["__tensor_shm__"] is True
-    # Whole storage per request before; handles only after.
-    storage_bytes = per_request.untyped_storage().nbytes()
-    assert unpacked_size > num_requests * storage_bytes / 2
-    assert packed_size < _SHM_TENSOR_THRESHOLD
-
-    unpack_diffusion_output_shm(output)
-    for i in range(num_requests):
-        torch.testing.assert_close(output[f"req-{i}"].result.trajectory_latents, batch[i])
-
-
 def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
    tensor = torch.arange(
        _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),