Support one-sided MoT biases

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
[Doc] Add production diffusion model skill (#6097 )
2026-09-03 17:05:57 +00:00 · 2026-09-03 20:09:28 +08:00 · 2026-09-03 18:48:43 +08:00 · 2026-09-03 00:54:36 -07:00 · 2026-09-03 06:44:42 +00:00 · 2026-09-02 23:19:08 -07:00
209 changed files with 16673 additions and 2227 deletions
@@ -1,19 +1,31 @@
 # Uploader-only hardware presets for CUDA / NPU test YAML.
 # Referenced as ``mirror_hardwares: l4_1`` on a step; ``upload_pipeline.py`` expands to
-# ``agents`` (+ optional top-level ``image`` for NPU) + ``plugins``, then strips
-# ``mirror_hardwares`` before Buildkite upload.
+# ``agents`` (+ optional top-level ``image`` for NPU) + ``plugins`` (+ ``retry``
+# for ``l4_*``), then strips ``mirror_hardwares`` before Buildkite upload.
+#
+# Omit the key: compose {chip}_{n} from pytest -m SKU + cards_n (max of several
+# positives; not cards_1 with no positive cards_* → that chip's highest existing preset).
+# Unset MIRROR_HW matches H100 or L4 in -m (both → H100); no match skips the step.
+# MIRROR_HW=b200 matches B200 in -m (else the step is skipped).
+# MIRROR_HW must be empty or b200 (case-insensitive); unknown values fail the upload.
+# CUDA preset strings (h100_4, l4_1) are omitted when MIRROR_HW=b200.
+# Unset MIRROR_HW keeps string presets. NPU presets ignore MIRROR_HW.
 #
 # CUDA naming: {chip}_{count}
-#   l4_1, l4_4, h100_1 … h100_4
+#   l4_1, l4_2, l4_3, l4_4, h100_1 … h100_4, b200_1 … b200_4
 # NPU naming: {queue}_{resource_class}
 #   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
+#   a3_npu_2    — ascend-a3, resource_class npu-2 (A3 CI image tag …:$BUILDKITE_COMMIT-a3)
 #
-# 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.
+# B200 (b200-k8s): matches vLLM device ``b200-k8s`` from ci-infra
+# (AgentQueue.B200_K8S + b200_plugin_template). Queue routes the K8s pool;
+# podSpec uses hostNetwork + /raid|/mnt/shared (no H100-style nodeSelector).
+#
+# The ``l4_*`` presets run on the EKS ``l4-k8s`` queue (agent-stack-k8s),
+# packed onto g6.12xlarge L4 nodes (``vllm.ci/gpu-pool=l4x4``). GPU / CPU /
+# memory / shm requests scale with card count (1 GPU = 10 CPU / 40Gi / 24Gi shm).
+# Retry covers lost agents / killed pods, not pytest or OOM (exit 1 / 137).

 x-npu-k8s-plugins: &npu_k8s_plugins
  - kubernetes:
@@ -21,39 +33,224 @@ x-npu-k8s-plugins: &npu_k8s_plugins
        imagePullSecrets:
          - name: swr-secret

+# Infra-only retry. Omit exit 1 (pytest failed) and 137 (cgroup OOM / SIGKILL
+# of the test process). ci-infra ``K8S_RETRY`` still retries exit 1.
+x-l4-k8s-retry: &l4_k8s_retry
+  automatic:
+    - exit_status: -1
+      limit: 1
+    - exit_status: 128
+      limit: 1
+    - signal_reason: agent_stop
+      limit: 1
+    - signal_reason: agent_refused
+      limit: 1
+

 mirror_hardwares:
  l4_1:
    agents:
-      queue: gpu_1_queue
+      queue: l4-k8s
+    retry: *l4_k8s_retry
    plugins:
-      - docker#v5.2.0:
-          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
-            - HF_TOKEN
-          volumes:
-            - /fsx/hf_cache:/fsx/hf_cache
+      - kubernetes:
+          metadata:
+            annotations:
+              cluster-autoscaler.kubernetes.io/safe-to-evict: "false"
+              karpenter.sh/do-not-disrupt: "true"
+          podSpec:
+            priorityClassName: ci
+            serviceAccountName: buildkite-gpu-jobs
+            tolerations:
+              - key: nvidia.com/gpu
+                operator: Exists
+                effect: NoSchedule
+            nodeSelector:
+              vllm.ci/gpu-pool: l4x4
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                  requests:
+                    nvidia.com/gpu: 1
+                    cpu: "10"
+                    memory: 40Gi
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /fsx/hf_cache
+                env:
+                  - name: HF_HOME
+                    value: /fsx/hf_cache
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+                  sizeLimit: 24Gi
+              - name: hf-cache
+                hostPath:
+                  path: /fsx/hf_cache
+                  type: DirectoryOrCreate
+
+  l4_2:
+    agents:
+      queue: l4-k8s
+    retry: *l4_k8s_retry
+    plugins:
+      - kubernetes:
+          metadata:
+            annotations:
+              cluster-autoscaler.kubernetes.io/safe-to-evict: "false"
+              karpenter.sh/do-not-disrupt: "true"
+          podSpec:
+            priorityClassName: ci
+            serviceAccountName: buildkite-gpu-jobs
+            tolerations:
+              - key: nvidia.com/gpu
+                operator: Exists
+                effect: NoSchedule
+            nodeSelector:
+              vllm.ci/gpu-pool: l4x4
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                  requests:
+                    nvidia.com/gpu: 2
+                    cpu: "20"
+                    memory: 80Gi
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /fsx/hf_cache
+                env:
+                  - name: HF_HOME
+                    value: /fsx/hf_cache
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+                  sizeLimit: 48Gi
+              - name: hf-cache
+                hostPath:
+                  path: /fsx/hf_cache
+                  type: DirectoryOrCreate
+
+  l4_3:
+    agents:
+      queue: l4-k8s
+    retry: *l4_k8s_retry
+    plugins:
+      - kubernetes:
+          metadata:
+            annotations:
+              cluster-autoscaler.kubernetes.io/safe-to-evict: "false"
+              karpenter.sh/do-not-disrupt: "true"
+          podSpec:
+            priorityClassName: ci
+            serviceAccountName: buildkite-gpu-jobs
+            tolerations:
+              - key: nvidia.com/gpu
+                operator: Exists
+                effect: NoSchedule
+            nodeSelector:
+              vllm.ci/gpu-pool: l4x4
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 3
+                  requests:
+                    nvidia.com/gpu: 3
+                    cpu: "30"
+                    memory: 120Gi
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /fsx/hf_cache
+                env:
+                  - name: HF_HOME
+                    value: /fsx/hf_cache
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+                  sizeLimit: 72Gi
+              - name: hf-cache
+                hostPath:
+                  path: /fsx/hf_cache
+                  type: DirectoryOrCreate

  l4_4:
    agents:
-      queue: gpu_4_queue
+      queue: l4-k8s
+    retry: *l4_k8s_retry
    plugins:
-      - docker#v5.2.0:
-          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
-            - HF_TOKEN
-          volumes:
-            - /fsx/hf_cache:/fsx/hf_cache
+      - kubernetes:
+          metadata:
+            annotations:
+              cluster-autoscaler.kubernetes.io/safe-to-evict: "false"
+              karpenter.sh/do-not-disrupt: "true"
+          podSpec:
+            priorityClassName: ci
+            serviceAccountName: buildkite-gpu-jobs
+            tolerations:
+              - key: nvidia.com/gpu
+                operator: Exists
+                effect: NoSchedule
+            nodeSelector:
+              vllm.ci/gpu-pool: l4x4
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 4
+                  requests:
+                    nvidia.com/gpu: 4
+                    cpu: "40"
+                    memory: 160Gi
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /fsx/hf_cache
+                env:
+                  - name: HF_HOME
+                    value: /fsx/hf_cache
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+                  sizeLimit: 96Gi
+              - name: hf-cache
+                hostPath:
+                  path: /fsx/hf_cache
+                  type: DirectoryOrCreate

  h100_1:
    agents:
@@ -195,6 +392,211 @@ mirror_hardwares:
                  path: /mnt/hf-cache
                  type: DirectoryOrCreate

+  # vLLM device: b200-k8s → agents.queue b200-k8s + Buildkite K8s agent-stack podSpec
+  b200_1:
+    agents:
+      queue: b200-k8s
+    plugins:
+      - kubernetes:
+          podSpec:
+            runtimeClassName: nvidia
+            hostNetwork: true
+            dnsPolicy: ClusterFirstWithHostNet
+            imagePullSecrets:
+              - name: k8s-ecr-login-renew-docker-secret
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                securityContext:
+                  capabilities:
+                    add:
+                      - IPC_LOCK
+                      - SYS_RESOURCE
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: raid
+                    mountPath: /raid
+                  - name: shared
+                    mountPath: /mnt/shared
+                env:
+                  - name: NCCL_CUMEM_HOST_ENABLE
+                    value: "0"
+                  - name: HF_HOME
+                    value: /raid/hf_cache
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: raid
+                hostPath:
+                  path: /raid
+                  type: DirectoryOrCreate
+              - name: shared
+                hostPath:
+                  path: /mnt/shared
+                  type: DirectoryOrCreate
+
+  b200_2:
+    agents:
+      queue: b200-k8s
+    plugins:
+      - kubernetes:
+          podSpec:
+            runtimeClassName: nvidia
+            hostNetwork: true
+            dnsPolicy: ClusterFirstWithHostNet
+            imagePullSecrets:
+              - name: k8s-ecr-login-renew-docker-secret
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                securityContext:
+                  capabilities:
+                    add:
+                      - IPC_LOCK
+                      - SYS_RESOURCE
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: raid
+                    mountPath: /raid
+                  - name: shared
+                    mountPath: /mnt/shared
+                env:
+                  - name: NCCL_CUMEM_HOST_ENABLE
+                    value: "0"
+                  - name: HF_HOME
+                    value: /raid/hf_cache
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: raid
+                hostPath:
+                  path: /raid
+                  type: DirectoryOrCreate
+              - name: shared
+                hostPath:
+                  path: /mnt/shared
+                  type: DirectoryOrCreate
+
+  b200_3:
+    agents:
+      queue: b200-k8s
+    plugins:
+      - kubernetes:
+          podSpec:
+            runtimeClassName: nvidia
+            hostNetwork: true
+            dnsPolicy: ClusterFirstWithHostNet
+            imagePullSecrets:
+              - name: k8s-ecr-login-renew-docker-secret
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 3
+                securityContext:
+                  capabilities:
+                    add:
+                      - IPC_LOCK
+                      - SYS_RESOURCE
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: raid
+                    mountPath: /raid
+                  - name: shared
+                    mountPath: /mnt/shared
+                env:
+                  - name: NCCL_CUMEM_HOST_ENABLE
+                    value: "0"
+                  - name: HF_HOME
+                    value: /raid/hf_cache
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: raid
+                hostPath:
+                  path: /raid
+                  type: DirectoryOrCreate
+              - name: shared
+                hostPath:
+                  path: /mnt/shared
+                  type: DirectoryOrCreate
+
+  b200_4:
+    agents:
+      queue: b200-k8s
+    plugins:
+      - kubernetes:
+          podSpec:
+            runtimeClassName: nvidia
+            hostNetwork: true
+            dnsPolicy: ClusterFirstWithHostNet
+            imagePullSecrets:
+              - name: k8s-ecr-login-renew-docker-secret
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 4
+                securityContext:
+                  capabilities:
+                    add:
+                      - IPC_LOCK
+                      - SYS_RESOURCE
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: raid
+                    mountPath: /raid
+                  - name: shared
+                    mountPath: /mnt/shared
+                env:
+                  - name: NCCL_CUMEM_HOST_ENABLE
+                    value: "0"
+                  - name: HF_HOME
+                    value: /raid/hf_cache
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: raid
+                hostPath:
+                  path: /raid
+                  type: DirectoryOrCreate
+              - name: shared
+                hostPath:
+                  path: /mnt/shared
+                  type: DirectoryOrCreate
+
  a2b3_npu_1:
    agents:
      queue: "ascend-a2b3"
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Shared Buildkite skip-ci and CI-YAML level targeting logic.

 ``resolve_ci_decision(changed_files, ...)`` buckets explicit file lists (tests, dry-run).
@@ -10,12 +11,6 @@ CLI (for bootstrap scripts; exit code is the signal):
  python3 skip_ci.py gate <platform> <level>
      Exit 0 when bootstrap should stop (skip-all, or that L2/L3 target is off).
      Prints ``skip-all`` or ``skip-l23`` on stdout; logs the decision once.
-  python3 skip_ci.py check-skip-all
-      Exit 0 when only docs/skip-mark changes → bootstrap skips entire CI upload.
-  python3 skip_ci.py check-skip-l2-l3 <platform> <level>
-      Exit 0 when that L2/L3 pipeline should be skipped (CI config YAML-only diffs).
-  python3 skip_ci.py print-annotate
-      Print decision message for Buildkite annotation.
 """

 from __future__ import annotations
@@ -258,11 +253,10 @@ class CiDecision:

@dataclass(frozen=True)
 class CiContext:
-    """Skip-ci decision plus git diff inputs used to compute it."""
+    """Skip-ci decision plus changed files used for pipeline filtering."""

    decision: CiDecision
    changed_files: list[str] | None
-    diff_range: str | None


 def _log(message: str) -> None:
@@ -332,7 +326,7 @@ def resolve_ci_context_from_git() -> CiContext:
    diff_range = _resolve_diff_range()
    changed_files = _changed_files_for_diff_range(diff_range)
    decision = resolve_ci_decision(changed_files, diff_range=diff_range)
-    return CiContext(decision=decision, changed_files=changed_files, diff_range=diff_range)
+    return CiContext(decision=decision, changed_files=changed_files)


 # --- Doc / skip-mark helpers ---
@@ -571,23 +565,6 @@ def main() -> int:
    gate_parser.add_argument("platform", help="cuda, amd, intel, or npu")
    gate_parser.add_argument("level", choices=("l2", "l3"), help="test level")

-    subparsers.add_parser(
-        "check-skip-all",
-        help="exit 0 when docs/skip-mark-only diff should skip entire CI upload",
-    )
-
-    skip_parser = subparsers.add_parser(
-        "check-skip-l2-l3",
-        help="exit 0 when PLATFORM l2/l3 pipeline should be skipped (CI config YAML-only diffs)",
-    )
-    skip_parser.add_argument("platform", help="cuda, amd, intel, or npu")
-    skip_parser.add_argument("level", choices=("l2", "l3"), help="test level")
-
-    subparsers.add_parser(
-        "print-annotate",
-        help="print human-readable skip reason for Buildkite annotation",
-    )
-
    args = parser.parse_args()
    decision = resolve_ci_context_from_git().decision

@@ -600,16 +577,6 @@ def main() -> int:
            return 0
        return 1

-    if args.command == "check-skip-all":
-        return 0 if decision.skip_all else 1
-
-    if args.command == "check-skip-l2-l3":
-        return 0 if not decision.is_run(args.platform, args.level) else 1
-
-    if args.command == "print-annotate":
-        print(decision.message or "CI will run normally")
-        return 0
-
    return 1


@@ -12,8 +12,19 @@ Bootstrap mode (``bootstrap-upload-steps.yml``):

 Test pipeline mode (e.g. test-merge.yml):
  - Drop steps whose ``source_file_dependencies`` do not match changed files.
-  - Expand uploader-only ``mirror_hardwares: l4_1`` into ``agents`` (+ optional ``image``
+  - Expand uploader-only ``mirror_hardwares`` into ``agents`` (+ optional ``image``
    for NPU) + ``plugins`` (see ci_mirror_hardwares.yml).
+  - Omit ``mirror_hardwares`` to compose ``{chip}_{n}`` from pytest ``-m`` SKU
+    markers plus ``cards_n`` (max of several positives; ``not cards_1`` with no
+    positive ``cards_*`` uses that chip's highest existing preset). An explicit
+    preset string (``h100_2``) is used as-is (marks are ignored). Names that
+    are not keys in ``ci_mirror_hardwares.yml`` fail the upload.
+
+    Inferred chip: unset ``MIRROR_HW`` matches ``H100`` or ``L4`` in ``-m``
+    (both → H100); no match skips the step. ``MIRROR_HW=b200`` matches
+    ``B200`` in ``-m`` (otherwise skipped). ``-m`` is not rewritten.
+    ``MIRROR_HW`` must be empty or ``b200``; unknown values fail the upload.
+    A CUDA preset string such as ``h100_4`` is omitted when ``MIRROR_HW=b200``.

 Usage:
  python3 upload_pipeline.py [--upload] [--all | --e2e] <pipeline.yml>
@@ -25,11 +36,13 @@ from __future__ import annotations

 import argparse
 import copy
+import os
+import re
 import subprocess
 import sys
 from functools import lru_cache
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal

 try:
    import yaml
@@ -45,6 +58,14 @@ from skip_ci import (
    resolve_ci_context_from_git,
 )

+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from tests.helpers.mark import (  # noqa: E402
+    get_skus_for_platform,
+    get_supported_card_counts,
+)
+
 # --- Constants ---

 LOG = "upload_pipeline"
@@ -233,32 +254,257 @@ def _load_mirror_hardwares() -> dict[str, dict[str, Any]]:
    return presets


-def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any]:
-    """Replace uploader-only ``mirror_hardwares`` with preset fields from ci_mirror_hardwares.yml."""
-    hardware = step.get("mirror_hardwares")
-    if hardware is None:
-        return step
+@lru_cache(maxsize=1)
+def _cuda_mirror_chips() -> tuple[str, ...]:
+    """Lowercase CUDA SKUs that have at least one ``{chip}_{n}`` preset.

-    if not isinstance(hardware, str) or not hardware.strip():
+    Longer tokens first so ``h100`` is not matched as a prefix of ``h200``.
+    """
+    skus = get_skus_for_platform("cuda")
+    known = _load_mirror_hardwares()
+    chips = {
+        sku.lower() for sku in skus if any(name == sku.lower() or name.startswith(f"{sku.lower()}_") for name in known)
+    }
+    return tuple(sorted(chips, key=lambda chip: (-len(chip), chip)))
+
+
+def _read_cards_marks(expr: str) -> tuple[set[int], bool]:
+    """Parse registered ``cards_n`` / ``not cards_n`` from a pytest ``-m`` expr.
+
+    Returns positive card counts and whether any ``not cards_*`` is present.
+    """
+    counts = get_supported_card_counts()
+    alt = "|".join(str(n) for n in sorted(counts, key=lambda n: (-len(str(n)), -n))) if counts else "(?!)"
+    positives: set[int] = set()
+    has_not_cards = False
+    for match in re.finditer(rf"\b(?:(not)\s+)?cards_({alt})\b", expr):
+        if match.group(1):
+            has_not_cards = True
+        else:
+            positives.add(int(match.group(2)))
+    return positives, has_not_cards
+
+
+def _read_hardware_marks(expr: str) -> set[str]:
+    """Return lowercase positive CUDA SKUs that have a mirror preset.
+
+    ``not H100`` is ignored. SKUs with no ``{chip}_*`` preset (e.g. H200) are dropped.
+    """
+    skus = sorted(get_skus_for_platform("cuda"), key=lambda s: (-len(s), s))
+    if not skus:
+        return set()
+    alt = "|".join(re.escape(sku) for sku in skus)
+    mirror = set(_cuda_mirror_chips())
+    chips: set[str] = set()
+    for match in re.finditer(rf"\b(?:(not)\s+)?({alt})\b", expr):
+        if match.group(1):
+            continue
+        chip = match.group(2).lower()
+        if chip in mirror:
+            chips.add(chip)
+    return chips
+
+
+# CI policy: not pytest-mark registry.
+_SUPPORTED_MIRROR_HW_SELECTORS = frozenset({"b200"})
+# Unset MIRROR_HW: match H100 then L4 in -m (skip if neither).
+_DEFAULT_INFER_CHIPS = ("h100", "l4")
+_PYTEST_MARKER_ARG = re.compile(r"-m\s+(?:\"([^\"]*)\"|'([^']*)'|(\S+))")
+_CARDS_CHIP_MAX: Literal["max"] = "max"
+
+
+def _get_mirror_hw_selector() -> str:
+    """Return lowercase ``MIRROR_HW``, or empty to keep string presets / marker chips.
+
+    A set value must be ``b200`` (case-insensitive). Unknown values fail closed
+    so a typo cannot silently drop the CUDA pipeline.
+    """
+    selector = os.environ.get("MIRROR_HW", "").strip().lower()
+    if not selector:
+        return ""
+    if selector not in _SUPPORTED_MIRROR_HW_SELECTORS:
+        allowed = " or ".join(repr(s) for s in sorted(_SUPPORTED_MIRROR_HW_SELECTORS))
+        raise ValueError(f"unsupported MIRROR_HW={selector!r}; expected empty or {allowed}")
+    return selector
+
+
+def _read_pytest_marks(commands: Any) -> tuple[set[str], int | Literal["max"] | None]:
+    """Read pytest ``-m`` from step commands: CUDA SKU chips and cards count.
+
+    *chips* are lowercase SKUs from positive markers (``not H100`` is ignored).
+    *cards* is the max positive ``cards_n``, ``"max"`` when only ``not cards_*``
+    is present, or ``None``.
+    """
+    chunks: list[str] = []
+
+    def _collect(value: Any) -> None:
+        if value is None:
+            return
+        if isinstance(value, str):
+            chunks.append(value)
+        elif isinstance(value, list):
+            for part in value:
+                _collect(part)
+        else:
+            chunks.append(str(value))
+
+    _collect(commands)
+    text = "\n".join(chunks)
+
+    chips: set[str] = set()
+    positives: set[int] = set()
+    has_not_cards = False
+    for match in _PYTEST_MARKER_ARG.finditer(text):
+        expr = match.group(1) or match.group(2) or match.group(3) or ""
+        chips |= _read_hardware_marks(expr)
+        card_counts, not_cards = _read_cards_marks(expr)
+        positives.update(card_counts)
+        has_not_cards = has_not_cards or not_cards
+
+    if positives:
+        cards: int | Literal["max"] | None = max(positives)
+    elif has_not_cards:
+        cards = _CARDS_CHIP_MAX
+    else:
+        cards = None
+    return chips, cards
+
+
+def _compose_mirror_hardware_name(
+    chips: set[str],
+    cards: int | Literal["max"] | None,
+    *,
+    step_label: str,
+) -> str | None:
+    """Assemble ``{chip}_{n}`` from pytest ``-m`` chips/cards, or None to skip.
+
+    Unset ``MIRROR_HW``: ``H100`` if in ``-m``, else ``L4``, else skip.
+    Set ``MIRROR_HW``: that chip if in ``-m``, else skip.
+    ``cards`` is a registered ``cards_n``, ``"max"`` for that chip's highest
+    existing preset, or None (error after a chip matched — no ``cards_*``).
+    """
+    selector = _get_mirror_hw_selector()
+    if selector:
+        chip = selector if selector in chips else None
+        skip_reason = f"MIRROR_HW={selector!r} not in -m SKUs {sorted(chips)}"
+    else:
+        chip = next((c for c in _DEFAULT_INFER_CHIPS if c in chips), None)
+        skip_reason = f"no H100 or L4 in -m SKUs {sorted(chips)}"
+    if chip is None:
+        _log(f"skip {step_label}: {skip_reason}")
+        return None
+    if cards is None:
        raise ValueError(
-            f"mirror_hardwares must be a non-empty string in step {_get_step_label(step)!r}",
+            f"step {step_label!r} has a pytest -m SKU marker but no cards_* "
+            f"(or not cards_*); add cards_n or set mirror_hardwares to force a preset",
        )

-    preset = _load_mirror_hardwares().get(hardware)
+    count_source = "not cards_*" if cards == _CARDS_CHIP_MAX else f"cards_{cards}"
+    count = None if cards == _CARDS_CHIP_MAX else cards
+    registered = get_supported_card_counts()
+    if count is not None and count not in registered:
+        supported = ", ".join(f"cards_{n}" for n in sorted(registered)) or "(none)"
+        raise ValueError(
+            f"{count_source} in step {step_label!r} is not a registered cards_* mark; supported: {supported}",
+        )
+    known = _load_mirror_hardwares()
+    if count is None:
+        available = [n for n in sorted(registered) if f"{chip}_{n}" in known]
+        if not available:
+            raise ValueError(
+                f"{count_source} in step {step_label!r} has no CUDA count presets for {chip}",
+            )
+        count = max(available)
+    chosen = f"{chip}_{count}"
+    if chosen not in known:
+        raise ValueError(
+            f"{count_source} in step {step_label!r} has no preset {chosen!r}",
+        )
+    _log(f"{step_label}: {count_source} → {chosen!r}")
+    return chosen
+
+
+def _resolve_mirror_hardware_name(hardware: Any, *, step_label: str) -> str | None:
+    """Map ``mirror_hardwares`` to a ``ci_mirror_hardwares.yml`` preset, or None to skip.
+
+    The value is a preset name (``h100_2``). Unknown names fail later when the
+    registry is loaded. CUDA names are omitted when ``MIRROR_HW`` names a
+    different chip; NPU strings ignore ``MIRROR_HW``.
+    """
+    if isinstance(hardware, bool):
+        raise ValueError(f"mirror_hardwares must not be a boolean in step {step_label!r}")
+    if isinstance(hardware, str):
+        name = hardware.strip()
+    elif isinstance(hardware, int):
+        name = str(hardware)
+    else:
+        raise ValueError(
+            f"mirror_hardwares must be a preset string in step {step_label!r}",
+        )
+    if not name:
+        raise ValueError(f"mirror_hardwares must be a non-empty string in step {step_label!r}")
+
+    selector = _get_mirror_hw_selector()
+    token = name.lower()
+    chip = next(
+        (c for c in _cuda_mirror_chips() if token == c or token.startswith(f"{c}_")),
+        None,
+    )
+    if selector and chip is not None and chip != selector:
+        _log(
+            f"skip {step_label}: preset {name!r} is {chip} hardware; MIRROR_HW={selector!r}",
+        )
+        return None
+    return name
+
+
+def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None:
+    """Replace ``mirror_hardwares`` (or infer it from ``-m``) with preset agents/plugins.
+
+    Explicit preset string: used as-is; pytest marks are ignored.
+    Omitted key: match H100/L4 (or ``MIRROR_HW``) in ``-m``, then compose
+    ``{chip}_{n}`` from ``cards_n``. No match skips the step. Steps with no
+    SKU/cards and no ``mirror_hardwares`` are left unchanged (CPU jobs).
+
+    Returns None when the current ``MIRROR_HW`` selector does not apply.
+    """
+    step_label = _get_step_label(step)
+    has_pool = step.get("agents") is not None or step.get("plugins") is not None or step.get("image") is not None
+
+    if "mirror_hardwares" not in step:
+        if has_pool:
+            return step
+        chips, cards = _read_pytest_marks(step.get("commands"))
+        if not chips and cards is None:
+            return step
+        preset_name = _compose_mirror_hardware_name(chips, cards, step_label=step_label)
+    else:
+        if has_pool:
+            raise ValueError(
+                f"step {step_label!r} sets mirror_hardwares together with agents/plugins/image; "
+                f"use mirror_hardwares only",
+            )
+        preset_name = _resolve_mirror_hardware_name(
+            step.get("mirror_hardwares"),
+            step_label=step_label,
+        )
+
+    if preset_name is None:
+        return None
+
+    preset = _load_mirror_hardwares().get(preset_name)
    if preset is None:
        known = ", ".join(sorted(_load_mirror_hardwares()))
        raise ValueError(
-            f"unknown mirror_hardwares {hardware!r} in step {_get_step_label(step)!r}; known: {known}",
-        )
-
-    if step.get("agents") is not None or step.get("plugins") is not None or step.get("image") is not None:
-        raise ValueError(
-            f"step {_get_step_label(step)!r} sets mirror_hardwares together with agents/plugins/image; "
-            "use mirror_hardwares only",
+            f"unknown mirror_hardwares {preset_name!r} in step {step_label!r}; known: {known}",
        )

    expanded = copy.deepcopy(preset)
-    return {key: value for key, value in step.items() if key != "mirror_hardwares"} | expanded
+    merged = {key: value for key, value in step.items() if key != "mirror_hardwares"} | expanded
+    # Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out.
+    if "retry" in step:
+        merged["retry"] = step["retry"]
+    return merged


 def _match_source_file(changed_files: list[str], prefixes: list[str]) -> bool:
@@ -297,7 +543,7 @@ def _process_test_steps(
        nested = step.get("steps")
        if nested is not None:
            kept_nested = _process_test_steps(nested, changed_files)
-            if changed_files is not None and not kept_nested:
+            if not kept_nested:
                _log(f"omit empty group {_get_step_label(step)!r}")
                continue
            new_step = {key: value for key, value in step.items() if key != "source_file_dependencies"}
@@ -305,14 +551,11 @@ def _process_test_steps(
            processed.append(new_step)
            continue

-        if deps is not None:
-            processed.append(
-                _expand_mirror_hardwares(
-                    {key: value for key, value in step.items() if key != "source_file_dependencies"},
-                ),
-            )
-        else:
-            processed.append(_expand_mirror_hardwares(step))
+        leaf = {key: value for key, value in step.items() if key != "source_file_dependencies"}
+        expanded = _expand_mirror_hardwares(leaf)
+        if expanded is None:
+            continue
+        processed.append(expanded)

    return processed

@@ -338,6 +581,10 @@ def _render_test_pipeline(
    e2e_only: bool = False,
 ) -> dict[str, Any]:
    """Filter steps by PR diff and strip uploader-only ``source_file_dependencies`` metadata."""
+    # Validate MIRROR_HW once up front. Per-step skip must not run for typos
+    # (e.g. b20o), or a pipeline can shrink to leftover CPU-only steps and
+    # still upload successfully.
+    _get_mirror_hw_selector()
    steps = doc.get("steps")
    if not isinstance(steps, list):
        return doc
@@ -367,10 +614,7 @@ def _render_pipeline(

    text = path.read_text(encoding="utf-8")
    ctx = resolve_ci_context_from_git()
-    if force_all or e2e_only:
-        changed_files = None
-    else:
-        changed_files = ctx.changed_files
+    changed_files = None if force_all or e2e_only else ctx.changed_files

    doc = yaml.safe_load(text)
    if not isinstance(doc, dict):
@@ -1,3 +1,6 @@
+# CUDA CI L3 (`advanced_model` e2e) plus L1 CPU Simple Test (`core_model and cpu`).
+# Uploaded by Upload Merge Pipeline: PR label `merge-test`, or post-merge on main.
+
 env:
  VLLM_WORKER_MULTIPROC_METHOD: spawn
  VLLM_USE_FLASHINFER_MOE_FP16: 0
@@ -1,3 +1,7 @@
+# CUDA CI L4 (`full_model`): nightly and release share this file.
+# Uploaded by Upload Nightly Pipeline: NIGHTLY=1, or PR labels (nightly-test / omni-test / ...).
+# Hardware: omit mirror_hardwares to compose from -m SKU + cards_n; or set a preset string to force.
+
 env:
  VLLM_WORKER_MULTIPROC_METHOD: spawn
  VLLM_USE_FLASHINFER_MOE_FP16: 0
@@ -13,16 +17,14 @@ steps:
    steps:

      - label: ":full_moon: Omni · Function Test with H100 · Single-GPU"
-        timeout_in_minutes: 90
+        timeout_in_minutes: 120
        commands:
-          - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_1" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
-        mirror_hardwares: h100_1
+          - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and B200 and omni and cards_1" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py

      - label: ":full_moon: Omni · Function Test with H100 · 2-GPU"
        timeout_in_minutes: 90
        commands:
-          - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
-        mirror_hardwares: h100_2
+          - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and B200 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py

      - label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
        timeout_in_minutes: 50
@@ -30,11 +32,11 @@ steps:
          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model"
        mirror_hardwares: h100_1

-      - label: ":full_moon: Omni · Doc Test with L4"
+      - label: ":full_moon: Omni · Doc Test with L4 · 4-GPU"
        timeout_in_minutes: 90
        commands:
          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model"
+          - pytest -s -v tests/examples/ -m "full_model and omni and L4 and cards_4" --run-level "full_model"
        mirror_hardwares: l4_4

      - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU"
@@ -178,7 +180,7 @@ steps:
      build.pull_request.labels includes "nightly-test"
    steps:

-      - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations)"
+      - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 2-GPU"
        timeout_in_minutes: 30
        # NOTE: Single-GPU (core_model) tiny model tests run on the ready label,
        # so nightly only adds the multi-GPU parallelism configurations.
@@ -188,7 +190,14 @@ steps:
        # for the CI, since these tests should pass irrespective of whether or not we use
        # the full weights or not.
        commands:
-          - timeout 25m pytest -sv tests/model_tests/diffusion -m 'full_model and cuda' --run-level "core_model"
+          - pytest -sv tests/model_tests/diffusion -m 'full_model and cuda and cards_2' --run-level "core_model"
+        mirror_hardwares: l4_2
+
+      - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 4-GPU"
+        timeout_in_minutes: 30
+        # CFG+TP extra_test_groups request 2×2=4 devices (see get_required_device_count).
+        commands:
+          - pytest -sv tests/model_tests/diffusion -m 'full_model and cuda and cards_4' --run-level "core_model"
        mirror_hardwares: l4_4

  - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
@@ -220,10 +229,21 @@ steps:
            --run-level "full_model"
        mirror_hardwares: h100_2

-      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · Multi-GPU"
+      - label: ":full_moon: Diffusion X2I(&A&T) · Boogu Edit CFG · 3-GPU"
        timeout_in_minutes: 120
        commands:
-          - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and not cards_1" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - >-
+            pytest -sv
+            tests/e2e/online_serving/test_boogu_image_edit_expansion.py
+            -k "test_double_guidance_cfg_parallel and cfg_parallel_3_full"
+            -m "slow and diffusion and cuda and H100 and cards_3"
+            --run-level "full_model"
+        mirror_hardwares: h100_3
+
+      - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · 4-GPU"
+        timeout_in_minutes: 120
+        commands:
+          - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and cards_4" --run-level "full_model" --ignore=tests/e2e/accuracy
        mirror_hardwares: l4_4

      - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
@@ -259,7 +279,6 @@ steps:
        timeout_in_minutes: 180
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
          - export CACHE_DIT_VERSION=1.5.0
          - |
            set +e
@@ -362,7 +381,6 @@ steps:
        timeout_in_minutes: 180
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json -m "H100 and cards_1"
@@ -412,7 +430,6 @@ steps:
        timeout_in_minutes: 180
        commands:
          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
          - |
            set +e
            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json
@@ -427,10 +444,22 @@ steps:
    depends_on: upload-nightly-pipeline
    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
    steps:
-      - label: ":full_moon: Diffusion · Distributed Test with L4"
+      - label: ":full_moon: Diffusion · Distributed Test with L4 · 2-GPU"
        timeout_in_minutes: 60
        commands:
-          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
+          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_2' --run-level "full_model"
+        mirror_hardwares: l4_2
+
+      - label: ":full_moon: Diffusion · Distributed Test with L4 · 3-GPU"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_3' --run-level "full_model"
+        mirror_hardwares: l4_3
+
+      - label: ":full_moon: Diffusion · Distributed Test with L4 · 4-GPU"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_4' --run-level "full_model"
        mirror_hardwares: l4_4

      - label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU"
@@ -1,3 +1,6 @@
+# CUDA CI L2 (`core_model` e2e) plus L1 CPU Simple Test (`core_model and cpu`).
+# Uploaded by Upload Ready Pipeline: PR label `ready`, or main+WEEKLY=1 with --e2e.
+
 env:
  VLLM_WORKER_MULTIPROC_METHOD: spawn
  VLLM_USE_FLASHINFER_MOE_FP16: 0
@@ -51,17 +54,20 @@ steps:
          - timeout 15m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
        mirror_hardwares: l4_1

-      - label: "Diffusion · Distributed & Attention Test"
-        source_file_dependencies:
-          - tests/diffusion/distributed/
-          - tests/diffusion/attention/test_attention_sp.py
-          - vllm_omni/diffusion/distributed/
-          - vllm_omni/diffusion/attention/
-          - vllm_omni/diffusion/data.py
-          - vllm_omni/diffusion/models/
+      - label: "Diffusion · Distributed Test · Single-GPU"
        commands:
-          - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda' --run-level "core_model"
-          - timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda' --run-level "core_model"
+          - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_1' --run-level "core_model"
+        mirror_hardwares: l4_1
+
+      - label: "Diffusion · Distributed Test · 2-GPU"
+        commands:
+          - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_2' --run-level "core_model"
+        mirror_hardwares: l4_2
+
+      - label: "Diffusion · Distributed & Attention Test · 4-GPU"
+        commands:
+          - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_4' --run-level "core_model"
+          - timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda and L4 and cards_4' --run-level "core_model"
        mirror_hardwares: l4_4

      - label: "Diffusion · Model Test"
@@ -1,3 +1,6 @@
+# CUDA CI L5 (`slow` / reliability / extra coverage).
+# Uploaded by Upload Weekly Pipeline: WEEKLY=1 / NON_CRITICAL=1, or PR label `weekly-test`.
+
 env:
  VLLM_WORKER_MULTIPROC_METHOD: spawn
  VLLM_USE_FLASHINFER_MOE_FP16: 0
@@ -111,17 +114,24 @@ steps:
    if: build.env("NON_CRITICAL") == "1"
    steps:

-      # Qwen2.5 expansion (CUDA 4-card / online 2-card) stays on this l4_4 job.
-      - label: "Omni · L4 · Multi-GPU"
+      # Online Qwen2.5 expansion is 2-card; offline expansion + AutoRound are 4-card.
+      - label: "Omni · L4 · 2-GPU"
        timeout_in_minutes: 90
        commands:
          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and not cards_1" --run-level "core_model"
+          - pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and cards_2" --run-level "core_model"
+        mirror_hardwares: l4_2
+
+      - label: "Omni · L4 · 4-GPU"
+        timeout_in_minutes: 90
+        commands:
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and cards_4" --run-level "core_model"
          - >-
            pytest -s -v tests/e2e/
            --ignore=tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
            --ignore=tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
-            -m "slow and L4 and omni and not cards_1"
+            -m "slow and L4 and omni and cards_4"
            --run-level "full_model"
        mirror_hardwares: l4_4

@@ -194,12 +204,21 @@ steps:
            --run-level "full_model"
        mirror_hardwares: l4_1

-      - label: "Diffusion · L4 · Multi-GPU"
+      - label: "Diffusion · L4 · 2-GPU"
        timeout_in_minutes: 60
        commands:
          - >-
            pytest -sv
-            tests/e2e/ -m "slow and diffusion and L4 and not cards_1"
+            tests/e2e/ -m "slow and diffusion and L4 and cards_2"
+            --run-level "full_model"
+        mirror_hardwares: l4_2
+
+      - label: "Diffusion · L4 · 4-GPU"
+        timeout_in_minutes: 60
+        commands:
+          - >-
+            pytest -sv
+            tests/e2e/ -m "slow and diffusion and L4 and cards_4"
            --run-level "full_model"
        mirror_hardwares: l4_4

@@ -0,0 +1,458 @@
+---
+name: production-add-diffusion-model
+description: Productionize a vLLM-Omni diffusion model after its Day-0 vertical slice works. Use when work requires official API parity and input limits, feature-combination evidence, online FP8, distributed layerwise offload, per-request Cache-DiT, CUDA/ROCm/NPU/XPU recipes, faster or sparse attention, fused operators, USP or disaggregation analysis, continuous batching and abort handling, long-running RPS stability, or accuracy/performance/reliability CI. For initial architecture porting, registry wiring, and basic weight loading, use add-diffusion-model first.
+---
+
+# Productionizing a Diffusion Model
+
+## Scope
+
+Use this skill to move a working model integration from Day-0 support to a
+measured, fail-closed production deployment. It is deliberately separate from
+[`add-diffusion-model`](../add-diffusion-model/SKILL.md); do not modify or copy
+that skill when productionizing a model. If registry, pipeline, or basic loader
+work is incomplete, read that skill first and finish the Day-0 vertical slice.
+
+MiniMax-H3 [PR #5691](https://github.com/vllm-project/vllm-omni/pull/5691)
+is the merged Day-0 case study. Its open optimization
+[issue #5700](https://github.com/vllm-project/vllm-omni/issues/5700) is an
+actively maintained roadmap and evidence map, not a list of already supported
+features. It can lag merged implementation: read its current revision, cited
+PRs, target source, and maintained recipes, then leave every unverified row
+`not tested` until the target revision has scoped evidence.
+
+Read these references before the corresponding gate:
+
+- [API and hardware recipes](references/api-and-recipes.md)
+- [Production feature patterns](references/feature-patterns.md)
+- [Performance patterns and code examples](references/performance-patterns.md)
+- [Serving, soak, and CI validation](references/production-validation.md)
+
+## Non-negotiable evidence contract
+
+Create a matrix whose row key is at least:
+
+```text
+task x API mode x execution mode/capacity x shape/schedule x attention backend x
+packed layout x cache policy x quantization x offload mode x topology x
+hardware x dtype x output representation/transport
+```
+
+Use exactly these states:
+
+| State | Meaning |
+|---|---|
+| `validated` | The exact row passed correctness and measured deployment tests, with reproducible artifacts. |
+| `limited` | A bounded subset passed; the limitation and rejection/fallback behavior are explicit. |
+| `unsupported` | A stable architectural/platform restriction is proven and fails early with an actionable error. |
+| `not tested` | No adequate evidence exists. This is the default. |
+
+Never infer compatibility from import success, server startup, another task,
+another card, or another topology. Record the model/checkpoint revision,
+vLLM-Omni commit, commands, request assets and hashes, seed/schedule/shape,
+raw output artifacts, environment, and result for every `validated` row.
+
+## Workflow
+
+Follow the gates in order. Keep a dense BF16 single-device path as the
+correctness oracle until every advertised fast path has parity evidence.
+
+### Gate 0: Freeze the official contract
+
+Pin the official implementation and checkpoint revision. Build a matrix for:
+
+- constructor/call arguments, defaults, scheduler and sigma schedule;
+- task-to-source modality/count/order rules;
+- prompt, MIME, bytes, pixels, frames, duration, FPS, steps, seed, guidance,
+  output count, and output format limits;
+- offline API and every public sync/async serving endpoint;
+- checkpoint partitions and task selection.
+
+Normalize offline, sync, and async requests through one validation contract.
+Reject an invalid task/source combination before downloads, upload persistence,
+temporary files, async job creation, or engine submission. Use 400 for invalid
+semantics and 413 for payload limits; clean partial resources on every failure.
+See [API and hardware recipes](references/api-and-recipes.md).
+
+### Gate 1: Re-prove correctness under production shapes
+
+1. Make loading strict: track expected, loaded, unexpected, and intentionally
+   ignored tensors. A fused destination is complete only after every source
+   shard arrives. Missing retained parameters abort startup.
+2. Fix reference revision, asset hashes, seed, schedule, and shape. Compare
+   component outputs, intermediate denoise latents, and final artifacts.
+3. Test unequal packed samples. Verify `cu_seqlens`, padding exclusion,
+   modality/CFG ownership, split/gather boundaries, and one output per request.
+4. Test partial process groups with non-member ranks and group size smaller
+   than world size. Do not construct a private group per component or request.
+5. Run negative API cases before positive E2E tests to prove side-effect-free
+   rejection and cleanup.
+
+### Gate 2: Wire shared fast primitives
+
+Prefer explicit, testable vLLM-Omni/vLLM operators with native fallbacks.
+`torch.compile` remains a separately validated regional optimization; it is
+not a substitute for correct operator selection or lifecycle-safe boundaries.
+
+#### RMSNorm + RoPE baseline
+
+Use the shared cross-platform operators, preserving the official Q/K norm and
+RoPE order:
+
+```python
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.layers.rope import (
+    RotaryEmbedding,
+    apply_rope_to_qk,
+)
+
+self.norm_q = RMSNorm(head_dim, eps=eps)
+self.norm_k = RMSNorm(head_dim, eps=eps)
+self.rope = RotaryEmbedding(is_neox_style=False)
+
+query = self.norm_q(query)  # [B, S, Hq, D] or packed [T, Hq, D]
+key = self.norm_k(key)      # [B, S, Hkv, D] or packed [T, Hkv, D]
+query, key = apply_rope_to_qk(
+    self.rope,
+    query,
+    key,
+    (cos, sin),
+)
+```
+
+Verify `is_neox_style`, `half_head_dim`, partial rotary dimensions, cos/sin
+layout, and packed row ownership against the official model. `False` means
+interleaved/GPT-J style; do not assume the default is NeoX. Shared RMSNorm and
+RoPE are two dispatched operators, not automatically one fused kernel.
+
+For packed non-interleaved RoPE, use the shared fused boundary when its contract
+matches the model:
+
+```python
+from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
+
+query, key = fused_qk_norm_rope(
+    query,                    # [T, Hq, D]
+    key,                      # [T, Hkv, D]
+    self.norm_q.weight,       # [D]
+    self.norm_k.weight,       # [D]
+    rope_table,               # [T, rotary_dim] = [cos | sin]
+    self.norm_q.variance_epsilon,
+)
+```
+
+The current CUDA fast path specializes BF16 `head_dim=128`, `rotary_dim=96`;
+the public function keeps an eager fallback for unsupported inputs. Verify the
+packed row/frequency contract and trace the actual fast path. Do not reshape an
+incompatible official RoPE layout merely to enter this kernel.
+
+MiniMax-H3-style partial NeoX RoPE rotates a prefix and passes the suffix
+through:
+
+```python
+self.q_norm = RMSNorm(head_dim, eps=eps, dtype=torch.bfloat16)
+self.k_norm = RMSNorm(head_dim, eps=eps, dtype=torch.bfloat16)
+self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
+
+def apply_partial_rope(self, x, freqs, rot_dim):
+    x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
+    cos, sin = torch.cos(freqs).to(x.dtype), torch.sin(freqs).to(x.dtype)
+    return torch.cat((self.rope(x_rot, cos, sin), x_pass), dim=-1)
+```
+
+The partial-RoPE snippet assumes `import torch` and lives on the owning
+attention module; adapt names and validated frequency shapes rather than
+copying it at module scope.
+
+Call the modules normally so `CustomOp` selects CUDA/HIP/NPU/XPU/native
+implementations. Do not call `forward_cuda()` directly. Compare the shared
+path with native BF16 at operator, block, denoise-trajectory, and artifact
+levels before claiming either accuracy or speed.
+
+#### Faster BF16 attention and fused projections
+
+Use local TP head counts and role-aware shared attention:
+
+```python
+from vllm.model_executor.layers.linear import QKVParallelLinear
+from vllm_omni.diffusion.attention.layer import Attention
+
+self.qkv_proj = QKVParallelLinear(
+    hidden_size=hidden_size,
+    head_size=head_dim,
+    total_num_heads=num_heads,
+    total_num_kv_heads=num_kv_heads,
+    bias=False,
+    quant_config=quant_config,
+    prefix=f"{prefix}.qkv_proj",
+)
+self.attn = Attention(
+    num_heads=self.qkv_proj.num_heads,
+    num_kv_heads=self.qkv_proj.num_kv_heads,
+    head_size=head_dim,
+    softmax_scale=head_dim**-0.5,
+    causal=False,
+    qkv_layout="BSND",
+    role="self",
+    prefix=prefix,
+)
+```
+
+Prove the selected backend from runtime logs/trace and compare it with dense
+SDPA. Packed metadata must express real sample boundaries; do not build masks
+that the selected backend ignores. A backend that is faster for one shape or
+platform is not a global default. See [performance patterns](references/performance-patterns.md)
+for QKV, SwiGLU, AdaLN, sparse attention, redundancy, and fusion examples.
+
+### Gate 3: Add online FP8 without breaking the loader
+
+Route the active config to each component, then pass a stable runtime prefix
+and config into every quantizable vLLM linear:
+
+```python
+def resolve_component_quant_config(quant_config, component):
+    return quant_config.resolve(component) if hasattr(quant_config, "resolve") else quant_config
+
+transformer_quant_config = resolve_component_quant_config(
+    od_config.quantization_config,
+    "transformer",
+)
+self.transformer = MyDiT(
+    od_config,
+    quant_config=transformer_quant_config,
+    prefix="transformer",
+)
+
+self.qkv_proj = QKVParallelLinear(
+    hidden_size=hidden_size,
+    head_size=head_dim,
+    total_num_heads=num_heads,
+    total_num_kv_heads=num_kv_heads,
+    quant_config=quant_config,
+    prefix=f"{prefix}.qkv_proj",
+)
+```
+
+Use `build_quant_config()` for direct construction/testing:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+quant_config = build_quant_config({
+    "transformer": {"method": "fp8"},
+    "text_encoder": None,
+    "vae": None,
+    "default": None,
+})
+```
+
+Preserve each parameter's vLLM `weight_loader`; perform checkpoint layout
+conversion before invoking it. Account for every QKV/gate-up source shard.
+Keep precision-sensitive norm/modulation/embedders in BF16 unless separately
+proven, using the current `safe_quant_config` pattern where applicable.
+
+Validate resident DiT FP8 first: prove which modules are quantized, HBM saved,
+BF16-vs-FP8 trajectory/artifact quality, latency, and throughput. Text encoder,
+VAE, other hardware, TP/cache/DLO combinations are separate rows.
+Pre-quantized, pruned, rotated, or adapter-modified checkpoints are separate
+model/loading lanes; compare them with the released BF16 model without calling
+the result a pure quantization ablation. Read
+[`../quantization/SKILL.md`](../quantization/SKILL.md) if present and prefer the
+target revision's implementation/docs when sibling guidance differs.
+
+### Gate 4: Make every component offloadable
+
+Declare topology instead of relying on heuristic discovery:
+
+```python
+from typing import ClassVar
+
+import torch.nn as nn
+
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.offloader import OffloadPlan
+
+class MyPipeline(nn.Module, SupportsComponentDiscovery):
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+    _resident_modules: ClassVar[list[str]] = []
+    _offload_plan: ClassVar[OffloadPlan] = OffloadPlan(
+        on_demand_component_paths=frozenset({"text_encoder", "vae"}),
+        block_attrs={"transformer": ("blocks",)},
+        offload_submodules={"token_refiner": "blocks"},
+        resident_dit_paths=frozenset({"transformer"}),
+        encoder_block_attrs={"text_encoder": ("encoder.layers",)},
+    )
+```
+
+Add a test that resolves every dotted component/block path and fails if it is
+missing or is not the expected module/container; discovery may otherwise warn
+and skip. Validate ordinary layerwise offload and distributed layerwise
+offload (DLO) independently.
+
+Keep component lifecycle and host-weight storage separate. `OffloadPlan`
+declares what can be streamed; the diffusion loader owns any `HostWeightPlan`
+and hands the exact prevalidated plan to DLO. Do not make the offloader rescan
+checkpoint files or duplicate loader name/shape/dtype decisions. Direct mmap is
+a fail-closed optimization: current preflight requires TP1, no HSDP, no online
+quantization, complete bindings, and represented transforms. Otherwise the
+ordinary loader remains authoritative.
+
+If the target revision supports Host Weight Runtime (HWR), treat it as a
+separate, exact-identity startup/cache path rather than generic offload or
+zero-copy execution. Qualify `preferred` population/fallback and `required`
+consume-only failure independently; verify immutable manifests, final-layout
+restore, lease cleanup, corruption/quarantine, and node-local sharing. Current
+model contracts may be BF16 no-AllGather-only, so reject AllGather, online
+quantization, HSDP, LoRA/adapted weights, and unrecognized load formats unless
+the exact target revision explicitly adds them.
+
+On small-HBM cards, record load/materialization peak, resident peak, encode,
+each denoise wave, decode, transient per-rank HBM, host PSS, H2D, and collective
+time. Test DLO AllGather and `--dlo-no-use-allgather` as different deployments.
+Sweep the supported resident-block policies and publish the non-dominated
+latency-HBM frontier; a single minimum-memory or fastest point is not the DLO
+trade-off. Keep request-specific preparation replica-local. Only weight
+materialization may enter the declared DLO collective group; never use WORLD
+for per-request work in a DP deployment.
+Direct checkpoint mmap preflight is limited to TP1 without HSDP or online
+quantization; TP>1 falls back to the ordinary TP-aware loader and can still feed
+DLO AllGather or no-AllGather after scoped E2E. HSDP remains incompatible with
+DLO AllGather. A target revision may allow finalized per-tensor online FP8
+through the ordinary loader and AllGather; keep every other online quantizer
+fail-closed, and require model/card/topology E2E before promotion.
+Use deploy YAML for DP under `--omni`; vLLM DP CLI flags are rejected.
+
+### Gate 5: Add request-scoped acceleration safely
+
+Server-wide Cache-DiT support does not imply a per-request quality API. Only
+add request policy after defining calibrated quality tiers and a safe hook
+lifecycle using `SupportsRequestScopedCacheDiT`, `CacheDiTRequestSpec`, and
+`RequestScopedCacheDiTRuntime`. Validate alternating and concurrent quality
+tiers, real cache hits, refresh/reset, success, error, disconnect, abort, and
+the next uncached request. Do not batch different quality policies unless the
+batch compatibility key and hook ownership make that safe.
+
+Then test Cache-DiT against FP8, DLO, TP/SP, compile, and batching one
+combination at a time. See [production feature patterns](references/feature-patterns.md)
+for the lifecycle skeleton and compatibility matrix.
+
+### Gate 6: Optimize only from profiles
+
+Freeze one canonical workload manifest before profiling. Keep lossless runtime
+A/Bs, accelerated paths with numerical/quality trade-offs, and production
+topology studies in separate result lanes. Reconcile the client boundary as
+queue + encode + denoise + video/audio decode + output transport/codec +
+residual; do not hide a material residual in another stage.
+
+Profile stage time, GPU kernels, CPU/GPU synchronization, allocations, H2D,
+output payload/copies, and collectives before choosing work. Apply and A/B one
+change at a time:
+
+1. Remove synchronization and redundant denoise-loop work: hoist static prompt
+   conditioning/token refinement, RoPE tables, masks/metadata, bounded AdaLN
+   schedule projections, and reference scans; remove per-layer `.item()` and
+   unused masks/allocations.
+2. Select the fastest correct BF16 attention backend for each role and shape.
+3. Prefer explicit fused QKV, gate-up + `SiluAndMul`, RMSNorm/RoPE, AdaLN, and
+   layout/residual operators with guarded native fallbacks. Preserve every
+   materialized dtype/rounding boundary consumed by the unfused graph.
+4. Enable sparse attention only when platform, shape, topology, realized
+   sparsity, quality, fallback, and end-to-end speedup are all proven.
+5. Profile USP all-to-all count/bytes, packing, contiguous copies, overlap, and
+   scaling efficiency. Audit replicated cross-attention before using
+   `skip_sequence_parallel`. Qualify regular and accelerated Ulysses transport
+   independently; activation logs, JIT/readiness time, workspace growth,
+   stream ownership, maximum-shape warmup, and numerical drift are gates.
+   Preserve the checkpoint's Q:KV head ratio when padding GQA for Ulysses, and
+   validate strided QKV staging rather than assuming packed contiguity.
+6. Evaluate text-encoder or VAE disaggregation only if stage share, transfer
+   volume, reuse/concurrency, and independent scaling justify it. Tiling,
+   offload, or patch parallelism is not disaggregation.
+
+Do not combine independent optimizations into one performance claim. Preserve
+raw before/after runs with fixed workload and warmup policy.
+
+### Gate 7: Make serving interruptible and load-stable
+
+Treat request-mode batching and step-wise execution as separate capabilities:
+
+- request batching requires `supports_request_batch = True` and
+  `DiffusionRequestBatch -> list[DiffusionOutput]`, exactly one output per
+  request;
+- step execution requires the complete `SupportsStepExecution` contract:
+  `prepare_encode`, `denoise_step`, `step_scheduler`, and `post_decode`.
+
+Keep all mutable scheduler, RNG, latent, mask, cache, and step state request
+scoped. First validate `--step-execution --max-num-seqs 1` and step-boundary
+abort; only then test heterogeneous multi-request waves. Step continuous
+batching is experimental and is not automatically compatible with Cache-DiT.
+Before recommending it, predeclare a useful case and success threshold, then
+A/B it against request mode under the same arrival process. Structural support
+or a generic bridge is not latency, throughput, or cancellation evidence for a
+model.
+
+Test queued and in-flight abort, client disconnect, timeout, OOM, worker error,
+and restart. Require one terminal result, idempotent cleanup, no post-decode
+after abort, no cache/temp/VRAM/request-ID leak, and a successful next request.
+Inject late worker results/exceptions after cancellation and prove the result
+pump stays alive instead of raising `InvalidStateError`. For large batched or
+long-window video outputs, freeze dtype, range, layout, contiguity, ownership,
+payload size, serializer/codec limits, and offline-versus-HTTP semantics.
+Validate device-side output preparation, D2H/IPC or shared handles, remote
+encoding, event-loop responsiveness, and client materialization as separate
+stages, including abort and consumer failure cleanup.
+A model/VAE callback that publishes ordered chunks is only a producer contract;
+it is not transport, backpressure, cancellation, public streaming, or
+time-to-first-frame evidence. Distinguish complete-response faster-than-playback
+(`client E2E / output duration <= 1`) from streaming and report first-chunk,
+steady cadence, finalization, and complete-artifact latency separately.
+Run below-, near-, and above-saturation mixed-RPS soaks and report success/error
+rate, throughput, queue time, memory slope, temp growth, and worker health.
+Report p50/p95/p99 only from a declared arrival model with enough measured
+requests for the claimed percentile. See
+[serving and CI validation](references/production-validation.md).
+
+### Gate 8: Publish recipes and four independent CI tracks
+
+For every target vendor/card/topology, publish exact environment, serve/deploy
+command, one complete JSON or multipart curl per supported task, expected
+output checks, benchmark command, warmup/repeat/concurrency policy, raw metrics,
+per-rank HBM/host PSS, accuracy artifact, and limitations. Mark each
+CUDA/ROCm/NPU/XPU row `validated`, `unsupported`, or `not tested`; never inherit
+support from the platform abstraction or another card. Re-run DLO on each
+small-HBM recipe.
+
+Maintain separate CI gates:
+
+1. Function: strict load, API/offline positive and negative contracts.
+2. Accuracy: fixed reference revision/assets/seed/schedule and scoped tolerance.
+3. Performance: fixed best deployment plus memory-constrained/DLO row, raw JSON,
+   explicit warmup, enough samples for each claimed statistic, and owned
+   regression thresholds.
+4. Reliability: long mixed-RPS soak, abort/disconnect/fault cleanup, memory and
+   temporary-resource trend.
+
+## Production Definition of Done
+
+A model is production-ready only when:
+
+- every official advertised task passes API parity, offline/online parity,
+  strict loading, fixed-reference quality, and negative validation;
+- the recommended deployment row passes correctness, isolation, abort/error
+  cleanup, benchmark, and soak gates on named hardware;
+- online FP8 and DLO are each implemented and either validated in scoped rows
+  or explicitly limited/unsupported with fail-fast behavior;
+- per-request cache quality tiers have measured speed/quality frontiers and
+  request isolation, if advertised;
+- CUDA/ROCm/NPU/XPU recipes describe evidence rather than inferred support;
+- Function, Accuracy, Performance, and Reliability CI have owners, artifacts,
+  and actionable failure output;
+- accelerated paths identify every precision/quality trade-off, and the raw
+  offline plus online encoded output contracts are versioned and tested;
+- all unverified task/backend/cache/quant/offload/topology/hardware combinations
+  remain `not tested` and are not presented as supported.
+
+Day-0 support is still a valid milestone, but title, docs, and matrix must say
+which representative task passed and which production gates remain.
@@ -0,0 +1,4 @@
+interface:
+  display_name: "Production Add Diffusion Model"
+  short_description: "Productionize vLLM-Omni diffusion models"
+  default_prompt: "Use $production-add-diffusion-model to take this integrated diffusion model from Day-0 support to evidence-backed production readiness."
@@ -0,0 +1,315 @@
+# API Contract and Hardware Recipes
+
+## Contents
+
+1. [Freeze the upstream contract](#freeze-the-upstream-contract)
+2. [Validate before side effects](#validate-before-side-effects)
+3. [Task request templates](#task-request-templates)
+4. [Serving benchmark template](#serving-benchmark-template)
+5. [Hardware recipe contract](#hardware-recipe-contract)
+6. [Best deployment selection](#best-deployment-selection)
+
+## Freeze the upstream contract
+
+Pin three revisions independently:
+
+- official source implementation commit or package version;
+- checkpoint/model-card revision and partition;
+- vLLM-Omni target commit.
+
+Build a table before editing code. One row is one official task/partition.
+
+| Field | Required content |
+|---|---|
+| Task and partition | Official task name, checkpoint subdirectory, task selection rules |
+| Call contract | Constructor and call arguments, types, required/optional, official defaults |
+| Prompt/input | Prompt count/length, modality, MIME, order, count, decoded shape |
+| Size/time | Width/height alignment, pixels, frames, duration, FPS, start offset |
+| Denoise | Scheduler, sigma schedule, steps, CFG/guidance, seed behavior |
+| Output | Count, format, codec/sample rate, return type |
+| API | Offline mapping, sync endpoint, async endpoint/job lifecycle |
+| Rejection | Status code and stable error class/message for every invalid boundary |
+
+Official API parity includes limits and defaults, not just names. Do not add a
+generic default that changes a distilled checkpoint's fixed schedule, or accept
+an input combination the official task rejects.
+
+For every accepted boundary, add the neighboring rejected values. Examples:
+`max_images` and `max_images + 1`, smallest/largest legal duration, aligned and
+misaligned dimensions, omitted and illegal steps, one and too many outputs.
+
+## Validate before side effects
+
+Use one normalization/validation path for offline, sync, and async requests.
+The ordering must be:
+
+```text
+parse -> normalize aliases -> validate task/source matrix -> validate declared
+limits -> stream with byte caps -> inspect/decode with frame/pixel/duration caps
+-> persist only if needed -> submit to engine
+```
+
+Fail before download, persistence, job creation, or engine submit whenever the
+available metadata is sufficient. For streamed uploads, enforce:
+
+- total reference count and per-modality count before reading bodies;
+- per-file bytes while streaming;
+- aggregate bytes across all files;
+- decoded frames, duration, pixels, and media type before engine submission.
+
+On rejection, close streams and remove every partial temporary resource. Use
+HTTP 400 for invalid task/field combinations and 413 for payload limits. Test
+both sync and async paths; an invalid async request must not create a job ID.
+
+Do not trust client MIME alone. Sniff/validate content and keep ordered repeated
+multipart fields ordered. Avoid logging prompts, raw references, signed URLs,
+or base64 payloads.
+
+## Task request templates
+
+These are endpoint-shape templates, not model recipes. Replace placeholders
+with the exact official defaults and validated assets for the target model.
+Every production recipe must include all advertised tasks and an expected
+output check.
+
+### Text to image: JSON
+
+```bash
+curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/images/generations \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "<model>",
+    "prompt": "<validated prompt>",
+    "size": "1024x1024",
+    "num_inference_steps": 30,
+    "seed": 42,
+    "response_format": "b64_json"
+  }' | jq -er '.data[0].b64_json' | base64 -d > output.png
+file output.png
+```
+
+### Image edit: multipart
+
+```bash
+curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/images/edits \
+  -F 'model=<model>' \
+  -F 'prompt=<validated edit prompt>' \
+  -F 'image=@/path/to/reference.png;type=image/png' \
+  -F 'size=1024x1024' \
+  -F 'num_inference_steps=30' \
+  -F 'seed=42' \
+  -F 'output_format=png' \
+  | jq -er '.data[0].b64_json' | base64 -d > edited.png
+file edited.png
+```
+
+Use repeated `image` fields only if the model's official input matrix allows
+multiple ordered sources. Never describe a file upload as JSON-only.
+
+### Public request fields versus internal sampling fields
+
+Use the field exposed by the target revision's endpoint schema; do not rename
+it from the pipeline's internal attribute. For example, the current video JSON
+and multipart contracts expose `extra_params`, while serving translates that
+object into `DiffusionSamplingParams.extra_args`. Other request paths may
+expose canonical `extra_args` or nest it under `extra_body`. Inspect the exact
+Pydantic model/FastAPI form dependency and exercise the generated curl before
+publishing a recipe. A deprecation warning in a different endpoint is not
+evidence that an unrecognized multipart form field works here.
+
+Current endpoint map (re-check it at the target commit):
+
+| Endpoint | Public model-specific field |
+|---|---|
+| `/v1/videos`, `/v1/videos/sync` | Multipart `extra_params` JSON string |
+| `/v1/images/generations` | JSON `extra_params` object |
+| `/v1/chat/completions` diffusion | `extra_body.extra_args` |
+| `/v1/images/edits` | Declared typed form fields only |
+| `/v1/audio/generate` | Declared typed JSON fields only |
+
+Prefer declared top-level fields. Use a generic model-specific envelope only
+when the endpoint actually exposes it.
+
+### Text to video: multipart
+
+```bash
+curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
+  -H 'Accept: video/mp4' \
+  -F 'model=<model>' \
+  -F 'prompt=<validated prompt>' \
+  -F 'size=1280x720' \
+  -F 'num_frames=81' \
+  -F 'fps=16' \
+  -F 'num_inference_steps=30' \
+  -F 'seed=42' \
+  -F 'extra_params={"task":"<official-task>"}' \
+  -o output.mp4
+ffprobe -v error -show_entries stream=codec_name,width,height,r_frame_rate \
+  -of json output.mp4
+```
+
+### Image/video to video: multipart
+
+```bash
+curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
+  -H 'Accept: video/mp4' \
+  -F 'model=<model>' \
+  -F 'prompt=<validated prompt>' \
+  -F 'num_frames=81' \
+  -F 'fps=16' \
+  -F 'num_inference_steps=30' \
+  -F 'seed=42' \
+  -F 'extra_params={"task":"<official-task>","start_time_seconds":0}' \
+  -F 'input_reference=@/path/to/reference.png;type=image/png' \
+  -o i2v.mp4
+
+curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
+  -H 'Accept: video/mp4' \
+  -F 'model=<model>' \
+  -F 'prompt=<validated prompt>' \
+  -F 'num_frames=81' \
+  -F 'num_inference_steps=30' \
+  -F 'seed=43' \
+  -F 'extra_params={"task":"<official-task>","start_time_seconds":0}' \
+  -F 'input_reference=@/path/to/reference.mp4;type=video/mp4' \
+  -o v2v.mp4
+```
+
+For multiple references, repeat the model's accepted multipart field in
+official order. Include complete downloadable or repository fixtures plus
+SHA256 hashes. A placeholder-only recipe is not runnable evidence.
+
+### Text to audio: JSON
+
+```bash
+curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/audio/generate \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "model": "<model>",
+    "input": "<validated audio prompt>",
+    "audio_length": 10.0,
+    "num_inference_steps": 100,
+    "seed": 42,
+    "response_format": "wav"
+  }' -o output.wav
+ffprobe -v error -show_entries stream=codec_name,sample_rate,channels \
+  -of json output.wav
+```
+
+Use `/v1/audio/generate` for general diffusion audio. Do not substitute the
+speech API unless the model is actually served through the speech contract.
+
+### Async video lifecycle
+
+Alongside `/v1/videos/sync`, include `POST /v1/videos`, status polling,
+content retrieval, cancellation/abort if public, and cleanup assertions. The
+same request must normalize to the same engine sampling parameters in both
+modes.
+
+## Serving benchmark template
+
+Run a smoke first, then fixed-concurrency and arrival-rate tests. Video task
+names for the harness are `t2v`, `i2v`, `ti2v`, or `v2v`; model-specific task
+selection belongs in `--extra-body` when required.
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+  --base-url http://127.0.0.1:8091 \
+  --endpoint /v1/videos \
+  --model '<model>' \
+  --dataset random \
+  --task t2v \
+  --num-prompts 100 \
+  --width 1280 --height 720 \
+  --num-frames 81 --fps 16 \
+  --num-inference-steps 30 \
+  --request-rate 0.2 \
+  --max-concurrency 4 \
+  --warmup-requests 4 \
+  --warmup-concurrency 4 \
+  --extra-body '{"extra_params":"{\"task\":\"<official-task>\"}"}' \
+  --output-file result.json
+```
+
+The video backend currently stringifies each multipart value. Pre-serialize
+the nested `extra_params` object as above; passing it as a nested object would
+produce Python dict syntax rather than valid JSON. If the target revision
+changes the backend to JSON-encode dict/list values, update and smoke-test this
+command against that revision.
+
+For i2v/v2v, use the harness's validated dataset/trace path or provide the
+documented media source; do not silently benchmark t2v and label it i2v. Keep
+prompt, assets, seed, dimensions, frames, steps, guidance, warmup, concurrency,
+and arrival process identical across A/B runs.
+
+Keep three benchmark lanes distinct:
+
+1. lossless runtime A/B with the released schedule/precision/attention policy;
+2. accelerated-path A/B that changes one adapter, quantization, cache, or
+   sparsity policy and adds same-seed quality evidence;
+3. production-topology studies that vary placement or parallelism and report a
+   latency/throughput/memory frontier rather than a kernel speedup.
+
+For a fixed single-request A/B, retain every measured run and report median or
+mean with range after an explicit warmup; do not infer p95/p99 from a few
+repetitions. Tail latency requires a declared arrival process, request count,
+concurrency, and enough successful samples. Reconcile client E2E with queue,
+encoder, denoise, video/audio decode, D2H/IPC, codec, and residual timings.
+
+For generated media, define `real-time` rather than relying on the label. A
+complete-response claim uses `client E2E / validated output duration <= 1` and
+does not imply streaming or low time-to-first-frame. Report first chunk/fragment
+and cadence separately when a public streaming route exists. Never add gains
+across benchmark lanes with different schedules, prompts, seeds, artifacts, or
+topologies.
+
+The diffusion serving harness currently covers image/video task families, not
+`/v1/audio/generate`. For text-to-audio, use the target model's task-specific
+audio benchmark if one exists; otherwise add a serving harness that records
+request rate, concurrency, latency, generated-audio seconds per second, error
+rate, and output validation. Do not relabel the TTS harness or a video workload
+as diffusion-audio evidence.
+
+## Hardware recipe contract
+
+Create a separate row and runnable section for each vendor/card/topology:
+
+| Evidence | Required detail |
+|---|---|
+| Identity | Vendor, exact SKU, device count, HBM, interconnect |
+| Software | Driver, CUDA/ROCm/CANN/oneAPI, framework/container, dependency versions |
+| Source | vLLM-Omni commit, official/checkpoint revision and hashes |
+| Deployment | dtype, attention backend per role, TP/SP/CFG/HSDP/DP, cache, quant, offload |
+| Commands | Complete deploy YAML/serve command, every task curl, benchmark command |
+| Method | Warmup exclusion, repetitions or request count, concurrency, arrival distribution, duration, statistic definitions |
+| Results | Success rate, p50/p95/p99, RPS/throughput per device, stage latency |
+| Memory | Load/materialize, resident, encode, denoise, decode, transient per-rank HBM and host PSS |
+| Quality | Metric/tolerance, artifact links/hashes, temporal/audio checks |
+| Output | Raw/offline dtype, range, layout and ownership; online codec, payload bytes, transport route and client boundary |
+| Limits | `validated`/`limited`/`unsupported`/`not tested`, fallback/rejection behavior |
+
+Use current paths such as `recipes/<vendor>/`, `recipes/README.md`, supported
+models, and feature compatibility tables. Top-level documentation must link to
+the scoped matrix; it must not use an unqualified checkmark that hides a ROCm,
+NPU, XPU, task, topology, or dtype limitation.
+
+DLO evidence is per card, especially on small-HBM devices. A DLO result on one
+CUDA card does not establish another CUDA SKU, ROCm, NPU, or XPU support.
+
+## Best deployment selection
+
+Benchmark a correctness-approved candidate matrix, not every theoretical
+combination. Choose the recommended row by the production objective—latency,
+throughput, cost, or small-HBM fit—and document that objective.
+
+Do not recommend a row until it passes:
+
+- all advertised task/API contracts;
+- fixed-reference quality;
+- request isolation and abort/error cleanup;
+- sustained arrival-rate tests;
+- per-rank HBM and host PSS limits;
+- hardware-specific accuracy and performance CI.
+
+Keep the most memory-constrained DLO row even when it is not the fastest; it is
+a separate production recipe, not a failed latency candidate.
@@ -0,0 +1,462 @@
+# Production Feature Patterns
+
+## Contents
+
+1. [Compatibility matrix](#compatibility-matrix)
+2. [Online FP8](#online-fp8)
+3. [Strict fused-weight loading](#strict-fused-weight-loading)
+4. [Distributed layerwise offload](#distributed-layerwise-offload)
+5. [Per-request Cache-DiT](#per-request-cache-dit)
+6. [Combination gates](#combination-gates)
+
+## Compatibility matrix
+
+Create rows at the granularity below. Add evidence links and a reason for every
+non-validated state.
+
+| Task/shape | Execution mode/capacity | Backend/layout | Quant | Cache policy | Offload | Topology | Hardware/dtype | Output contract/transport | State | Evidence |
+|---|---|---|---|---|---|---|---|---|---|---|
+| `<task>/<shape>/<schedule>` | serial request / 1 | SDPA, padded | BF16 | none | resident | 1 device | `<card>` BF16 | raw float / local | `not tested` | — |
+| `<task>/<shape>/<schedule>` | step / N | `<fast>`, packed | BF16 | none | resident | `<TP/SP>` | `<card>` BF16 | encoded / subprocess | `not tested` | — |
+| `<task>/<shape>/<schedule>` | step / N | `<fast>`, packed | FP8 | high | DLO no-AG | `<TP/SP>` | `<card>` | device uint8 / SHM | `not tested` | — |
+
+Initialize every row as `not tested`; the table is a work queue, not evidence.
+Promote the dense BF16 row to the oracle only after its scoped parity passes.
+Add one axis at a time. If the runtime cannot
+enforce a known-incompatible combination, add construction/admission
+validation before writing a recipe.
+
+## Online FP8
+
+### Component routing
+
+The public builder accepts one method or a longest-prefix component map:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+config = build_quant_config({
+    "transformer": {
+        "method": "fp8",
+        "ignored_layers": ["transformer.final_layer"],
+    },
+    "text_encoder": None,
+    "vae": None,
+    "default": None,
+})
+
+assert config.resolve("transformer.blocks.0.attn.qkv_proj") is not None
+assert config.resolve("text_encoder.layers.0.mlp") is None
+assert config.resolve("vae.decoder.conv_out") is None
+```
+
+Use actual runtime prefixes. Multi-DiT or nested pipelines may not use the
+generic names `transformer` and `vae`. A prefix that falls through to a default
+is not component validation. Add tests for every included and ignored prefix.
+
+For CLI serving, first validate resident DiT-only FP8:
+
+```bash
+vllm serve '<model>' --omni --quantization fp8 --max-num-seqs 1
+```
+
+For scoped routing:
+
+```bash
+vllm serve '<model>' --omni \
+  --diffusion-quantization-config \
+  '{"<dit-runtime-prefix>":{"method":"fp8","ignored_layers":["<exact-linear-prefix>"]},"<vae-runtime-prefix>":null,"default":null}' \
+  --max-num-seqs 1
+```
+
+Do not publish these placeholders. A recipe must contain the target model's
+verified prefixes.
+
+### Quantizable module wiring
+
+Pass `quant_config` and a stable prefix through every vLLM linear. Keep
+precision-sensitive modulation/norm layers safe unless proven otherwise.
+
+```python
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+
+self.qkv_proj = QKVParallelLinear(
+    hidden_size=hidden_size,
+    head_size=head_dim,
+    total_num_heads=num_heads,
+    total_num_kv_heads=num_kv_heads,
+    bias=False,
+    quant_config=quant_config,
+    prefix=f"{prefix}.attn.qkv_proj",
+)
+self.gate_up_proj = MergedColumnParallelLinear(
+    hidden_size,
+    [intermediate_size, intermediate_size],
+    bias=False,
+    quant_config=quant_config,
+    prefix=f"{prefix}.mlp.gate_up_proj",
+)
+self.act_fn = SiluAndMul()
+self.down_proj = RowParallelLinear(
+    intermediate_size,
+    hidden_size,
+    bias=False,
+    input_is_parallel=True,
+    quant_config=quant_config,
+    prefix=f"{prefix}.mlp.down_proj",
+)
+```
+
+Set required checkpoint parameters to fail rather than silently initialize if
+the current loader supports `missing_param_init = "error"`.
+
+### FP8 evidence
+
+For each task/hardware row, capture:
+
+- startup proves the FP8 quant method is attached to intended linear modules;
+- no unexpected BF16 fallback and no meta/uninitialized parameters;
+- same-seed component, trajectory, and final-artifact comparison with BF16;
+- HBM at load, resident, encode, denoise, decode, and peak;
+- cold load time; fixed-work warm-latency raw runs; and serving p50/p95/p99 plus
+  throughput only from a declared arrival load with enough samples;
+- named ignored layers and why their BF16 retention is required.
+
+DiT FP8 evidence does not cover the text encoder or VAE. CUDA evidence does not
+cover ROCm/NPU/XPU. A pre-quantized checkpoint is a different loading path from
+online FP8 and needs its own state. If it is also pruned, rotated, distilled, or
+adapter-modified, report that model delta explicitly; BF16-versus-artifact
+quality and speed are deployment comparisons, not pure quantization ablations.
+
+## Strict fused-weight loading
+
+Never use `param.data.copy_()` for quantizable vLLM parameters. Transform the
+checkpoint layout before invoking the parameter loader so the online quant
+loader remains outermost:
+
+```python
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+param = params[target_name]
+weight_loader = getattr(param, "weight_loader", default_weight_loader)
+
+if target_name.endswith(".qkv_proj.weight"):
+    packed_qkv = reorder_checkpoint_qkv(checkpoint_weight)
+    weight_loader(param, packed_qkv)
+elif target_name.endswith(".gate_up_proj.weight"):
+    gate, up = checkpoint_weight.chunk(2, dim=0)
+    weight_loader(param, gate, 0)
+    weight_loader(param, up, 1)
+else:
+    weight_loader(param, checkpoint_weight)
+```
+
+Track source shards per fused destination:
+
+```python
+expected = {
+    "blocks.0.attn.qkv_proj.weight": {"q", "k", "v"},
+    "blocks.0.mlp.gate_up_proj.weight": {0, 1},
+}
+seen = {name: set() for name in expected}
+
+# After mapping/loading each source:
+seen[target_name].add(shard_id)
+
+incomplete = {
+    name: sorted(shards - seen[name], key=str)
+    for name, shards in expected.items()
+    if shards - seen[name]
+}
+if incomplete:
+    raise RuntimeError(f"incomplete fused checkpoint parameters: {incomplete}")
+```
+
+Also compare retained model parameters with the returned loaded set. Explicit
+checkpoint ignores need a reason and test; an unknown or missing weight is not
+an ignore policy.
+
+Load-time adapter fusion creates a dedicated model, not a request-switchable
+LoRA. Pin the base and adapter revisions, validate every low-rank and full-rank
+delta target, freeze the adapter-declared schedule/task, and fuse before
+sharding only when the loader contract proves identical rank-local weights.
+Reject offload or host-weight paths that bypass that fusion rather than serving
+the unfused base model silently.
+
+## Distributed layerwise offload
+
+### Declare component topology
+
+Use `SupportsComponentDiscovery` and `OffloadPlan`. Validate all dotted paths
+at construction/test time because discovery may otherwise warn and skip.
+
+Treat a future/requested `--layerwise-offload-components` selector as a public
+API only if it exists in the target revision. Today, express component topology
+with `OffloadPlan` (`on_demand_component_paths`, `block_attrs`,
+`encoder_block_attrs`, `resident_dit_paths`, and `offload_submodules`). Do not
+invent a CLI flag or claim component-selective DLO from discovery declarations
+alone; validate the selector separately when it lands.
+
+```python
+from operator import attrgetter
+
+import torch
+
+def assert_offload_path(root, path):
+    value = attrgetter(path)(root)
+    if not isinstance(value, torch.nn.Module):
+        raise TypeError(f"offload path {path!r} is not an nn.Module")
+    return value
+
+for path in (
+    *pipeline._dit_modules,
+    *pipeline._encoder_modules,
+    *pipeline._vae_modules,
+    *pipeline._offload_plan.on_demand_component_paths,
+):
+    assert_offload_path(pipeline, path)
+```
+
+For `block_attrs`, resolve the DiT first and then each declared block attribute.
+Require an indexable block container with at least one `nn.Module`. Validate
+`offload_submodules`, `resident_dit_paths`, and `encoder_block_attrs` similarly.
+
+### Keep host storage loader-owned
+
+Current DLO accepts a loader-produced `HostWeightPlan`. Pipeline authors should
+not instantiate this plan or add ad hoc mmap capability flags. The loader must
+preflight complete runtime/checkpoint bindings, names, shapes, dtypes, persistent
+buffers, and represented transforms before it skips ordinary materialization.
+The exact accepted plan is consumed once by DLO; a rejected plan records an
+observable fallback reason and continues through the ordinary loader.
+
+Direct checkpoint mmap currently requires TP1 without HSDP or online
+quantization. TP>1 falls back to the ordinary TP-aware loader and may feed DLO
+AllGather or no-AllGather; the current design records a bounded DP2xTP2
+AllGather smoke, but that path does not receive direct-mmap host-page sharing.
+Per-tensor online FP8 must use the ordinary loader. Current revisions can admit
+its finalized weight/scale layout to DLO AllGather, including transposed runtime
+weights, while unsupported online quantizers still fail closed. Direct mmap
+remains unavailable, so capture the temporary full-model host materialization
+before rank sharding.
+
+Host Weight Runtime is a distinct opt-in path. Current final-layout consumers
+use no-AllGather DLO and exact model-declared BF16 representations. In
+`preferred` mode, consume an exact local hit or canonically load and publish for
+a future startup; in `required` mode, consume only and fail on a miss or unusable
+artifact. HWR must not interact with DLO AllGather, online quantization, HSDP,
+LoRA/adapted weights, or an unrecognized load format unless the target revision
+adds an exact producer/restore contract. It is immutable node-local CPU backing,
+not zero-copy GPU execution; the DLO transport still owns staging or direct H2D.
+Qualify identity inputs, cold population, warm reuse, atomic concurrent
+publication, source digests, corruption/quarantine, lease cleanup, capacity,
+node-local PSS, H2D, and E2E latency. A host-memory saving that materially
+regresses transfer or request latency is a separate memory-first deployment.
+
+### DLO deployment modes
+
+Treat these paths separately:
+
+| Path | Weight/runtime behavior | Required boundaries |
+|---|---|---|
+| DLO AllGather | Rank-sharded pinned host tensors; reconstruct layer through collective | TP>1 uses ordinary TP-aware loader output; HSDP and unsupported online quantizers are rejected; finalized per-tensor online FP8 is eligible only when recognized by the target revision; concurrent ranks must participate consistently |
+| DLO no-AllGather | Loader-approved checkpoint mmap or ordinary runtime tensors; each rank streams a complete block | Direct mmap is TP1/non-HSDP/non-online-quant only; TP/HSDP/online quant use ordinary runtime tensors and require scoped E2E |
+| SP + DLO | SP group can supply the DLO collective group | Packed boundaries and collective order require E2E |
+
+Single-stage examples:
+
+```bash
+# SP + DLO AllGather candidate; validate exact model/card before publishing.
+vllm serve '<model>' --omni \
+  --enable-distributed-layerwise-offload \
+  --usp 4 \
+  --max-num-seqs 1
+
+# TP + DLO rank-local candidate; not a support statement.
+vllm serve '<model>' --omni \
+  --tensor-parallel-size 2 \
+  --enable-distributed-layerwise-offload \
+  --dlo-no-use-allgather \
+  --max-num-seqs 1
+```
+
+Under `--omni`, configure DP in a deploy YAML rather than passing vLLM DP CLI
+flags:
+
+```yaml
+pipeline: <registered-pipeline>
+async_chunk: false
+data_parallel_size: 4
+stages:
+  - stage_id: 0
+    devices: "0,1,2,3"
+    max_num_seqs: 1
+    enable_distributed_layerwise_offload: true
+    dlo_use_allgather: true
+    parallel_config:
+      tensor_parallel_size: 1
+      sequence_parallel_size: 1
+```
+
+```bash
+vllm serve '<model>' --omni --deploy-config /path/to/dlo_dp4.yaml
+```
+
+Keep the YAML next to the validated recipe. Match its schema to the target
+revision; do not assume a model accepts another pipeline's deploy config.
+
+### DLO validation sequence
+
+1. Resident BF16 reference.
+2. Ordinary layerwise offload parity.
+3. DLO enable, warmup, same-seed generation, disable/restart.
+4. AllGather and no-AllGather separately.
+5. Concurrent requests with different prompts and identical collective-affecting
+   parameters; then reject or schedule incompatible steps/shapes safely.
+6. Queued/in-flight abort, error, idle rank, worker exit, and next request.
+7. TP, SP, cache, compile, online FP8 one axis at a time.
+8. Loader-plan selection/fallback reason, cold/warm storage lifecycle, and
+   corruption/stale-entry recovery when mmap/cache backing is advertised.
+9. Replica-local request preparation plus DLO-group-only weight collectives in
+   DP; exercise concurrent replicas to catch accidental WORLD rendezvous.
+10. Sweep resident block counts; publish raw latency/HBM pairs, dominated
+    points, and the non-dominated Pareto frontier.
+11. Per-rank HBM, host PSS for the full process tree, H2D and collective trace.
+
+On a small-HBM card, measure transient encode/VAE peaks as well as resident DiT
+weights. If a non-DiT component still exceeds HBM, combine only independently
+validated on-demand component staging, VAE tiling/patching, or another topology;
+do not hide the remaining peak behind the phrase “DLO enabled.”
+
+## Per-request Cache-DiT
+
+### Policy and lifecycle skeleton
+
+First verify that the target revision contains the shared protocol/runtime.
+If not, follow the nearest current model's lifecycle without inventing a new
+public API.
+
+The following is a **serial/exclusive request-mode sketch only**. Pipeline-wide
+hooks are mutable; concurrent requests need an admission/batching key and
+runner-owned transition serialization before this shape is safe.
+
+```python
+import torch.nn as nn
+
+from vllm_omni.diffusion.cache.cachedit import (
+    CacheDiTBackend,
+    CacheDiTRequestSpec,
+    RequestScopedCacheDiTRuntime,
+)
+
+class MyPipeline(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        # ... components ...
+        self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
+
+    def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
+        self._cache_dit_runtime.adopt(
+            backend,
+            installation_key="my_model.generic",
+        )
+
+    def is_cache_dit_enabled(self) -> bool:
+        return self._cache_dit_runtime.is_enabled
+
+    def _cache_spec(self, quality, steps):
+        if quality == "lossless":
+            return None
+        if quality == "high":
+            return CacheDiTRequestSpec(
+                installation_key="my_model.high.v1",
+                cache_config=self._validated_high_config,
+                num_inference_steps=steps,
+            )
+        raise ValueError(f"unsupported quality tier: {quality!r}")
+
+    def forward(self, req):
+        quality, steps = validate_and_resolve_request(req)
+        self._cache_dit_runtime.prepare(self._cache_spec(quality, steps))
+        return self._run_validated_request(req)
+```
+
+The snippet shows the model-owned transition for exclusive execution. Put
+disable/restore handling in the runtime/runner lifecycle that owns success,
+exception, disconnect, and asynchronous abort; a `try/finally` inside
+`forward()` alone cannot observe every server cancellation boundary.
+
+The `installation_key` must include every policy dimension that changes hook
+installation; changes in step count can use runtime refresh. `None` disables
+installed hooks. Do not mutate the global startup cache config for a request.
+
+Centralize cleanup in the owning request/runner boundary. An exception, abort,
+disconnect, or admission failure must not leave a changed hook policy visible
+to the next request. The exact owner depends on the target revision; add a
+failure-injection test rather than relying on `forward()` alone to observe
+asynchronous abort.
+
+Example model-specific request form after calibration:
+
+```bash
+curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
+  -F 'model=<model>' \
+  -F 'prompt=<prompt>' \
+  -F 'quality=lossless' \
+  -F 'num_inference_steps=<validated-steps>' \
+  -o lossless.mp4
+
+curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
+  -F 'model=<model>' \
+  -F 'prompt=<prompt>' \
+  -F 'quality=high' \
+  -F 'num_inference_steps=<validated-steps>' \
+  -o high.mp4
+```
+
+Do not expose these tier names merely because another model uses them.
+
+### Cache-DiT evidence
+
+For every tier/task/shape/schedule:
+
+- show real cache hits/skipped blocks, not only “backend enabled” logs;
+- measure fixed-work speed and quality against lossless/native;
+- publish a speed/quality frontier and chosen tier mapping;
+- alternate lossless/high/lossless and vary steps to prove restore/refresh;
+- interleave tasks and shapes; run concurrent requests only when scheduling
+  guarantees exclusive compatible hook state;
+- inject validation error, denoise error, disconnect, and abort, then prove the
+  next request starts with its requested policy;
+- bound any cached tensors and clear request-local state.
+
+## Combination gates
+
+Start every combination as `not tested`.
+
+| Combination | Default production posture |
+|---|---|
+| Direct checkpoint mmap + TP>1/HSDP/online quant | Preflight falls back to the ordinary loader; do not claim direct-mmap savings |
+| Per-tensor online FP8 + DLO AllGather | Generic compatibility exists for finalized weights/scales; each model/card/topology remains a candidate until E2E, and other online quantizers fail closed |
+| Online FP8 + DLO no-AllGather | Candidate only after model/card/topology E2E |
+| TP>1 + DLO AllGather | Ordinary TP-aware loader only; bounded DP2xTP2 smoke exists, but each new model/card/topology still needs E2E |
+| TP>1 + DLO no-AllGather | Ordinary TP-aware loader only; candidate with limited generic coverage |
+| HSDP + DLO AllGather | Reject; it would double-shard HSDP-managed parameters |
+| HSDP + DLO no-AllGather | Configuration-compatible candidate with limited E2E coverage |
+| Cache-DiT + online FP8 | Separate trajectory, hit-rate, HBM, latency test |
+| Cache-DiT + DLO | Verify hooks, streamed blocks, memory peaks, abort cleanup |
+| Cache-DiT + step execution | Do not advertise until lifecycle/batching support is explicit |
+| Sparse attention + Ring | Verify backend execution; reject silently ignored sparse settings |
+| Load-time fused adapter + offload/HWR | Reject unless the selected loader path applies and validates the same fusion before sharding/storage |
+| Packed attention + any topology | Unequal-sample boundary parity is mandatory |
+| Device-side uint8 preparation + offline caller | Version the dtype/range/layout change; HTTP MP4 parity does not preserve the raw tensor contract |
+| Device-side preparation + remote codec | Validate route selection, strided/contiguous layouts, payload bytes, ordering, fallback, and encoded-media parity |
+
+A startup-only test never upgrades a combination. Require request output,
+parity/quality, actual fast-path evidence, resource cleanup, and a benchmark on
+the named hardware.
@@ -0,0 +1,481 @@
+# Performance Patterns and Code Examples
+
+## Contents
+
+1. [Measurement loop](#measurement-loop)
+2. [Faster BF16 attention](#faster-bf16-attention)
+3. [Sparse attention](#sparse-attention)
+4. [Fusion-first operator patterns](#fusion-first-operator-patterns)
+5. [Remove redundant computation](#remove-redundant-computation)
+6. [USP communication optimization](#usp-communication-optimization)
+7. [Text encoder and VAE disaggregation](#text-encoder-and-vae-disaggregation)
+8. [Performance acceptance](#performance-acceptance)
+
+## Measurement loop
+
+Do not start with a list of optimizations. Start with a fixed correctness and
+performance workload:
+
+```text
+model/checkpoint revision + prompt/media hashes + seed + dimensions/frames +
+scheduler/sigma/steps/guidance + task + dtype + backend + topology + hardware
+```
+
+Capture:
+
+- stage time: load, encode, conditioning, denoise, VAE decode/postprocess;
+- GPU trace: kernel duration/count, launch gaps, allocations, H2D, collectives;
+- CPU trace: preprocessing, per-layer sync, Python dispatch, temporary files;
+- memory: per-rank HBM peaks and process-tree host PSS;
+- output: device preparation, D2H/IPC, payload bytes, codec wall/process CPU,
+  event-loop wait, client materialization, and signed residual;
+- serving: queue time, throughput, errors, request mix, and distribution samples.
+
+For a fixed-work single-request A/B, run one explicit warmup policy and at
+least three measured repetitions, keep every raw run, and report median or mean
+with range. Do not derive p95/p99 from three samples. For tail latency, declare
+the arrival model and collect enough successful requests for the percentile and
+confidence claimed. Implement one optimization at a time, re-run parity, and
+report both local kernel and end-to-end deltas. A kernel win that moves no E2E
+metric is not a production speed claim.
+
+Use the pipeline profiler already exposed by the target revision when
+available, and use PyTorch/CUDA/ROCm/NPU/XPU profiling tools appropriate to the
+platform. Do not force every platform through a CUDA-only workflow.
+
+## Faster BF16 attention
+
+### Shared attention wiring
+
+Use the shared role-aware attention layer and give it the true local head
+counts and QKV layout:
+
+```python
+import torch
+
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm_omni.diffusion.attention.layer import Attention
+
+class SelfAttention(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_size,
+        head_dim,
+        num_heads,
+        num_kv_heads,
+        quant_config=None,
+        prefix="",
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_kv_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.attention = Attention(
+            num_heads=self.qkv_proj.num_heads,
+            num_kv_heads=self.qkv_proj.num_kv_heads,
+            head_size=head_dim,
+            softmax_scale=head_dim**-0.5,
+            causal=False,
+            qkv_layout="BSND",
+            role="self",
+            prefix=prefix,
+        )
+        self.out_proj = RowParallelLinear(
+            num_heads * head_dim,
+            hidden_size,
+            bias=False,
+            input_is_parallel=True,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+        )
+
+    def forward(self, x, attn_metadata=None):
+        qkv, _ = self.qkv_proj(x)
+        q_size = self.qkv_proj.num_heads * self.head_dim
+        kv_size = self.qkv_proj.num_kv_heads * self.head_dim
+        q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
+        q = q.unflatten(-1, (self.qkv_proj.num_heads, self.head_dim))
+        k = k.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim))
+        v = v.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim))
+        out = self.attention(q, k, v, attn_metadata)
+        out = out.flatten(-2)
+        out, _ = self.out_proj(out)
+        return out
+```
+
+Add Q/K norm and RoPE in the official order before the `Attention` call. Keep
+checkpoint mapping for the fused QKV and output projection complete.
+
+### Backend qualification
+
+For each role (`self`, cross-attention category, token refiner, etc.), record:
+
+- actual selected backend from logs/trace;
+- supported dtype, head size, QKV layout, mask/metadata, packed varlen behavior;
+- platform and compute capability;
+- TP/SP/Ring/AllGather-KV restrictions;
+- fallback behavior and whether fallback changes the claimed performance.
+
+Use unequal packed samples and compare with dense SDPA. Verify `cu_seqlens`,
+max sequence length, padding exclusion, modality boundaries, CFG ownership, and
+output slicing. Do not pass an `attn_mask` and assume it is consumed. Inspect
+the backend contract and trace.
+
+TRTLLM, Flash Attention, cuDNN, NPU varlen, and other backends are independent
+qualification rows. The shared `Attention` class does not make every backend
+packed-safe or faster automatically. Ring and hybrid paths require their own
+boundary parity. AllGather-KV and specific backend combinations may fail fast;
+keep that behavior visible in docs.
+
+## Sparse attention
+
+Sparse attention is a backend/model/platform optimization, not a synonym for
+low-precision dense attention. Enable it only after dense BF16 parity.
+
+For every candidate task/shape/topology:
+
+1. State the sparsity policy, selected backend/kernel, block/head dimensions,
+   supported hardware, external kernel/package revision, and fallback.
+2. Prove the sparse kernel executed and report realized sparsity; an enabled
+   flag is insufficient.
+3. Compare same-seed per-block outputs, denoise trajectory, and final artifact
+   with dense attention.
+4. Add modality-specific metrics: temporal consistency and audio/video sync for
+   video, spectral/audio quality for audio, and prompt/image metrics as relevant.
+5. Report attention time and full E2E time. Include preprocessing/index-build
+   overhead, peak memory, and the exact dense control.
+6. Validate each task and shape. Do not extrapolate a T2V result to I2V/Ref2V,
+   or NPU sparse support to CUDA/ROCm/XPU.
+
+If Ring bypasses the selected backend or a sparse setting would be silently
+ignored, reject the combination. Prefer a visible dense fallback only when the
+recipe labels the row `limited` and does not claim sparse speed.
+
+## Fusion-first operator patterns
+
+Prefer current shared/custom operators with explicit guards and native
+fallbacks. Prove parity at operator, block, trajectory, and artifact levels,
+then use a trace to prove fewer launches/bytes or lower latency.
+
+### Q/K RMSNorm + RoPE
+
+```python
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
+
+self.norm_q = RMSNorm(head_dim, eps=eps)
+self.norm_k = RMSNorm(head_dim, eps=eps)
+# Example only: set these booleans from the official checkpoint contract.
+use_neox_layout = False
+cos_sin_are_half_dim = True
+self.rope = RotaryEmbedding(
+    is_neox_style=use_neox_layout,
+    half_head_dim=cos_sin_are_half_dim,
+)
+
+q = self.norm_q(q)
+k = self.norm_k(k)
+q, k = apply_rope_to_qk(self.rope, q, k, (cos, sin))
+```
+
+Replace the illustrative booleans with the official layout in real code.
+Preserve partial rotary dimensions by splitting the rotated prefix and concatenating
+the pass-through suffix. If 3D cos/sin includes a batch axis, the current
+shared path assumes the values are shared across the batch; per-sample
+different positions require a proven representation/path.
+
+For packed `[tokens, heads, head_dim]` tensors with a non-interleaved
+`[cos | sin]` table, the current shared fused boundary is:
+
+```python
+from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
+
+q, k = fused_qk_norm_rope(
+    q,
+    k,
+    self.norm_q.weight,
+    self.norm_k.weight,
+    rope_table,
+    self.norm_q.variance_epsilon,
+)
+```
+
+Its CUDA fast path currently specializes BF16 `head_dim=128` and
+`rotary_dim=96`; other valid inputs take the eager implementation. The separate
+shared RMSNorm and RotaryEmbedding still provide the general cross-platform
+baseline. Treat NPU/MUSA and other vendor fusions as independent qualification
+rows, and do not generalize a model-specific Triton kernel without compatible
+shape, layout, frequency partition, dtype, and platform evidence.
+
+### Fused QKV projection
+
+Use `QKVParallelLinear` as shown above. Qualification includes:
+
+- `num_heads` and `num_kv_heads` divisibility by TP;
+- local split sizes from the layer, never reusing global counts after TP;
+- Q/K/V checkpoint order conversion before the vLLM loader;
+- all Q, K, and V shards present;
+- quantization scale/loader mapping when FP8 or another method is enabled.
+
+### Gate-up projection + fused SwiGLU
+
+```python
+from vllm.model_executor.layers.activation import SiluAndMul
+from vllm.model_executor.layers.linear import (
+    MergedColumnParallelLinear,
+    RowParallelLinear,
+)
+
+self.gate_up_proj = MergedColumnParallelLinear(
+    input_size=hidden_size,
+    output_sizes=[intermediate_size, intermediate_size],
+    bias=False,
+    quant_config=quant_config,
+    prefix=f"{prefix}.gate_up_proj",
+)
+self.act_fn = SiluAndMul()
+self.down_proj = RowParallelLinear(
+    input_size=intermediate_size,
+    output_size=hidden_size,
+    bias=False,
+    input_is_parallel=True,
+    quant_config=quant_config,
+    prefix=f"{prefix}.down_proj",
+)
+
+gate_up, _ = self.gate_up_proj(x)
+x = self.act_fn(gate_up)
+x, _ = self.down_proj(x)
+```
+
+Map both checkpoint `gate_proj` and `up_proj` shards into the merged parameter
+and require both. A merged projection followed by `F.silu(gate) * up` is not
+the same claim as the `SiluAndMul` fused activation kernel; identify what the
+trace actually uses.
+
+### AdaLN
+
+Use a shared layer only when its mathematical/return contract matches the
+official model:
+
+```python
+from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
+
+self.adaln = AdaLayerNorm(
+    hidden_size,
+    elementwise_affine=False,
+    eps=1e-6,
+)
+x = self.adaln(x, scale[:, None, :], shift[:, None, :])
+```
+
+`AdaLayerNormZero`, `AdaLayerNormZeroSingle`, and continuous variants include
+different projections, broadcasts, and tuple outputs; they are not drop-in
+replacements. A shared AdaLN may have an NPU fast path while using native ops on
+CUDA/HIP/XPU. Do not call it a universal fused AdaLN kernel.
+
+For sigma schedules with a small finite set of timesteps, profile whether
+precomputing the timestep/AdaLN projection once per request is worthwhile.
+Cache keys must include the exact schedule/config/device/dtype and stay bounded.
+
+### Layout, residual, and other candidates
+
+After profiling, inspect:
+
+- bias/dropout/add or gated residual chains;
+- norm + linear or norm + modulation boundaries;
+- repeated transpose/reshape/contiguous around attention;
+- VAE normalization/activation/convolution boundaries;
+- MoE routing and fused experts for models that actually use MoE.
+
+For fused residual-plus-norm or modulation chains, preserve the unfused
+materialization boundary. If the residual was stored in BF16 before RMSNorm,
+round to BF16 and promote again inside the fused kernel; using the unrounded
+FP32 temporary changes the denoising trajectory even when the stored residual
+looks correct.
+
+Keep DLO materialize/release hooks, process-group collectives, request cache
+transitions, and cleanup outside compiled/fused regions unless lifecycle safety
+is explicitly tested. Compile only stable regions and retain eager fallback.
+
+## Remove redundant computation
+
+### Hoist request invariants
+
+Audit the denoise loop for work that depends only on the request rather than
+the current timestep:
+
+```python
+# Before the denoise loop.
+condition = self.condition_proj(text_embeddings)
+refined_condition = self.token_refiner(condition, packed_metadata)
+rope_tables = self.build_rope_tables(position_ids)
+attn_metadata = self.build_attention_metadata(layout)
+
+for step, sigma in enumerate(sigmas):
+    timestep_state = self.prepare_timestep_state(sigma)
+    noise = self.transformer(
+        latents,
+        refined_condition,
+        rope_tables=rope_tables,
+        attn_metadata=attn_metadata,
+        timestep_state=timestep_state,
+    )
+    latents = self.scheduler_step(noise, sigma, latents)
+```
+
+Candidates include prompt conditioning, token-refiner output, reference
+encoding/scans, position IDs/RoPE tables, packed masks/metadata, static CFG
+branches, and finite-schedule AdaLN projections. First prove they are invariant
+for all official tasks, LoRA/adapters, prompt-update features, shapes, and
+schedules. In continuous batching, store request-specific results on request
+state rather than the pipeline singleton.
+
+### Eliminate per-layer synchronization
+
+Avoid device-to-host scalar extraction such as `.item()` inside every layer:
+
+```python
+# Bad: forces a GPU/CPU synchronization in every attention layer.
+max_len = int(cu_seqlens[-1].item())
+
+# Better: compute validated host metadata once during request packing and pass it.
+packed = PackedLayout(
+    cu_seqlens=cu_seqlens,
+    max_seqlen=max(sample_lengths),
+)
+for block in self.blocks:
+    hidden = block(hidden, packed_layout=packed)
+```
+
+Trace to verify the synchronization disappears. Do not replace a correct
+dynamic value with a stale constant merely to remove `.item()`.
+
+### Remove unused masks and allocations
+
+If packed attention consumes `cu_seqlens` and ignores a dense mask, do not
+construct the dense mask. Reuse immutable request metadata and preallocated
+buffers when safe. Prefer direct indexed scatter/gather over constructing large
+temporary tensors only when duplicate-index semantics and gradients (normally
+disabled for inference) remain correct.
+
+Any cross-request cache needs a bounded lifetime and a complete key containing
+shape, layout, device, dtype, schedule, task, adapter/LoRA, and any value that
+changes the tensor. Otherwise keep it request-scoped.
+
+### VAE and reference work
+
+Profile duplicate VAE decode/postprocess, repeated reference decoding/scans,
+unnecessary host-device round trips, and full-frame materialization. Apply VAE
+tiling, slicing, patch parallelism, or on-demand staging only with seam,
+ordering, color/range, and memory parity tests.
+
+## USP communication optimization
+
+“USP works” is not evidence that USP communication is optimized. For Ulysses,
+Ring, hybrid, and AllGather-KV separately, profile:
+
+- collective count, bytes, duration, stream, and overlap with compute;
+- Q/K/V and output all-to-all boundaries;
+- pack/unpack, padding, metadata broadcast, transpose, and contiguous copies;
+- head divisibility or advanced uneven-head support;
+- scaling efficiency from 1 to N devices on the same workload/interconnect;
+- load balance for unequal packed samples.
+
+Declare `_sp_plan` split/gather boundaries around meaningful module outputs,
+not arbitrary Python statements. Keep packed metadata sharded/gathered in the
+same ownership model as tokens.
+
+When the target revision offers accelerated Ulysses transport such as
+`--ulysses-a2a-permute`, qualify it separately from regular Ulysses. Prove the
+strict scatter-heads/gather-sequence layout selected the fast path, record
+one-time JIT/readiness cost, prewarm the maximum workspace shape before CUDA
+graph capture, keep the grow-only workspace on one stream, and exercise cleanup
+before process-group destruction. Compare collective and E2E deltas after the
+dense backend and topology are frozen; keep unsupported layouts on an explicit
+regular-Ulysses fallback.
+
+For GQA, validate `Q_heads % KV_heads == 0`, pad KV heads first, and derive Q
+padding from the original ratio. Exercise both contiguous and row-strided QKV
+layouts; a fast copy path must preserve bitwise values and packed boundaries.
+
+Cross-attention sometimes uses replicated text K/V while only video/audio Q is
+sequence-sharded. In that specific case, `skip_sequence_parallel=True` can
+avoid incorrect or wasteful Q/K/V redistribution:
+
+```python
+self.cross_attention = Attention(
+    num_heads=local_heads,
+    num_kv_heads=local_kv_heads,
+    head_size=head_dim,
+    softmax_scale=head_dim**-0.5,
+    causal=False,
+    role="cross",
+    skip_sequence_parallel=True,
+    prefix=f"{prefix}.cross_attn",
+)
+```
+
+Use it only after proving K/V are replicated and Q/output ownership is correct.
+Otherwise it disables needed communication. Report whether the optimization
+reduces collective bytes/time and improves E2E scaling.
+
+## Text encoder and VAE disaggregation
+
+Disaggregation is a measured architecture decision. First collect:
+
+| Question | Required measurement |
+|---|---|
+| Is the stage large enough? | Encode/decode share at target concurrency and shape |
+| Is transfer economical? | Tensor/media bytes, serialization, connector latency and bandwidth |
+| Is there reuse? | Prompt/reference reuse rate, cacheability, fan-out |
+| Can it scale independently? | Queue saturation and resource utilization per stage |
+| Is the contract stable? | Shape, dtype, mask/tag/position metadata, ordering |
+
+### Text encoder candidate
+
+Define an explicit tensor contract containing hidden states, masks, tags or
+position IDs, dtype, shape, task/checkpoint revision, and request ID. Support
+disabling the local encoder so the DiT cannot accidentally recompute it. Test
+all tasks, prompt lengths, partial TP group membership, retries, backpressure,
+abort, connector cleanup, and independent replica scaling.
+
+### VAE candidate
+
+Separate VAE encode and decode contracts if both exist. Define latent shape,
+scale/shift, dtype/range, tiling/chunk order, output media ordering, and stage
+metadata. Measure latent/media transfer against local decode cost. Test tiling
+seams, RNG/determinism where relevant, chunk ordering, codec/postprocess,
+backpressure, abort, and worker failure.
+
+Do not label VAE tiling, patch parallelism, compile, or offload as VAE
+disaggregation. Likewise, a separate text-encoder TP group is not a separate
+stage. If no reusable generic stage/connector exists in the target revision,
+produce a feasibility result or RFC and keep the capability `not tested` rather
+than inventing a model-local production architecture.
+
+## Performance acceptance
+
+For each accepted optimization, require:
+
+- same fixed workload and correctness artifacts before/after;
+- actual fast-path execution proven by logs/trace/counters;
+- local metric plus E2E latency/throughput and memory;
+- at least three raw fixed-work repetitions after explicit warmup, or a
+  declared arrival-load sample set large enough for every tail percentile;
+- named hardware, backend, topology, dtype, and task scope;
+- fallback/rejection tests for unsupported shapes/platforms;
+- no regression in abort/error cleanup or long-run memory trend.
+
+Update the feature matrix narrowly. Do not convert a profile-driven direction
+into a generic “supported” feature merely because a model-specific PR exists.
@@ -0,0 +1,464 @@
+# Production Serving and Validation
+
+## Contents
+
+1. [Execution-mode decision](#execution-mode-decision)
+2. [Request-level batching](#request-level-batching)
+3. [Step execution and continuous batching](#step-execution-and-continuous-batching)
+4. [Abort and cleanup](#abort-and-cleanup)
+5. [Mixed-RPS soak](#mixed-rps-soak)
+6. [Four CI tracks](#four-ci-tracks)
+7. [PR evidence checklist](#pr-evidence-checklist)
+
+## Execution-mode decision
+
+Request batching and step-wise execution solve different problems:
+
+| Mode | Pipeline contract | Scheduling/abort boundary |
+|---|---|---|
+| Serial request | `forward(DiffusionRequestBatch)` for one request | Whole request; in-denoise preemption is unavailable |
+| Request batching | `supports_request_batch = True`; one forward consumes compatible independent requests | Whole batch/request |
+| Single-request step | Full `SupportsStepExecution`; `max_num_seqs=1` | Between denoise steps |
+| Step continuous batching | Full step contract plus heterogeneous batched-step evidence | Between denoise waves/steps |
+
+Do not expose four stub methods as an opt-out. The current protocol is
+structural; if a native pipeline is not step-capable, do not define the four
+step methods. If it is capable, implement and test all four.
+
+Treat structural support, a generic request-batch bridge, and model benefit as
+three separate claims. Before recommending step execution, name a hypothesis
+that request mode cannot satisfy, define a success threshold, and run a matched
+request-mode control. If throughput or latency regresses, retain step execution
+as functional/limited and keep request mode as the production recommendation.
+
+Choose compatibility keys from every value that changes tensor shapes,
+semantics, hooks, or scheduler behavior: task, output shape/count, frames,
+steps/schedule, guidance/CFG, dtype/backend, LoRA/adapter, cache quality policy,
+and model-specific conditions. Reject or schedule separately rather than
+silently using the first request's settings.
+
+## Request-level batching
+
+A request-batched pipeline returns exactly one `DiffusionOutput` per logical
+request. This is model-owned pseudocode: `validate_batch_compatibility()` and
+`_generate_batch()` are placeholders to implement from the official contract,
+not shared APIs to import.
+
+```python
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.worker.request_batch import (
+    DiffusionRequestBatch,
+    split_diffusion_output_by_request,
+)
+
+class MyPipeline(torch.nn.Module):
+    supports_request_batch = True
+
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
+        validate_batch_compatibility(req.sampling_params_list, req.prompts)
+        common = req.sampling_params_list[0]
+
+        text_prompts: list[str] | None = [
+            prompt
+            if isinstance(prompt, str)
+            else (prompt.get("prompt") or "")
+            for prompt in req.prompts
+        ]
+
+        prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
+            req.prompts,
+            {"prompt_embeds": None, "prompt_mask": None},
+        )
+        if prompt_fields["prompt_embeds"] is not None:
+            text_prompts = None
+        generators = req.collate_request_generators(
+            common.num_outputs_per_prompt,
+            default_generator=None,
+        )
+        output = self._generate_batch(
+            prompt=text_prompts,
+            prompt_fields=prompt_fields,
+            generators=generators,
+            sampling=common,
+        )
+        return split_diffusion_output_by_request(
+            DiffusionOutput(output=output),
+            req,
+            num_outputs_per_prompt=common.num_outputs_per_prompt,
+        )
+```
+
+The prompt-field collator only handles tensor fields; extract raw text prompts
+and model-specific negative/alias fields separately, as current in-tree batched
+pipelines do. Set a raw prompt field to `None` when the matching precomputed
+embedding is supplied. The exact aliases and generator device depend on the
+model. Use current collators rather than manual `torch.cat` when they cover the
+contract; they reject mixed present/missing tensors and incompatible
+shape/dtype/device.
+
+Before reading `sampling_params_list[0]`, validate every field that must be
+common. Tests must cover:
+
+- two and N requests, independent seeds/generators, one output each;
+- mixed prompt lengths and packed/padded boundaries;
+- multiple outputs per prompt and correct result slicing;
+- compatible versus incompatible shapes, steps, guidance, LoRA, and quality;
+- one request failure without cross-request output or state corruption;
+- output ordering and request IDs under different completion times.
+
+After the model advertises support, candidate serving is:
+
+```bash
+vllm serve '<model>' --omni \
+  --max-num-seqs 4 \
+  --request-batch-max-wait-ms 20
+```
+
+The wait window trades admission latency for batch formation. Benchmark at 0
+and candidate nonzero values under the target arrival process.
+
+## Step execution and continuous batching
+
+### Full protocol skeleton
+
+Use the current shared state types and make every mutable value request-local:
+
+The helpers beginning with `_` and `validate_step_batch()` below are
+model-owned placeholders; scheduler initialization is scheduler-specific.
+
+```python
+import copy
+from collections.abc import Sequence
+from typing import Any, ClassVar
+
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.interface import SupportsStepExecution
+from vllm_omni.diffusion.worker.input_batch import InputBatch
+from vllm_omni.diffusion.worker.utils import StepRequestState
+
+class MyPipeline(torch.nn.Module, SupportsStepExecution):
+    supports_step_execution: ClassVar[bool] = True
+
+    def prepare_encode(
+        self,
+        state: StepRequestState,
+        **kwargs: Any,
+    ) -> StepRequestState:
+        prepared = self._prepare_request(state.prompt, state.sampling)
+        state.prompt_embeds = prepared.prompt_embeds
+        state.latents = prepared.latents
+        state.timesteps = prepared.timesteps
+        state.scheduler = copy.deepcopy(self.scheduler)
+        state.scheduler.set_begin_index(0)
+        state.step_index = 0
+        state.extra["task_state"] = prepared.task_state
+        return state
+
+    def denoise_step(
+        self,
+        input_batch: InputBatch,
+        *,
+        states: Sequence[StepRequestState] | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor | None:
+        validate_step_batch(input_batch, states)
+        return self.transformer(**build_batched_denoise_inputs(input_batch))
+
+    def step_scheduler(
+        self,
+        state: StepRequestState,
+        noise_pred: torch.Tensor,
+        **kwargs: Any,
+    ) -> None:
+        timestep = state.current_timestep
+        if timestep is None:
+            raise RuntimeError("step_scheduler called without a current timestep")
+        state.latents = state.scheduler.step(
+            noise_pred,
+            timestep,
+            state.latents,
+            return_dict=False,
+        )[0].to(state.latents.dtype)
+        state.step_index += 1
+
+    def post_decode(
+        self,
+        state: StepRequestState,
+        **kwargs: Any,
+    ) -> DiffusionOutput:
+        return DiffusionOutput(output=self.decode(state.latents))
+```
+
+This is a shape guide, not drop-in code: verify the target revision's
+`StepRequestState.extra`, scheduler signature, CFG handling, state fields, and
+decode contract. Follow a current native pipeline such as Qwen-Image for the
+exact revision.
+
+Never use the shared `self.scheduler` to advance a request. Store scheduler,
+latents, embeddings, timesteps, masks, RNG/generator, task metadata, cache
+policy, adapter, and output assembly state on `StepRequestState` or its
+request-local extension. Pipeline fields may hold immutable weights/config and
+bounded immutable caches only.
+
+### Qualification sequence
+
+1. Compare request-mode and step-mode at `max_num_seqs=1` for every task using
+   identical seed/schedule/shape and intermediate latents.
+2. Abort queued and in-flight requests at several step boundaries.
+3. Increase to two compatible requests; verify independent scheduler/RNG/state
+   and output order.
+4. Increase to the target capacity and mix prompt lengths.
+5. Exercise mixed shape, frames, steps/schedule, task, guidance, and quality.
+   Batch only compatible rows; otherwise split/reject predictably.
+6. Inject a failure in encode, denoise, scheduler, and decode independently.
+7. Run the production soak and fault matrix.
+
+Candidate commands only after capability validation:
+
+```bash
+# Conservative step execution and step-boundary abort.
+vllm serve '<model>' --omni \
+  --step-execution \
+  --max-num-seqs 1
+
+# Experimental batched-step candidate after multi-request E2E.
+vllm serve '<model>' --omni \
+  --step-execution \
+  --max-num-seqs 4
+```
+
+If an attention backend is only correct for single-request step mode, fail
+early or keep the recipe at `max_num_seqs=1`. Do not describe generic runtime
+support as model support.
+
+## Abort and cleanup
+
+### Required semantics
+
+For queued and in-flight abort/disconnect:
+
+- emit exactly one terminal aborted/error result;
+- do not run `post_decode` after abort unless a documented partial-output mode
+  explicitly requires it;
+- do not increment `step_index` or apply a scheduler result after cancellation;
+- remove scheduler/state cache entries and request IDs;
+- release latents, prompt/reference tensors, cache hooks, DLO buffers, connector
+  messages, uploaded/temp files, and output buffers;
+- keep other requests alive;
+- make cleanup idempotent when disconnect, timeout, and worker error race;
+- drop late results/exceptions for already-resolved or cancelled futures
+  without terminating the shared result-pump thread;
+- complete the next valid request with its requested cache/offload policy.
+
+Test cancellation before engine admission, while queued, after encode, during
+early/middle/final denoise steps, before decode, and while returning/persisting
+output. Include async job cancellation and client socket disconnect.
+
+### Failure injection matrix
+
+| Location | Injection | Assertions |
+|---|---|---|
+| Validation/upload | bad MIME, count/bytes overflow, truncated media | 400/413, no job/engine submit, no temp leak |
+| Load/start | missing fused shard, unsupported topology | startup fails, actionable error, no silent init |
+| Encode | exception/OOM/timeout | one error, request state/temp released |
+| Denoise | exception/OOM/worker exit | peers handled, hooks/buffers released, health behavior explicit |
+| Scheduler | invalid state/shape | no latent/step partial commit |
+| Decode/post | exception/codec error | no corrupt successful result, partial output removed |
+| Result/IPC | late result after cancel, large split tensor, orphan SHM | pump remains alive, one terminal result, exact ownership cleanup |
+| Connector/disagg | timeout/retry/duplicate | idempotency, backpressure, no orphan payload |
+
+After each injection, check server health and issue a known-good request. If
+the process must restart, document and test the restart policy; do not claim
+transparent recovery.
+
+### Large-output transport
+
+Treat worker-to-engine IPC and engine-to-client transport as separate hops.
+For batched video, prove that each logical output does not accidentally retain
+or serialize the complete batch storage. Exercise the target revision's shared
+memory or artifact-handle path above and below its threshold, including
+non-contiguous views, multiple outputs, consumer failure, timeout, abort, and
+worker exit. Verify every shared segment/file is released exactly once.
+
+Freeze the media contract at each hop: dtype, value range, shape/layout,
+contiguity, request ownership, encoding/color model, payload bytes, and pending
+consumers. Device-side float-to-uint8 preparation, D2H/IPC, planar/interleaved
+conversion, CPU encode/mux, and delivery off the asyncio event loop are
+different optimizations. Validate their route selection and fallbacks
+independently. Byte-identical HTTP MP4 output does not prove that a raw offline
+tensor kept its dtype or range.
+
+For long-window output, validate segment ordering and audio/video continuity,
+then exercise payloads around serializer, shared-memory, HTTP, and codec size
+limits. Chunking must reassemble exactly, reject malformed/missing chunks, and
+release partial state after abort, timeout, or peer failure.
+
+Qualify chunk production, transport, and public streaming as separate layers.
+A VAE callback must define owned representation, frame offsets, final marker,
+rank ownership, source/config gate, and collective-safe failure, but it does not
+prove SHM/ZMQ delivery, bounded backpressure, cancellation, codec reuse, or
+client streaming. Measure time to first owned chunk, first encoded fragment,
+steady cadence, mux finalization, and complete client artifact independently.
+
+Measure remote MP4 encoding, serialization/copies, network transfer, and client
+materialization separately from denoise/decode. An inference-kernel speedup does
+not establish serving throughput if output transport dominates or leaks.
+
+## Mixed-RPS soak
+
+Use a realistic request mix and arrival distribution. Include small/large
+shapes, short/long prompts, all advertised tasks, cache tiers, and a small rate
+of abort/disconnect/error injections. Run phases below, near, and above measured
+saturation so queue/backpressure behavior is visible.
+
+Example serving benchmark phase:
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+  --base-url http://127.0.0.1:8091 \
+  --endpoint /v1/videos \
+  --model '<model>' \
+  --dataset trace \
+  --dataset-path /path/to/versioned-production-trace.jsonl \
+  --task t2v \
+  --num-prompts 500 \
+  --request-rate 0.5 \
+  --max-concurrency 8 \
+  --warmup-requests 8 \
+  --warmup-concurrency 8 \
+  --output-file soak-near-saturation.json
+```
+
+Use a task-appropriate trace/media dataset. Run the same phase at rates selected
+from an initial saturation curve; a literal `0.5` is illustrative only.
+
+On supported GPU hosts, wrap a test or soak with the repository monitor:
+
+```bash
+bash tests/dfx/stability/scripts/resource_monitor.sh run --backend gpu -- \
+  pytest -s -v tests/dfx/stability/scripts/test_stability_<model>.py -m slow
+```
+
+Check the target revision's monitor backends. Do not claim CPU/NPU monitoring
+from a reserved CLI choice unless collection is implemented and verified.
+
+Report:
+
+- offered and achieved RPS, success/error/abort rates;
+- p50/p95/p99 and queue/service/stage times, with the successful sample count
+  and arrival process used for each percentile;
+- throughput per device and batch/wave utilization;
+- per-rank HBM and process-tree PSS trend/slope;
+- cache size/hit rate, temp-file/disk growth, open FDs/handles;
+- outstanding shared-memory/artifact handles and encoded/network bytes;
+- worker/engine health, restarts, OOM/timeouts, backpressure response;
+- total duration, request count/mix, raw logs/JSON/monitor artifact hashes.
+
+A short burst or repeated serial curl loop is not a stability test.
+
+## Four CI tracks
+
+Keep the tracks independent so a failure answers a specific question. These
+four readiness tracks are not the repository's numbered CI levels. Map them
+onto the current L1-L5 taxonomy:
+
+| Readiness track | Repository CI placement |
+|---|---|
+| Function | L1 unit/logic plus L2 basic online/offline E2E; extend real-model scenarios in L3/L4 |
+| Accuracy | L3 real-model gates for important cases; deeper and broader nightly coverage in L4 |
+| Performance | Important thresholds in L3 where they fit the time budget; full baselines/regression jobs in nightly L4 |
+| Reliability | Cheap rejection/recovery assertions can run earlier, but long stability, fault injection, and reliability suites belong to weekly L5 |
+
+Distributed topology is a test case within the appropriate level, not a
+replacement definition for a level.
+
+### 1. Function CI
+
+Cover:
+
+- strict unquantized and advertised quantized loading;
+- fused source-shard completeness and prefix routing;
+- official task positive boundaries and neighboring negative cases;
+- offline, sync, async normalization/parity;
+- single-device reference plus every advertised topology smoke;
+- output type/count/order and request ID isolation;
+- feature rejection paths that must fail early.
+
+### 2. Accuracy CI
+
+Pin official implementation/checkpoint/vLLM-Omni revisions, prompts and media
+hashes, seed, scheduler/sigma/steps, dimensions/frames/FPS, guidance, dtype,
+backend, and topology. Compare at suitable levels:
+
+- component outputs or selected blocks;
+- one or more intermediate denoise latents;
+- final image/video/audio artifact metrics;
+- temporal consistency and audio/video synchronization where relevant.
+
+Set tolerance per modality/hardware/backend with rationale. Keep failure
+artifacts and metadata. A single visual inspection does not define CI.
+
+For quantization trajectory tooling, inspect the target revision's current
+`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` and
+reuse it when compatible instead of creating an ad hoc comparator.
+
+### 3. Performance CI
+
+Pin workload and best deployment. Include a resident recommended row and a
+small-HBM/DLO row when both are production targets. Declare:
+
+- one explicit warmup exclusion and at least three measured fixed-work
+  repetitions;
+- raw result JSON and environment identity;
+- median/mean plus range for small fixed-work A/Bs; p50/p95/p99 only for an
+  arrival-load run with enough samples, plus throughput/RPS, stage time,
+  HBM/PSS;
+- regression metric, threshold, baseline owner, and update procedure;
+- correctness/accuracy check linked to the same configuration.
+
+Use the current harness config style, for example:
+
+```bash
+pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
+  --test-config-file tests/dfx/perf/tests/test_<model>_vllm_omni.json
+```
+
+Verify parser options on the target revision. Do not copy a JSON schema from a
+different model without matching its endpoint, task, dataset, and deploy path.
+
+### 4. Reliability CI
+
+Run a bounded CI soak plus a longer scheduled soak. Cover:
+
+- below/near/above-saturation arrival phases;
+- queued/in-flight abort and disconnect;
+- OOM/timeout and worker/process failure appropriate to the platform;
+- DLO/cache/connector cleanup;
+- health, fast-fail/backpressure, restart/recovery policy;
+- memory/temp/FD trend and successful known-good request after each fault.
+
+Keep reliability thresholds independent from latency thresholds. A fast server
+that leaks memory or hangs after cancellation fails production readiness.
+
+## PR evidence checklist
+
+The PR description must include or link:
+
+- official source/checkpoint revisions and API/limit matrix;
+- scoped feature compatibility matrix with the four allowed states;
+- implementation notes for FP8, DLO, request cache, attention/fusion, batching;
+- exact environment, deploy/serve, every task curl, benchmark and soak commands;
+- fixed accuracy inputs and artifacts;
+- raw before/after performance runs and memory traces;
+- output-contract manifests and stage accounting from device preparation
+  through the complete client artifact;
+- abort/disconnect/fault results and cleanup evidence;
+- Function/Accuracy/Performance/Reliability CI files and ownership;
+- limitations/TODOs that remain `not tested` rather than implied support.
+
+Separate the Day-0 milestone from production readiness. If the PR only closes
+some gates, title and description should say which scoped gates it closes.
@@ -37,6 +37,10 @@ self-check.
 - [`precheck-pr`](precheck-pr/SKILL.md): self-checks a branch before creating a
  PR by validating title format, dead code, simplification opportunities,
  accuracy and performance claims, and merge readiness
+- [`production-add-diffusion-model`](production-add-diffusion-model/SKILL.md):
+  takes a working Day-0 diffusion integration through API/limit parity,
+  feature-combination evidence, quantization/offload, performance, production
+  serving, hardware recipes, and CI readiness
 - [`quantization`](quantization/SKILL.md): guides quantization method selection,
  model integration, checkpoint loading, and quality/performance validation
  for vLLM-Omni
@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Benchmark diffusion LoRA expand-and-accumulate implementations.
+
+This compares the previous functional expand path with direct accumulation into
+the selected output slice. It is a synthetic kernel microbenchmark and does not
+load a diffusion model.
+
+Example:
+    python benchmarks/kernels/benchmark_diffusion_lora_expand.py \
+        --m 1,16,128,512,4096,16384 \
+        --in-dim 4096 --out-dim 4096 --rank 64 \
+        --dtype bfloat16 --packed --warmup 30 \
+        --small-iterations 500 --large-iterations 100
+"""
+
+import argparse
+import math
+import statistics
+from collections.abc import Callable
+
+import torch
+
+MiB = 1024**2
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--m", default="1,16,128,512,4096,16384")
+    parser.add_argument("--in-dim", type=int, default=4096)
+    parser.add_argument("--out-dim", type=int, default=4096)
+    parser.add_argument("--rank", type=int, default=64)
+    parser.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
+    parser.add_argument(
+        "--packed",
+        action="store_true",
+        help="Accumulate into the second half of a packed output tensor.",
+    )
+    parser.add_argument("--warmup", type=int, default=30)
+    parser.add_argument("--small-iterations", type=int, default=500)
+    parser.add_argument("--large-iterations", type=int, default=100)
+    parser.add_argument("--large-m-threshold", type=int, default=1024)
+    return parser.parse_args()
+
+
+def percentile(samples: list[float], probability: float) -> float:
+    ordered = sorted(samples)
+    position = probability * (len(ordered) - 1)
+    lower = math.floor(position)
+    upper = math.ceil(position)
+    if lower == upper:
+        return ordered[lower]
+    weight = position - lower
+    return ordered[lower] * (1 - weight) + ordered[upper] * weight
+
+
+def measure_latency_ms(operation: Callable[[], None], warmup: int, iterations: int) -> list[float]:
+    for _ in range(warmup):
+        operation()
+    torch.accelerator.synchronize()
+
+    starts = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)]
+    ends = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)]
+    for start, end in zip(starts, ends, strict=True):
+        start.record()
+        operation()
+        end.record()
+    torch.accelerator.synchronize()
+    return [start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)]
+
+
+def measure_extra_peak_mib(operation: Callable[[], None]) -> float:
+    torch.accelerator.synchronize()
+    torch.accelerator.reset_peak_memory_stats()
+    allocated = torch.accelerator.memory_allocated()
+    operation()
+    torch.accelerator.synchronize()
+    peak = torch.accelerator.max_memory_allocated()
+    return (peak - allocated) / MiB
+
+
+def print_header(args: argparse.Namespace, dtype: torch.dtype) -> None:
+    device_name = torch.cuda.get_device_name()
+    print(f"device={device_name} torch={torch.__version__} dtype={dtype}")
+    print(f"in_dim={args.in_dim} out_dim={args.out_dim} rank={args.rank} packed={args.packed} warmup={args.warmup}")
+    print("M       iters  old_p10  old_p50  old_p90  new_p10  new_p50  new_p90  speedup  old_MiB  new_MiB  max_abs")
+
+
+def benchmark_shape(m: int, args: argparse.Namespace, dtype: torch.dtype) -> None:
+    device = torch.device("cuda")
+    output_width = args.out_dim * (2 if args.packed else 1)
+    offset = args.out_dim if args.packed else 0
+    output_slice = slice(offset, offset + args.out_dim)
+    scale = 1 / math.sqrt(args.in_dim)
+
+    x = torch.randn(m, args.in_dim, device=device, dtype=dtype)
+    lora_a = torch.randn(args.rank, args.in_dim, device=device, dtype=dtype) * scale
+    lora_b = torch.randn(args.out_dim, args.rank, device=device, dtype=dtype) / math.sqrt(args.rank)
+    initial_output = torch.randn(m, output_width, device=device, dtype=dtype)
+
+    def baseline(output: torch.Tensor) -> None:
+        buffer = x @ lora_a.t()
+        delta = buffer @ lora_b.t()
+        selected = output[:, output_slice]
+        selected[:] = selected + delta
+
+    def optimized(output: torch.Tensor) -> None:
+        buffer = x @ lora_a.t()
+        selected = output[:, output_slice]
+        torch.addmm(selected, buffer, lora_b.t(), out=selected)
+
+    reference = initial_output.clone()
+    candidate = initial_output.clone()
+    baseline(reference)
+    optimized(candidate)
+    torch.testing.assert_close(candidate, reference, rtol=1e-2, atol=1e-2)
+    max_abs = (candidate - reference).abs().max().item()
+
+    old_output = initial_output.clone()
+    new_output = initial_output.clone()
+    iterations = args.large_iterations if m >= args.large_m_threshold else args.small_iterations
+
+    old_peak = measure_extra_peak_mib(lambda: baseline(old_output))
+    new_peak = measure_extra_peak_mib(lambda: optimized(new_output))
+    old_samples = measure_latency_ms(lambda: baseline(old_output), args.warmup, iterations)
+    new_samples = measure_latency_ms(lambda: optimized(new_output), args.warmup, iterations)
+
+    old_p50 = statistics.median(old_samples)
+    new_p50 = statistics.median(new_samples)
+    print(
+        f"{m:<7} {iterations:<6} "
+        f"{percentile(old_samples, 0.1):>8.5f} {old_p50:>8.5f} {percentile(old_samples, 0.9):>8.5f} "
+        f"{percentile(new_samples, 0.1):>8.5f} {new_p50:>8.5f} {percentile(new_samples, 0.9):>8.5f} "
+        f"{old_p50 / new_p50:>8.3f} {old_peak:>8.2f} {new_peak:>8.2f} {max_abs:>8.5f}"
+    )
+
+
+def main() -> None:
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("This benchmark requires a CUDA GPU")
+    if args.warmup < 0 or args.small_iterations <= 0 or args.large_iterations <= 0:
+        raise ValueError("warmup must be non-negative and iteration counts must be positive")
+
+    dtype = getattr(torch, args.dtype)
+    torch.manual_seed(0)
+    print_header(args, dtype)
+    with torch.inference_mode():
+        for m in (int(value) for value in args.m.split(",")):
+            if m <= 0:
+                raise ValueError(f"M must be positive, got {m}")
+            benchmark_shape(m, args, dtype)
+
+
+if __name__ == "__main__":
+    main()
@@ -132,3 +132,14 @@ models:
    endpoint: /v1/audio/speech
    task_extra_body:
      voice_design: {}
+
+  # OmniVoice serves plain text-to-speech over /v1/audio/speech; voice cloning
+  # and voice design are offline-only for this model, so default_voice is the
+  # only task the server can answer.
+  k2-fsa/OmniVoice:
+    supported_tasks: [default_voice]
+    backend: openai-audio-speech
+    endpoint: /v1/audio/speech
+    trust_remote_code: true
+    task_extra_body:
+      default_voice: {}
@@ -179,7 +179,7 @@ collection, examples, bug reports, or logs.

 ## Model-specific variables

-The audit found 56 variables read by a single model or pipeline family. They are
+The audit found 58 variables read by a single model or pipeline family. They are
 not listed as public usage options here because doing so would turn implementation
 escape hatches into an accidental compatibility contract.

@@ -188,7 +188,7 @@ Every audited model-specific name has a migration disposition in the

 | Disposition | Count | Required outcome |
 | --- | ---: | --- |
-| Promote | 33 | Move a stable setting into typed stage or model configuration. |
+| Promote | 35 | Move a stable setting into typed stage or model configuration. |
 | Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
 | External | 0 | Retain only when a supported third-party library owns the contract. |
 | Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |
@@ -57,13 +57,11 @@ Canonical layout (prefer these paths for new changes):
 | Bootstrap vs test YAML | **Bootstrap** (`pipeline*.yml`) builds images and uploads **child** test pipelines. **Test** YAML (`test-*.yml`) lists pytest steps only. |
 | Register CI YAML in `skip_ci.py` | If you add a new whitelisted test pipeline file, update `L2_YAML_FILES`, `L3_YAML_FILES`, or `L45_YAML_FILES` in `.buildkite/common/scripts/skip_ci.py` so skip-ci paths stay correct. |

-There are still **legacy copies** at `.buildkite/*.yaml` (without the `cuda/` prefix). Treat `.buildkite/cuda/*` as source of truth.
-
 ## Platform comparison

 | Platform | Bootstrap entry | Test job files | Upload mechanism | Job hardware in YAML |
 | -------- | ---------------- | -------------- | ---------------- | -------------------- |
-| **CUDA** | `cuda/pipeline.yml` | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml`, `test-weekly.yml` | `upload_pipeline.py --upload` (expands uploader-only keys) | `mirror_hardwares: <preset>` (string) |
+| **CUDA** | `cuda/pipeline.yml` | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml`, `test-weekly.yml` | `upload_pipeline.py --upload` (expands uploader-only keys) | omit `mirror_hardwares` (from `-m`) / preset string + `MIRROR_HW` |
 | **NPU** | `npu/pipeline-npu.yml` | `test-npu-ready.yml`, `test-npu-nightly.yml` | `upload_pipeline.py --upload` | `mirror_hardwares: a2b3_npu_1` / `a2b3_npu_4` / `a3_npu_2` |
 | **AMD** | `amd/scripts/bootstrap-amd-omni.sh` | `test-amd-ready.yml`, `test-amd-merge.yml` | Jinja (`test-template-amd-omni.j2`) → `pipeline upload` | `agent_pool` + `mirror_hardwares: [amdproduction]` (array, template filter) |
 | **Intel** | `intel/scripts/bootstrap-intel-omni.sh` | `intel/pipeline-intel.yml` (steps inline) | Direct `pipeline upload` | Inline `agents.queue` on each step |
@@ -92,7 +90,23 @@ There are still **legacy copies** at `.buildkite/*.yaml` (without the `cuda/` pr

    **Upload:** `upload_pipeline.py --upload` expands uploader-only keys before Buildkite upload.

-    **Hardware in YAML:** `mirror_hardwares: <preset>` (string)—preset names in [`common/ci_mirror_hardwares.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/ci_mirror_hardwares.yml). Do **not** set `agents` / `plugins` on the same step.
+    **Hardware in YAML:** omit `mirror_hardwares` to compose `{chip}_{n}` from pytest `-m` (SKU + `cards_n`). Several positive `cards_*` take the **max** n (`cards_2 and cards_3` → `{chip}_3`). `not cards_1` with no positive `cards_*` uses that chip's **highest existing** preset in `ci_mirror_hardwares.yml` (L4 → `l4_4`, H100 → `h100_4`, B200 → `b200_4`). Set a **preset string** to force a pool (marks are ignored). Do **not** set `agents` / `plugins` on the same step as `mirror_hardwares`. `MIRROR_HW=b200` remaps inferred jobs whose `-m` includes `B200`, and omits CUDA `h100_*` / `l4_*` strings. Jobs whose `-m` does not match (unset: no H100/L4; set: no B200) are **skipped**.
+
+    ```yaml
+    # inferred: H100+B200 + cards_2 → h100_2, or b200_2 when MIRROR_HW=b200
+    commands:
+      - pytest -sv tests/e2e -m "full_model and H100 and B200 and omni and cards_2"
+    # no mirror_hardwares
+
+    # inferred: not cards_1 (no positive cards_*) → that chip's highest pool (l4_4 / h100_4 / b200_4)
+    commands:
+      - pytest -sv tests/e2e -m "full_model and diffusion and L4 and not cards_1"
+
+    # forced preset (ignores -m SKU/cards for pool selection)
+    mirror_hardwares: h100_1
+    ```
+
+    Inferred form: unset/`empty` `MIRROR_HW` matches `H100` or `L4` in `-m` (both present → H100); no match skips the step. `MIRROR_HW=b200` must appear in `-m` or the step is skipped. The uploader does **not** rewrite pytest `-m`; B200 collection comes from the YAML expression and from tests that declare `B200` in `hardware_test` / `hardware_marks` (for example `res={"cuda": ["H100", "B200"]}`). `MIRROR_HW` must be empty or `b200` (case-insensitive); unknown values (for example `b20o`) **fail the upload**. A CUDA preset string such as `mirror_hardwares: h100_4` is skipped when `MIRROR_HW=b200`. Unset `MIRROR_HW` keeps string presets. NPU presets ignore `MIRROR_HW`. A `mirror_hardwares` name that is not a key in `ci_mirror_hardwares.yml` **fails the upload**.

    **Conventions**

@@ -103,9 +117,9 @@ There are still **legacy copies** at `.buildkite/*.yaml` (without the `cuda/` pr

    **Adding a job**

-    1. Pick the level file (ready / merge / nightly / weekly).
+    1. Pick the level file (`test-ready.yml` / `test-merge.yml` / `test-nightly.yml` / `test-weekly.yml`).
    2. Add a step under the right **group** (usually **E2E Test** for model pytest).
-    3. Set `label`, `commands`, `mirror_hardwares`, `depends_on: upload-<level>-pipeline`.
+    3. Set `label`, `commands`, `mirror_hardwares`, `depends_on: upload-ready-pipeline` (or `upload-merge-pipeline` / `upload-nightly-pipeline` / `upload-weekly-pipeline`).
    4. For L2/L3 E2E, add `source_file_dependencies` (pytest + model + deploy YAML prefixes).
    5. Dry-run:

@@ -414,10 +428,12 @@ 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.
+The upload depends on `buildkite-agent` being callable inside the container.
+CUDA (`l4_*`, `h100_*`) and NPU (`*_npu_*`) presets all use the `kubernetes`
+plugin, which provides the agent. `l4_*` jobs run on the EKS `l4-k8s` queue
+(agent-stack-k8s); they no longer use the docker plugin or
+`mount-buildkite-agent`. A new **docker** preset that runs a coverage job
+still needs `mount-buildkite-agent: true`.

 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
@@ -136,7 +136,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
        <a href="test_writing_guide.md#l4-level-testing-full-functionality-performance-and-documentation-testing">L4</a><br>
        L4: Purpose, Test Content, Directory Location, Example
      </td>
-      <td>Nightly</td>
+      <td>Nightly / Days before Release<br>(additional GPU SKUs)</td>
      <td>GPU</td>
    </tr>
    <tr>
@@ -175,6 +175,15 @@ Before entering specific testing levels, the project establishes two common spec

 ## Notes

+### L4 cadence (nightly and pre-release)
+
+CUDA **L4** (`full_model`) jobs are defined once in [`.buildkite/cuda/test-nightly.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-nightly.yml). That file is reused at two cadences:
+
+- **Nightly:** scheduled `main` builds with `NIGHTLY=1` (and PR labels such as `nightly-test` / `omni-test`). Default fleet is typically H100.
+- **Pre-release:** the same L4 suite runs again on extra GPU SKUs before a release (for example B200 via `MIRROR_HW=b200`), so `full_model` coverage is not limited to the nightly machine type.
+
+How count-form `mirror_hardwares` and `MIRROR_HW` pick a SKU: [CI Settings](./ci_settings.md).
+
 ### L2 / L3 diff-aware CI (CUDA)

 On CUDA **L2** and **L3**, **E2E Test** Buildkite jobs may be omitted at pipeline upload when the PR diff does not touch their path prefixes; other groups still always upload. Full two-layer mechanics and YAML examples: [CI Settings — Diff-aware CI](./ci_settings.md#diff-aware-ci) ([step filtering](./ci_settings.md#step-filtering)).
@@ -1,90 +1,197 @@
 # Skip-Softmax

-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.
+Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. It implements
+[BLASST](https://arxiv.org/abs/2512.12087) (Dynamic BLocked Attention Sparsity via Softmax
+Thresholding), which adds a per-tile skip test to the FlashAttention main loop. This page describes the algorithm, how the user configuration is resolved
+into the value the kernel consumes, and how checkpoint calibration and timestep gating participate
+in that resolution. Configuration keys and operating-point guidance are covered in
+[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax); the
+backend selection contract is covered in
+[Diffusion Attention Backend Selection](attention_backend_selection.md).

 ## Motivation

 In a long attention row (a video DiT can have tens of thousands of keys), the softmax weight
 concentrates on a small fraction of the keys; the rest receive near-zero weight and barely move the
 output. Computing softmax and the value-weighted sum over those keys is wasted work. Skip-Softmax
-detects, per block of keys, when a block cannot matter and skips its softmax and its value
-multiply.
+detects, per tile of keys, when a tile cannot matter and skips its softmax and its value multiply.

-It is approximate: a skipped block still carries a small non-zero contribution, so the mode is
+It is approximate: a skipped tile still carries a small non-zero contribution, so the mode is
 opt-in and off by default.

-## The online-softmax pass
+## Algorithm

-Attention is computed in a single streaming pass over the keys, in tiles of 128. Per query row the
+![The BLASST algorithm: FlashAttention with a per-tile skip test on the local maximum](../figures/skip_softmax/blasst_algorithm.jpg)
+
+FlashAttention computes attention in a single streaming pass over KV tiles. Per query row the
 kernel maintains three running values:

 - `m` — the largest score seen so far,
 - `l` — the running denominator `Σ exp(sⱼ − m)`,
 - `O` — the running numerator `Σ exp(sⱼ − m)·vⱼ`,

-and returns `O / l` at the end. For each key tile it computes the tile's scores `Q · K_jᵀ · scale`,
-updates `m`, and accumulates that tile's contribution into `l` and `O`. Rescaling `l` and `O` when
-`m` grows keeps the single pass numerically exact.
+and returns `O / l` at the end. For each KV tile it computes the scores `QK_j^T`, updates `m`, and
+accumulates the tile's contribution into `l` and `O`. Rescaling `l` and `O` when `m` grows keeps the
+dense pass numerically exact.

-## The skip test
-
-Once a tile's scores are known, its largest score `tile_max` is compared against the running
-maximum:
+A tile is a block of query rows against a block of keys, and the kernel processes it as a unit.
+BLASST inserts one test between computing a tile's scores and accumulating them. For each query
+row `i` in the tile, let `tile_max[i]` be the row's largest score within the tile, `running_max[i]`
+the row's running maximum, and `λ` the effective threshold:

 ```text
-if exp(tile_max − running_max) < threshold:
-    skip this tile          # do not compute its softmax weights or its P·V accumulation
+if all(exp(tile_max[i] - running_max[i]) < λ for every row i in the tile):
+    skip this tile          # do not compute its Softmax or PV contribution
 ```

-`exp(tile_max − running_max)` is an upper bound on the softmax weight any key in the tile can
-receive: if even the tile's best key is far below the current maximum, every key in the tile is
-negligible, and both the softmax (the exponentials) and the `P·V` accumulation for that tile can be
-dropped. The tile's contribution to `l` and `O` is simply skipped.
+`exp(tile_max[i] − running_max[i])` is an upper bound on the softmax weight any key in the tile can
+receive in row `i`: if even the row's best key in this tile is far below the row's current maximum,
+the tile is unimportant for that row. Only when this holds for every row in the tile can the
+Softmax and `PV` work for the tile be skipped, in which case its contribution to `l` and `O` is
+simply omitted. A larger `λ` makes the test more aggressive. The figure states the same test in
+log space, `m̃ − m < ln(λ)`, which is how the kernel evaluates it.

-## What this bounds
+The test needs the tile's scores, so `QK_j^T` always runs and only the Softmax and `PV` work can
+be skipped. Since the two matmuls have the same FLOP count, skipping every eligible tile removes at
+most roughly half of the attention arithmetic, and the kernel-level speedup stays well under 2×.

-Two properties of the test shape the achievable speedup:
+## From configuration to the kernel threshold

- **`Q · K_jᵀ` always runs.** The test needs `tile_max`, which comes from the tile's scores, so the
-  score matmul is never skipped — only the softmax and the `P·V` accumulation are. The score matmul
-  and the value matmul are comparable in cost, so skipping every eligible tile removes at most
-  roughly half the attention work; the kernel-level speedup is bounded well under 2×, not unbounded.
-
- **The decision is per tile, not per key.** A tile is skipped only when *all* of its keys are
-  collectively negligible; a single important key keeps the whole tile. How many tiles actually
-  qualify depends on the data and rounds down to tile granularity, so `target_sparsity` selects an
-  operating point on a calibrated curve — it is not a promise that a fixed fraction of tiles is
-  skipped.
-
-## The threshold
-
-The per-tile threshold is normalized by sequence length:
+The FlashInfer kernel does not take `λ` directly. It takes a `threshold_scale_factor` and divides
+it by the KV sequence length of the call:

 ```text
-threshold = factor / seqlen
+λ = threshold_scale_factor / max_kv_len
 ```

-`factor` comes from one of two sources:
+vLLM-Omni resolves the factor from one of two mutually exclusive user controls:

- **Calibrated** — `factor = a · exp(b · target_sparsity)`, where `a` and `b` are fit per model (and
-  per expert, for a multi-expert model) so that a given `target_sparsity` lands near that fraction of
-  skipped tiles on the calibration data.
- **Direct** — `factor = skip_softmax.threshold · seqlen`, the calibration-free path; the user sets
-  the threshold themselves.
+| Control | Factor passed to the kernel | Resulting `λ` |
+| --- | --- | --- |
+| `threshold` | `threshold * max_kv_len` | `threshold` |
+| `target_sparsity=s` | `a * exp(b * s)` | `a * exp(b * s) / max_kv_len` |
+
+`threshold` is therefore `λ` itself and is independent of sequence length: the same value yields
+the same per-tile test at any resolution or frame count. `threshold=0` skips no tiles. Because the
+left-hand side of the skip test lies in `(0, 1]`, values in `(0, 1)` are the meaningful range; the
+schema accepts any finite non-negative value.
+
+`target_sparsity` selects a point on a curve fitted per model by
+[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity),
+so that `s` lands near that fraction of skipped tiles on the calibration data. The coefficients
+`a` and `b` come from the checkpoint; the achieved sparsity on another prompt, shape, or layer can
+differ.
+
+### Skip-Softmax calibration config
+
+ModelOpt stores the calibration in each transformer component's `config.json` under
+`sparse_attention_config`. The following is taken from the `transformer/config.json` of
+[`nvidia/Wan2.2-T2V-A14B-Diffusers-FP8`](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8),
+with the `ignore` list abridged:
+
+```json
+{
+  "sparse_attention_config": {
+    "config_groups": {
+      "group_0": {
+        "algorithm": "skip_softmax",
+        "targets": ["WanAttention"],
+        "ignore": ["blocks.0.attn1", "blocks.0.attn2", "blocks.1.attn1", "..."],
+        "threshold_scale_factor": {
+          "formula": "a * exp(b * target_sparsity)",
+          "coefficients": {"a": 2822.6209952221557, "b": 4.388250623520793}
+        },
+        "target_sparsity": 0.7,
+        "disabled_until_timestep": 0.93
+      }
+    }
+  }
+}
+```
+
+A checkpoint with several transformer components carries one such config per component, each
+with its own coefficients; in this example, `transformer_2/config.json` holds the low-noise
+expert's curve.
+
+vLLM-Omni consumes `formula`, `coefficients`, and `ignore`, and ignores `targets`,
+`target_sparsity`, and `disabled_until_timestep`:
+
+- Only the formula `a * exp(b * target_sparsity)` is accepted; a checkpoint with any other formula
+  is rejected at startup. If a spec requests `target_sparsity` and the checkpoint carries no
+  calibration, startup fails and the error points to `threshold` as the calibration-free
+  alternative.
+- Each component's curve is applied to that component's attention layers. If the config for a
+  secondary component such as `transformer_2` cannot be read, that component stays dense and a
+  warning is logged.
+- `ignore` holds fnmatch patterns matched against both the full module name and the name relative
+  to the component, so `blocks.0.attn1` matches `transformer.blocks.0.attn1` as well as
+  `transformer_2.blocks.0.attn1`. Matching layers receive no coefficients and therefore stay dense
+  regardless of the user configuration.
+
+At each attention call, a layer with a `threshold` uses it directly; a layer with
+`target_sparsity` and stamped coefficients evaluates the curve; a layer with `target_sparsity` but
+no coefficients runs dense. The timestep gate below is applied to the result before it reaches the
+kernel.

 ## Timestep gating

-`disabled_until_timestep = D` keeps the mode off during the early, high-noise denoise steps and
-turns it on once the normalized timestep `t` drops to `t ≤ D` (`t` runs `1.0` → `0.0` over the
-schedule). The early steps set the global structure of the output and their errors propagate through
-every later step, so keeping them dense costs a few skipped-tile opportunities but protects fidelity.
+The early, high-noise denoising steps set the global structure of the output, and their errors
+propagate through every later step. Keeping just these steps dense markedly improves the quality
+of the generated video at a small cost in skipped work. `disabled_until_timestep` implements this:
+on each attention call the backend compares the current normalized timestep `t` against the
+configured cutoff and passes no skip factor to the kernel while `t` is above it. Skip-Softmax
+becomes active on the first step whose `t` is at or below the cutoff and stays active for the rest
+of the run, since `t` only decreases.

-`t` is the scheduler's own timestep divided by `num_train_timesteps`, published by the pipeline via
-`DenoiseProgressMixin.record_denoise_step`. It is deliberately not derived from the step index,
-because schedulers space their steps non-uniformly. A pipeline that does not publish a timestep
-stays dense when `disabled_until_timestep` is set.
+The default `disabled_until_timestep = 0` is a sentinel rather than a cutoff: the gate is off, the
+timestep is never read, and Skip-Softmax runs on every step. Any positive value enables the gate.
+`disabled_until_timestep = 1.0` therefore also yields no dense steps on a pipeline that publishes
+`t`, but on a pipeline that does not publish it the backend stays dense and logs a warning once
+rather than guessing from the step index.
+
+### Where `t` comes from
+
+The pipeline publishes `t` once per denoising step before running the transformer. It is the
+scheduler's own position in `[0, 1]`, decreasing from near `1.0` (pure noise) to `0.0` (clean
+sample): scheduler-based pipelines publish the scheduler timestep divided by `num_train_timesteps`,
+and rectified-flow pipelines publish the current sigma. `t` is deliberately not the step index
+divided by the step count. Schedulers place their steps non-uniformly in `t`, so expressing the
+cutoff in `t` ties it to the same noise level across step counts and schedulers.
+
+### Mapping the cutoff to denoising steps
+
+For a run whose published sequence is `t[0], ..., t[N-1]`, the number of dense steps is
+
+```text
+dense_steps = count(t[i] > disabled_until_timestep)
+skip_softmax_steps = N - dense_steps
+```
+
+This count **depends on the schedule**, and schedules differ from model to model: the
+scheduler family, the number of steps, and any time shift all change where the steps land in `t`.
+The same cutoff can therefore gate a very different fraction of the run on two models. Take
+MiniMax-H3 as an example. Its video branch uses a flow-shifted rectified-flow schedule that applies
+a shift `s` to `N + 1` uniform positions `u` from `1` to `0`,
+
+```text
+t = s * u / (1 + (s - 1) * u)
+```
+
+and a large `s` pushes most positions toward `t ≈ 1`. With the default `s = 12` and
+`num_inference_steps = 50` (50 sigma points, `N = 49` denoiser forwards), the published sequence
+starts `1.000, 0.998, 0.996, 0.995, 0.993, ...` and stays above `0.9` for 28 steps, giving:
+
+| `disabled_until_timestep` | Dense steps | Skip-Softmax steps |
+| :---: | ---: | ---: |
+| `1.00` | 0 | 49 |
+| `0.99` | 6 | 43 |
+| `0.97` | 14 | 35 |
+| `0.95` | 19 | 30 |
+| `0.90` | 28 | 21 |
+| `0.86` | 33 | 16 |
+
+These numbers hold only for that schedule. A model with a smaller shift, a different step count,
+or a scheduler that is not rectified flow produces a different sequence and a different table; a
+request that overrides `flow_shift` or the step count changes it for MiniMax-H3 too.
+Choose the cutoff by counting against the schedule actually served, not by reusing a value from
+another model.
@@ -95,7 +95,7 @@ response.stream_to_file("output.wav")

 ### Endpoint

-```
+```text
 POST /v1/audio/speech
 Content-Type: application/json
 ```
@@ -105,7 +105,7 @@ Content-Type: application/json
 #### OpenAI Standard Parameters

 | Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
+| ----------- | ------ | --------- | ------------- |
 | `input` | string | **required** | The text to synthesize into speech |
 | `model` | string | server's model | Model to use (optional, should match server if specified) |
 | `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
@@ -115,9 +115,9 @@ Content-Type: application/json
 #### vLLM-Omni Extension Parameters

 | Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
+| ----------- | ------ | --------- | ------------- |
+| `task_type` | string | null (inferred) | TTS task type: CustomVoice, VoiceDesign, or Base. For Qwen3-TTS, `ref_audio` or `ref_text` infers Base when this field is omitted, while an uploaded/precomputed voice always selects Base. Other omitted values select CustomVoice; VoiceDesign must be specified explicitly. |
 | `sample_rate` | integer | model native | Target output sample rate. Qwen3-TTS supports 8000 or 24000 Hz; the model remains native 24 kHz internally and is resampled before encoding. |
-| `task_type` | string | "CustomVoice" | TTS task type: CustomVoice, VoiceDesign, or Base |
 | `language` | string | "Auto" | Language (see supported languages below) |
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
@@ -131,7 +131,7 @@ Content-Type: application/json
 #### Voice Clone Parameters (Base task)

 | Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
+| ----------- | ------ | --------- | ------------- |
 | `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`). Local files fold `mtime_ns` and `size` into cache keys to automatically reload on-disk edits; HTTP URLs and base64 URIs remain cached by string locator. |
 | `ref_text` | string | null | Transcript of reference audio |
 | `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
@@ -205,7 +205,7 @@ The `usage` object on `speech.audio.done` is the same shape returned per item by

 ### Voices Endpoint

-```
+```text
 GET /v1/audio/voices
 ```

@@ -230,7 +230,7 @@ Lists available voices for the loaded model.

 `uploaded_voices` is always present (empty list when no custom voices have been uploaded). Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.

-```
+```text
 POST /v1/audio/voices
 Content-Type: multipart/form-data
 ```
@@ -240,7 +240,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 **Form Parameters:**

 | Parameter | Type | Required | Description |
-|-----------|------|----------|-------------|
+| ----------- | ------ | ---------- | ------------- |
 | `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
 | `consent` | string | Yes | Consent recording ID |
 | `name` | string | Yes | Name for the new voice |
@@ -289,7 +289,7 @@ The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and
 Client -> Server:

 | Message | Description |
-|---------|-------------|
+| --------- | ------------- |
 | `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) |
 | `{"type": "input.text", "text": "..."}` | Text chunk |
 | `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open |
@@ -298,7 +298,7 @@ Client -> Server:
 Server -> Client:

 | Message | Description |
-|---------|-------------|
+| --------- | ------------- |
 | `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
 | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
 | `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
@@ -312,17 +312,17 @@ text, emits `session.done`, and then waits on the same connection for the next
 utterance, so a client that speaks repeatedly (for example one driven by an
 upstream LLM) pays the WebSocket handshake once instead of once per utterance.

-* The session config is sticky. Send `input.text` again straight after
+- The session config is sticky. Send `input.text` again straight after
  `session.done` to reuse it, or send another `session.config` first to change
  voice, format, or reference audio. A `session.config` sent while text is
  still buffered is rejected so no pending input is silently dropped.
-* An utterance is the flush unit, not a linguistic one: it is whatever text was
+- An utterance is the flush unit, not a linguistic one: it is whatever text was
  buffered when `input.done` arrived, of any length, synthesized as one request.
  `utterance_index` counts those flushes across the connection, so it tells you
  which `input.done` a frame belongs to. `sentence_index` counts within one
  flush and so pairs with `total_sentences`, which means every utterance reports
  `sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer).
-* End the connection with `session.close`, or by closing the socket. An idle
+- End the connection with `session.close`, or by closing the socket. An idle
  connection is still closed after the server's idle timeout, which now also
  applies to the gap between utterances.

@@ -331,10 +331,9 @@ upstream LLM) pays the WebSocket handshake once instead of once per utterance.
 All REST API parameters are supported, plus:

 | Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
+| ----------- | ------ | --------- | ------------- |
 | `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |

-
 ```bash
 DELETE /v1/audio/voices/{name}
 ```
@@ -344,7 +343,7 @@ Delete an uploaded voice sample.
 **Path Parameters:**

 | Parameter | Type | Required | Description |
-|-----------|------|----------|-------------|
+| ----------- | ------ | ---------- | ------------- |
 | `name` | string | Yes | Name of the voice to delete |

 **Response Example:**
@@ -434,6 +433,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 ### Upload Voice

 Upload voice (speaker embedding only):
+
 ```bash
 curl -X POST http://localhost:8091/v1/audio/voices \
  -F "audio_sample=@/path/to/voice_sample.wav" \
@@ -442,6 +442,7 @@ curl -X POST http://localhost:8091/v1/audio/voices \
 ```

 Upload voice with transcript (in-context cloning, higher quality):
+
 ```bash
 curl -X POST http://localhost:8091/v1/audio/voices \
  -F "audio_sample=@/path/to/voice_sample.wav" \
@@ -451,6 +452,7 @@ curl -X POST http://localhost:8091/v1/audio/voices \
 ```

 ### Use Uploaded Voice
+
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
@@ -460,6 +462,11 @@ curl -X POST http://localhost:8091/v1/audio/speech \
    }' --output cloned.wav
 ```

+For Qwen3-TTS, an uploaded voice is a Base voice-cloning input. The server
+infers `task_type="Base"` when `voice` names an uploaded entry, so the request
+must be sent to a Base checkpoint. Built-in presets such as `vivian` and
+`ryan` remain CustomVoice speakers and require a CustomVoice checkpoint.
+
 ### Voice Storage & Caching

 Uploaded voices are persisted to disk as a single `.safetensors` file per voice
@@ -479,6 +486,9 @@ once.
 ### Precomputed Custom Voices

 Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
+Qwen3-TTS precomputed voices follow the same Base task and checkpoint-matching
+rules as uploaded voices.
+
 Generate a directory containing `custom_voice_manifest.json` plus one
 `.safetensors` file per voice, then set the pipeline-wide deploy config field:

@@ -517,7 +527,7 @@ by `GET /v1/audio/voices`. Valid precomputed voices can be used in
 **Configuration (environment variables):**

 | Variable | Default | Description |
-|----------|---------|-------------|
+| ---------- | --------- | ------------- |
 | `SPEAKER_SAMPLES_DIR` | `~/.cache/vllm-omni/speakers` | Directory for persisted uploaded speakers (`.safetensors` files). |
 | `SPEAKER_MAX_UPLOADED` | `1000` | Maximum number of uploaded speakers kept on disk. Upload requests past the cap return 400. |

@@ -529,7 +539,7 @@ The batch endpoint synthesizes multiple texts in a single request, returning all

 ### Endpoint

-```
+```text
 POST /v1/audio/speech/batch
 Content-Type: application/json
 ```
@@ -537,7 +547,7 @@ Content-Type: application/json
 ### Request Parameters

 | Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
+| ----------- | ------ | --------- | ------------- |
 | `items` | array | **required** | List of items to synthesize (1–32) |
 | `model` | string | server's model | Model to use |
 | `voice` | string | null | Default voice for all items |
@@ -554,7 +564,7 @@ Content-Type: application/json
 Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:

 | Field | Type | Description |
-|-------|------|-------------|
+| ------- | ------ | ------------- |
 | `input` | string | **required** — text to synthesize |
 | `voice` | string | Override voice for this item |
 | `response_format` | string | Override format for this item |
@@ -686,7 +696,7 @@ for result in response.json()["results"]:
 ### Configuration

 | Parameter | Source | Default | Description |
-|-----------|--------|---------|-------------|
+| ----------- | -------- | --------- | ------------- |
 | `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request |

 All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_num_seqs` and queues the rest — no client-side throttling needed.
@@ -709,17 +719,17 @@ The bundled config also sets `initial_codec_chunk_frames: 1`. This emits only th
 ### Qwen3-TTS

 | Model | Task Type | Description |
-|-------|-----------|-------------|
+| ------- | ----------- | ------------- |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
-| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio or an uploaded/precomputed voice |
 | `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
-| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster Base variant for voice cloning, including uploaded/precomputed voices |

 ### Fish Speech S2 Pro

 | Model | Description |
-|-------|-------------|
+| ------- | ------------- |
 | `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. |

 Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech section of the online TTS hub](../user_guide/examples/online_serving/text_to_speech.md#fish-speech-s2-pro) for details.
@@ -727,31 +737,31 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
 ### Voxtral TTS

 | Model | Description |
-|-------|-------------|
+| ------- | ------------- |
 | `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |

 ### CosyVoice3

 | Model | Description |
-|-------|-------------|
+| ------- | ------------- |
 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | Voice cloning from `ref_audio` + `ref_text`. No built-in voice presets — upload a voice or pass `ref_audio`/`ref_text` per request. |

 ### OmniVoice

 | Model | Description |
-|-------|-------------|
+| ------- | ------------- |
 | `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. |

 ### VoxCPM2

 | Model | Description |
-|-------|-------------|
+| ------- | ------------- |
 | `openbmb/VoxCPM2` | TTS + voice cloning with built-in speaker presets and uploaded-voice support. Accepts `voice` (preset or uploaded) or `ref_audio` + optional `ref_text`. |

 ### MOSS-TTS-Nano

 | Model | Description |
-|-------|-------------|
+| ------- | ------------- |
 | `OpenMOSS-Team/MOSS-TTS-Nano` | Voice cloning only. Requires `ref_audio` (or an uploaded `voice`); no built-in voice presets. `ref_text` is accepted but ignored — upstream's `voice_clone` mode does not consume a transcript. |

 ## Error Responses
@@ -771,6 +781,19 @@ Invalid parameters:
 }
 ```

+Qwen3-TTS task/checkpoint mismatch:
+
+```json
+{
+    "error": {
+        "message": "Qwen3-TTS CustomVoice checkpoint does not support task_type='Base'. Use task_type='CustomVoice' or load the matching Base checkpoint.",
+        "type": "BadRequestError",
+        "param": null,
+        "code": 400
+    }
+}
+```
+
 ### 404 Not Found

 Model not found:
@@ -791,10 +814,15 @@ Model not found:
 ### "TTS model did not produce audio output"

 Ensure you're using the correct model variant for your task type:
+
 - CustomVoice task → CustomVoice model
 - VoiceDesign task → VoiceDesign model
 - Base task → Base model

+Uploaded and precomputed Qwen3-TTS voices select the Base task, even when the
+client supplies a different `task_type`; serve them with a Base checkpoint.
+Built-in Qwen3-TTS presets remain CustomVoice voices.
+
 ### Server Not Running

 ```bash
@@ -805,6 +833,7 @@ curl http://localhost:8091/v1/audio/voices
 ### Out of Memory

 If you encounter OOM errors:
+
 1. Use smaller model variant: `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`
 2. Reduce `--gpu-memory-utilization`

@@ -823,7 +852,7 @@ final-output drain to a condition-variable wakeup at the same time.
 **Configuration (environment variables):**

 | Variable | Default | Description |
-|----------|---------|-------------|
+| --- | --- | --- |
 | `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |

 Set it on the process that runs the orchestrator (stage 0 of an omni
@@ -20,7 +20,7 @@ For the internal selector, registry, and platform contract, see
 | 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) |
-| Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
+| Use FastVideo VSA with FastWan2.2-TI2V-5B or FastH3 MiniMax-H3 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |

 ## Backend options

@@ -36,7 +36,7 @@ For the internal selector, registry, and platform contract, see
 | `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) |
-| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
+| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B and FastH3 MiniMax-H3 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |

 ## Configuration

@@ -1,8 +1,9 @@
 # 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.
+Use these backends to run dense attention. Optional quantized or sparse modes
+remain inactive unless they are explicitly configured. 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).
@@ -20,10 +21,11 @@ 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.
+is FA4 only (`flash_attn.cute` / `vllm-omni[fa4]`). If FA4 is unavailable,
+explicit `FLASH_ATTN` is rejected and automatic selection continues to another
+compatible Blackwell backend. Hopper-only FA2/FA3 wheels are not used, even if
+they import. 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
@@ -38,8 +40,25 @@ 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.
+Blackwell. On Blackwell, `FLASH_ATTN` is FA4 only: Hopper-only FA2/FA3
+wheels are not used, even if they import. If FA4 is missing, explicit
+`FLASH_ATTN` is rejected and automatic selection continues to other
+Blackwell backends.
+
+## `TRTLLM_ATTN`
+
+`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels and is the platform
+default on datacenter Blackwell for models that declare a compatible attention
+path. Selected without a `quant` or `skip_softmax` block, it runs dense BF16
+attention at FA4-level performance.
+
+```bash
+vllm serve <model> --omni \
+  --diffusion-attention-backend TRTLLM_ATTN
+```
+
+See [TRTLLM Attention](trtllm.md) for its requirements and for the optional
+SAGE quantization and Skip-Softmax modes.

 ## `CUDNN_ATTN`

@@ -73,6 +92,13 @@ DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
 option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
 is installed but cuDNN is too old for `CUDNN_ATTN`.

+On Blackwell, `auto` resolves to FlashInfer cute-dsl, which cannot run a
+nontrivial custom mask. Automatic selection may fall back to SDPA for those
+masks. An explicit `FLASHINFER_ATTN` selection does not: sequence-parallel
+auto-padding that would create a padding mask is rejected at capability
+preflight. Use `TORCH_SDPA`, pin `quant.flashinfer_backend` to `fa2`/`fa3`
+where those kernels exist, or choose a mask-capable backend.
+
 ```bash
 vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN
 ```
@@ -7,11 +7,18 @@ attention only against the selected top-k blocks.

 The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.

+MiniMax-H3 served with a FastH3 VSA adapter uses a second, H3-specific route
+through the same backend. See the
+[MiniMax-H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for its geometry,
+supported topologies, and commands.
+
 VSA is a CUDA-only, explicitly selected backend. It requires the
 `fastvideo-kernel` package and currently supports non-causal self-attention
 with equal query and key/value sequence lengths. Unsupported shapes, masks,
 dtypes, sequence-parallel execution, or kernel failures fall back to
-`TORCH_SDPA` and emit a warning with the reason.
+`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
+sequence-parallel context is one of those fallbacks; the H3 route supports pure
+Ulysses and rejects ring or all-gather sequence parallelism at startup.

 ## Enable the backend

@@ -119,7 +126,8 @@ the backend guarantees sparse execution:
 - `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
  executed; the warning includes the reason.

-The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks,
+The Wan route requires CUDA tensors in FP16 or BF16, 256-token blocks,
 standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
-and no active sequence-parallel context. NPU and XPU paths do not execute the
-FastVideo VSA CUDA kernel.
+and no active sequence-parallel context. The MiniMax-H3 route uses 64-token
+`(4, 4, 4)` blocks and supports pure Ulysses. NPU and XPU paths do not execute
+the FastVideo VSA CUDA kernel.
@@ -15,9 +15,9 @@ implementations to reduce numerical drift and sampling divergence.
 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`.
+If `kernels` is unavailable, an explicit `FLASH_ATTN_HUB` or
+`FLASH_ATTN_3_HUB` selection raises. Automatic / `FLASH_ATTN` selection in the
+Diffusers adapter may still walk hub then local FlashAttention backends.

 ## Usage

@@ -1,103 +1,274 @@
 # TRTLLM Attention

-`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
-BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
+`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels on datacenter
+Blackwell GPUs. Selected on its own it computes dense BF16 attention, and its
+dense performance is on par with
+[FlashAttention-4](dense_backends.md#flashattention-4-on-blackwell). It also
+provides two opt-in, lossy acceleration modes that can be enabled independently
+or together:

-For common selection and per-role configuration, see the
-[attention backend overview](../attention_backends.md).
+| Mode | Config key | What it changes |
+| --- | --- | --- |
+| [Skip-Softmax](#skip-softmax) | `skip_softmax` | Skips the Softmax and `PV` work of KV tiles whose scores are too low to matter |
+| [SAGE quantization](#sage-quantization) | `quant` | Runs `QK^T` in INT8 or FP8 and `PV` in FP8 instead of BF16 |
+
+Both keys are set through `--diffusion-attention-config`, as JSON or as
+vLLM-style dotted flags; the
+[attention backend overview](../attention_backends.md#configuration) covers
+both syntaxes and per-role resolution. The examples on this page use JSON.
+
+Attention computes scores `S = QK^T`, probabilities `P = softmax(S)`, and the
+output `O = PV`. SAGE lowers the precision of both matrix multiplications.
+Skip-Softmax keeps `QK^T` dense and removes the Softmax and `PV` work for
+tiles that `QK^T` shows to be unimportant. The two modes therefore compose:
+SAGE makes every tile cheaper, and Skip-Softmax reduces the number of tiles
+that reach the second half of the kernel.

 ## Requirements

-Selecting `TRTLLM_ATTN` requires:
+`TRTLLM_ATTN` requires all of the following:

- datacenter Blackwell, sm_100 or sm_103;
+- an `sm100a` or `sm103a` GPU (B200, B300, GB200, GB300); workstation Blackwell
+  (`sm120`/`sm121`) is not supported;
 - `head_dim=128`;
- FlashInfer; and
- a model-declared packed or mask-free compatible path.
+- a FlashInfer build that exposes the trtllm-gen kernels (0.6.16rc1 or newer
+  for SAGE);
+- an attention path that is mask-free or provides packed-padding metadata.
+  Structural suffix padding is expressed through that metadata rather than an
+  `attn_mask` tensor.

-It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
-explicit selection raises instead of silently changing the backend.
+An explicit selection that violates these requirements raises at startup
+instead of silently falling back to another backend.

-## Dense mode
+### Sequence parallelism

-Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
+`TRTLLM_ATTN` runs under no sequence parallelism or under pure Ulysses.
+Ulysses redistributes the sequence and attention heads around the attention
+call, but the local computation still goes through the configured backend, so
+both optional modes work unchanged. Ring and AllGather-KV do not:
+
+- Ring runs its own distributed attention and bypasses the backend. Combining
+  Ring with a `skip_softmax` key raises; a `quant` key would be silently
+  ignored, so do not combine them either.
+- AllGather-KV changes the Q/KV distribution and is rejected when
+  `TRTLLM_ATTN` is selected.
+
+Configure eight-way Ulysses alone as:

 ```bash
-vllm-omni serve <model> --diffusion-attention-backend TRTLLM_ATTN
+--usp 8 --ring 1 --allgather-degree 1
+```
+
+A degree of `1` disables that sequence-parallel mode. It does not limit the
+server to one GPU or affect tensor, pipeline, or VAE parallelism.
+
+## Quick start
+
+On datacenter Blackwell the platform selects `TRTLLM_ATTN` by default when the
+model declares a compatible path. To select it explicitly, for dense BF16:
+
+```bash
+vllm serve <model> --omni \
+  --diffusion-attention-backend TRTLLM_ATTN
+```
+
+To enable both optimizations with a conservative Skip-Softmax setting and FP8
+SAGE:
+
+```bash
+vllm serve <model> --omni \
+  --diffusion-attention-config '{
+    "default": {
+      "backend": "TRTLLM_ATTN",
+      "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16},
+      "skip_softmax": {"threshold": 0.05, "disabled_until_timestep": 0.97}
+    }
+  }'
+```
+
+The startup log reports the selection; look for one of:
+
+```text
+Defaulting to diffusion attention backend TRTLLM_ATTN (datacenter Blackwell ..., head_dim 128)
+Resolved diffusion attention backend 'TRTLLM_ATTN' for role='self' via attention_config.default
 ```

 ## 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.
+Skip-Softmax (BLASST) skips the Softmax and `PV` work of KV tiles whose
+scores fall far below the running row maximum. `QK^T` always runs, and how
+many tiles qualify depends on the input. The
+[feature design](../../../design/feature/skip_softmax.md) gives the skip test.

-| Key | Valid values | Meaning |
+### Configuration keys
+
+| Key | Range | 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` |
+| `threshold` | `>= 0`; useful values in `(0, 1)` | Skip threshold, independent of sequence length. Calibration-free. |
+| `target_sparsity` | `[0, 1]` | Requested operating point on the checkpoint's calibrated curve. Requires calibration metadata. |
+| `disabled_until_timestep` | `[0, 1]`; default `0` | Keeps attention dense while the normalized timestep `t > D`; `0` disables the gate. |
+
+`threshold` and `target_sparsity` are two ways to set the same kernel
+threshold; setting both is a configuration error. Exactly one of them enables
+Skip-Softmax.
+
+### Direct threshold
+
+Set `threshold` when the checkpoint carries no calibration. `threshold=0`
+skips nothing; larger values skip more tiles and lower output fidelity. Values
+around `0.05` are a reasonable first try for video DiTs; tune against dense
+output on the same prompt and seed. The value is independent of sequence
+length; see the
+[feature design](../../../design/feature/skip_softmax.md#from-configuration-to-the-kernel-threshold)
+for how it reaches the kernel.

 ```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}}}'
+vllm serve <model> --omni \
+  --diffusion-attention-config \
+  '{"default":{"backend":"TRTLLM_ATTN","skip_softmax":{
+    "threshold":0.05,"disabled_until_timestep":0.97}}}'
 ```

-```python
-from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
+### Calibrated target sparsity

-config = AttentionConfig(
-    default=AttentionSpec(
-        backend="TRTLLM_ATTN",
-        skip_softmax=SkipSoftmaxSpec(
-            target_sparsity=0.65,
-            disabled_until_timestep=0.86,
-        ),
-    ),
-)
+A fixed `threshold` does not produce a fixed fraction of skipped tiles because
+the score distribution changes with the model, prompt, and shape.
+[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer)
+can calibrate a per-model curve from sparsity to threshold and store it in the
+checkpoint's transformer `config.json` under `sparse_attention_config`.
+`target_sparsity` selects a point on that curve. The achieved sparsity still
+varies per prompt, layer, and denoising step; the calibration makes the
+requested value a meaningful target, not a guarantee.
+
+What vLLM-Omni takes from the checkpoint:
+
+- The curve coefficients. Only the `a * exp(b * target_sparsity)` form ModelOpt
+  writes is supported; another formula is rejected at startup.
+- The `ignore` list: attention layers matching those patterns stay dense.
+- For multi-expert Diffusers checkpoints, `transformer_2/config.json` is read
+  separately; if it is missing, `transformer_2` stays dense with a warning.
+
+Checkpoint-level `target_sparsity` and `disabled_until_timestep` defaults are
+not consumed; set them in the vLLM-Omni configuration. The
+[feature design](../../../design/feature/skip_softmax.md#skip-softmax-calibration-config)
+documents the checkpoint format and how the coefficients reach the kernel.
+
+To calibrate a checkpoint yourself, follow the
+[ModelOpt Skip-Softmax example](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity),
+which exports a copy of the checkpoint with the coefficients embedded. The
+fitted curve depends on the attention statistics of the calibration shape, so
+calibrate at the resolution and frame count you will serve.
+
+Requesting `target_sparsity` for a checkpoint without calibration is a startup
+error that names the `threshold` alternative. The
+[ModelOpt Wan2.2 FP8 checkpoint](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8/blob/main/transformer/config.json)
+is a calibrated example:
+
+```bash
+vllm serve nvidia/Wan2.2-T2V-A14B-Diffusers-FP8 --omni \
+  --diffusion-attention-config \
+  '{"default":{"backend":"TRTLLM_ATTN","skip_softmax":{
+    "target_sparsity":0.75,"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.
+### Timestep gating
+
+The early, high-noise denoising steps fix the global layout of the output, and
+their errors propagate through every later step. `disabled_until_timestep`
+keeps those steps dense and enables Skip-Softmax once the normalized timestep
+`t` reaches the configured cutoff. The default `0` is a sentinel that turns the
+gate off: Skip-Softmax runs on every step and no timestep needs to be
+published. `1.0` also leaves no step dense, but goes through the gate and
+therefore requires the pipeline to publish `t`.
+
+`t` is the scheduler's own timestep normalized to `[0, 1]`, published by the
+pipeline for each denoising step; for rectified-flow models it is the current
+sigma. It runs from near `1.0` down to `0.0`, but not linearly in the step
+index: flow-shifted schedules spend many of their steps at high `t`. The number
+of dense steps a given cutoff produces therefore depends on the model's
+schedule and step count, and follows from the actual timestep sequence:
+
+```text
+dense_steps = count(t[i] > disabled_until_timestep)
+```
+
+Derive the cutoff from the schedule of the model you are serving rather than
+reusing another model's value; the
+[feature design](../../../design/feature/skip_softmax.md#mapping-the-cutoff-to-denoising-steps)
+works through an example.
+
+A pipeline that does not publish `t` stays dense whenever
+`disabled_until_timestep > 0` is set, and logs a warning once. Pipelines
+publish it through `DenoiseProgressMixin.record_denoise_step`.

 ## SAGE quantization

-TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
-matrix multiplications use low precision.
+SAGE quantization runs both attention matrix multiplications in low precision:
+Q and K are quantized to INT8 or FP8 E4M3 for `QK^T`, and P and V use FP8 E4M3
+for `PV`. P is quantized inside the FMHA kernel; V is quantized per channel
+before the kernel call.
+vLLM-Omni exposes the Q/K dtype and the Q/K scale granularity. The P and V
+formats are fixed by the kernel, so the `quant` key has no V dtype for this
+backend. This mode is distinct from the standalone
+[SageAttention backends](sage.md), which use their own kernels.

-| Key | Valid values | Meaning |
+FP8 Q/K kernels exist on `sm100a` and `sm103a`; INT8 Q/K kernels exist on
+`sm100a` only.
+
+| Key | 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` |
+| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype. Setting it enables SAGE. |
+| `q_block_size` | `1`, `4`, `16` | Consecutive query tokens sharing one Q scale; default `1` |
+| `k_block_size` | `1`, `4`, `16` | Consecutive key tokens sharing one K scale; default `16` |

-V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
-Only block sizes with compiled kernels are accepted.
+Start with the defaults, `q_block_size=1` and `k_block_size=16`, and try
+smaller K blocks only if output quality does not meet expectations; smaller
+blocks give finer scales at some cost in speed. When a KV sequence in a call is
+shorter than `k_block_size`, that call falls back to dense attention and a
+warning is logged once.

 ```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}}}'
+vllm serve <model> --omni \
+  --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
+The `quant` key is shared with `FLASHINFER_ATTN`, but each backend validates
+its own fields: `float16`/`bfloat16` Q/K dtypes and `dtype_vo` are
+`FLASHINFER_ATTN` options and are rejected here.

-config = AttentionConfig(
-    default=AttentionSpec(
-        backend="TRTLLM_ATTN",
-        quant=AttnQuantSpec(
-            dtype_qk="fp8_e4m3",
-            q_block_size=1,
-            k_block_size=16,
-        ),
-    ),
-)
+## Composing both modes
+
+`skip_softmax` and `quant` may appear in the same `AttentionSpec`. Their
+quality effects compound, so establish a dense baseline, enable one mode at a
+time, and then evaluate the combination on the same prompts and seeds.
+
+Modes configured in `default` apply to every attention role that has no
+`per_role` entry. A `per_role` spec replaces the whole spec for that role and
+does not inherit `quant` or `skip_softmax` from `default`, so
+`{"backend":"TRTLLM_ATTN"}` is the way to keep a short or sensitive attention
+site dense while the long DiT sequence uses both modes:
+
+```bash
+vllm serve MiniMaxAI/MiniMax-H3 --omni \
+  --diffusion-attention-config '{
+    "default": {
+      "backend": "TRTLLM_ATTN",
+      "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16},
+      "skip_softmax": {"threshold": 0.05, "disabled_until_timestep": 0.97}
+    },
+    "per_role": {
+      "minimax_h3.token_refiner": {"backend": "TRTLLM_ATTN"}
+    }
+  }'
 ```

-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.
+Role names are declared by each model; the
+[attention backend overview](../attention_backends.md#configuration) covers
+the resolution order and the equivalent Python API.
+
+End-to-end speedup depends on the share of step time spent in attention, the
+sequence length, the chosen Q/K precision, and the tile sparsity the input
+actually yields. Benchmark the exact workload rather than extrapolating from
+another model.
@@ -94,7 +94,7 @@ In `DiffusionParallelConfig`


 !!! info
-    Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** is an exception: it supports `cfg_parallel_size=3`, which adds a third branch on rank 2 for full three-way CFG parallelism.
+    Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** and **Boogu-Image-Edit double guidance** also support `cfg_parallel_size=3` for full three-way branch execution. Boogu double guidance additionally supports size 2 via round-robin assignment of its three branches.


 ---
@@ -117,6 +117,7 @@ The following tables show which models support each feature:
 | Model                    | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Bagel**                |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Boogu-Image**          |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **FLUX.1-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.2-klein**         |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
@@ -162,7 +163,7 @@ The following tables show which models support each feature:
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2.5**                  |     ❌     | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) |         ❓         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |    ❓ (FP8)     |        ❌         |
-| **SANA-Video-2B T2V I2V**   |     ❓     |     ❓      |           ❓           |       ❓        |         ❓         |         ❓         |   ❓    |             ❓             |          ❓           |       ❓        |        ❌         |
+|  **SANA-Video-2B T2V I2V**   |     ❌     |     ❌      |          ❌           |       ✅       |  ✅ (TP=2 only)   |         ❌         |   ❌   |            ❌             |          ❌          |       ❌       |        ❌        |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **Cosmos3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
@@ -2,7 +2,6 @@

 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/text_to_speech>.

-
 vLLM-Omni exposes TTS models through the OpenAI-compatible
 [`POST /v1/audio/speech`](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md) endpoint,
 launched with `vllm serve <model> --omni`. Each TTS model has its own
@@ -16,7 +15,7 @@ For the full list of supported architectures across all modalities, see
 ## Supported Models

 | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
-|---|---|---|---|---|---|
+| --- | --- | --- | --- | --- | --- |
 | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
 | higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
@@ -104,6 +103,7 @@ For full request-shape documentation (all parameters, response formats, error co
 2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets.

 ### Prerequisites
+
 ```bash
 huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
 ```
@@ -111,6 +111,7 @@ huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
 If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class).

 ### Launch
+
 ```bash
 vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
 # or:
@@ -120,13 +121,16 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
 Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.

 ### CLI client
+
 The client defaults to the official upstream zero-shot prompt, so it runs without extra flags:
+
 ```bash
 cd examples/online_serving/text_to_speech/cosyvoice3
 python speech_client.py --text "收到好友从远方寄来的生日礼物。"
 ```

 Pass your own reference clip and transcript for a different voice:
+
 ```bash
 python speech_client.py --text "Hello, this is a cloned voice." \
    --ref-audio /path/to/reference.wav \
@@ -134,6 +138,7 @@ python speech_client.py --text "Hello, this is a cloned voice." \
 ```

 Stream PCM instead of WAV:
+
 ```bash
 python speech_client.py --text "Hello world" --stream --output output.pcm
 ```
@@ -141,6 +146,7 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
 The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`.

 ### Notes
+
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
 - Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
 - For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
@@ -152,19 +158,23 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.

 ### Prerequisites
+
 ```bash
 pip install fish-speech
 ```

 ### Launch
+
 ```bash
 vllm serve fishaudio/s2-pro --omni --port 8091
 # or:
 ./fish_speech/run_server.sh
 ```
+
 The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`).

 ### Voice cloning
+
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
@@ -177,6 +187,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 ```

 ### CLI client
+
 ```bash
 cd examples/online_serving/text_to_speech/fish_speech
 python speech_client.py --text "Hello, how are you?"
@@ -184,12 +195,14 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
 ```

 ### Gradio demo
+
 ```bash
 ./fish_speech/run_gradio_demo.sh             # launches server + Gradio
 python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server already running
 ```

 ### Notes
+
 - Output: 44.1 kHz mono.
 - Streaming PCM player command must use `-r 44100`.

@@ -208,12 +221,15 @@ pip install -U "transformers>=5.3.0"
 ```

 ### Launch
+
 ```bash
 GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
 ```
+
 Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.

 ### Sending requests
+
 ```bash
 # Plain TTS
 python higgs_audio_v2/batch_speech_client.py \
@@ -231,14 +247,17 @@ python higgs_audio_v2/batch_speech_client.py \
 ```

 ### Notes
+
 - Output: 24 kHz mono.
 - `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
 - Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
+
 ## GLM-TTS

 2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.

 ### Launch
+
 ```bash
 vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
 # or:
@@ -246,6 +265,7 @@ bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-T
 ```

 ### Sending requests
+
 ```bash
 # Voice cloning (required)
 python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
@@ -262,11 +282,13 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
 ```

 ### Gradio demo
+
 ```bash
 bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
 ```

 ### Notes
+
 - Output: 24 kHz mono WAV via HiFT vocoder.
 - `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
 - Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
@@ -278,12 +300,15 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
 Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.

 ### Prerequisites
+
 ```bash
 huggingface-cli download k2-fsa/OmniVoice
 ```
+
 Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`.

 ### Launch
+
 ```bash
 vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
 # or:
@@ -291,6 +316,7 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
 ```

 ### CLI client
+
 ```bash
 cd examples/online_serving/text_to_speech/omnivoice
 python speech_client.py --text "Hello, how are you?"
@@ -300,6 +326,7 @@ python speech_client.py --text "Bonjour, comment allez-vous?" --language French
 The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`.

 ### Notes
+
 - Voice cloning and voice design require offline inference; see the [offline OmniVoice section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#omnivoice).

 ---
@@ -309,7 +336,7 @@ The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--l
 Three model variants exposed via separate checkpoints:

 | Variant | HF repo | Use |
-|---|---|---|
+| --- | --- | --- |
 | CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions |
 | VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description |
 | Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio |
@@ -317,6 +344,7 @@ Three model variants exposed via separate checkpoints:
 Each variant ships smaller `0.6B` companions where available.

 ### Launch
+
 ```bash
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
 # or:
@@ -326,27 +354,30 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
 ```

 ### Executor backend
+
 Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.

 To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.

 ### Tuning stage 1 `max_num_seqs` per task type
+
 The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost.

 CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types:

 ```bash
-vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
    --stage-overrides '{"1": {"max_num_seqs": 1}}'
 ```

 ### Sending requests
+
 ```bash
 # CustomVoice with a predefined speaker
 python qwen3_tts/openai_speech_client.py \
    --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
    --text "今天天气真好" \
-    --voice ryan \
+    --speaker ryan \
    --instructions "用开心的语气说"

 # VoiceDesign with a style description
@@ -366,7 +397,9 @@ python qwen3_tts/openai_speech_client.py \
 ```

 ### Voices endpoint
+
 List available voices, or upload a custom one for Base cloning:
+
 ```bash
 # List
 curl http://localhost:8091/v1/audio/voices
@@ -379,10 +412,18 @@ curl -X POST http://localhost:8091/v1/audio/voices \
    -F "ref_text=The exact transcript of the audio sample." \
    -F "speaker_description=warm narrator"
 ```
-Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
+
+For Qwen3-TTS, uploaded voices are Base voice-cloning inputs and require a Base
+checkpoint. When a request names an uploaded voice, the server infers
+`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
+CustomVoice speakers and require a CustomVoice checkpoint.

 ### Precomputed custom voices
-For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
+
+For reused Base voice-cloning speakers, precompute the reference artifacts once
+and load them at server startup. Precomputed voices use the same Base task and
+checkpoint-matching rules as uploaded voices:
+
 ```bash
 python qwen3_tts/precompute_custom_voice.py \
    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
@@ -392,11 +433,15 @@ python qwen3_tts/precompute_custom_voice.py \
    --mode icl \
    --output-dir /path/to/custom_voices
 ```
+
 `--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
+
 ```yaml
 custom_voice_dir: /path/to/custom_voices
 ```
+
 Then start the server with that config and call the Speech API with only the voice name:
+
 ```bash
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml

@@ -407,6 +452,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 ```

 ### Streaming PCM
+
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
@@ -419,16 +465,20 @@ curl -X POST http://localhost:8091/v1/audio/speech \
        "response_format": "pcm"
    }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
 ```
+
 Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.

 ### Streaming WebSocket
+
 The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
+
 ```bash
 python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
 python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
 ```

 `input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
+
 ```bash
 python qwen3_tts/streaming_speech_client.py \
    --text "First utterance." \
@@ -436,6 +486,7 @@ python qwen3_tts/streaming_speech_client.py \
 ```

 ### Gradio demos
+
 ```bash
 ./qwen3_tts/run_gradio_demo.sh                              # CustomVoice (default)
 ./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign
@@ -447,12 +498,15 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
 ```

 ### Speaker embedding interpolation
+
 `qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage.

 ### Batch client
+
 `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.

 ### Notes
+
 - Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
 - With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
@@ -467,12 +521,15 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
 Single-stage native AR TTS at 48 kHz.

 ### Launch
+
 ```bash
 vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000
 ```
+
 Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks.

 ### Sending requests
+
 ```bash
 # Zero-shot synthesis
 python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2."
@@ -482,10 +539,13 @@ python voxcpm2/openai_speech_client.py \
    --text "This should sound like the reference speaker." \
    --ref-audio /path/to/reference.wav
 ```
+
 The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.

 ### Precomputed custom voices
+
 For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
+
 ```bash
 python voxcpm2/precompute_custom_voice.py \
    --model openbmb/VoxCPM2 \
@@ -495,16 +555,21 @@ python voxcpm2/precompute_custom_voice.py \
    --prompt-text "Original transcript of the reference audio" \
    --output-dir /path/to/custom_voices
 ```
+
 Add the output directory to the deploy config:
+
 ```yaml
 custom_voice_dir: /path/to/custom_voices
 ```
+
 After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.

 ### Gradio demo (gapless streaming via AudioWorklet)
+
 ```bash
 python voxcpm2/gradio_demo.py
 ```
+
 Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`.

 ---
@@ -514,24 +579,31 @@ Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free p
 Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific.

 ### Prerequisites
+
 Latest `mistral_common` with `SpeechRequest` support:
+
 ```bash
 pip install -e /path/to/mistral-common  # or upgrade from PyPI when available
 ```

 ### Launch
+
 ```bash
 vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
 ```
+
 Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`.

 ### Gradio demo
+
 ```bash
 python voxtral_tts/gradio_demo.py
 ```
+
 The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients).

 ### Notes
+
 - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
 - Voice cloning is gated upstream and may require a recent `mistral_common`.
 - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
@@ -46,20 +46,20 @@ More failure modes (for example, network jitter and network interruption) will b
 ## Fault Injection Scenario Matrix

 | Scenario | Fault Type | System Behavior | Current Status |
-|------|----------------------------|----------|----------|
-| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx |  |
-| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx |  |
+| ------ | ---------------------------- | ---------- | ---------- |
+| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | |
+| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | |
 | No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725) <br>[#43060](https://github.com/vllm-project/vllm/issues/43060) |
-| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released |  |
-| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released |  |
-| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly |  |
-| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed |  |
+| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released | |
+| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released | |
+| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly | |
+| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed | |
 | Under load | Send `SIGKILL` to Worker | In-flight requests are hard interrupted (5xx/connection drop); main service becomes unavailable | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
-| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable |  |
+| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable | |
 | Under load | Send `SIGKILL` to serve main process | serve is hard-killed and current connections are interrupted; in-flight requests fail; after cleanup there are no residual processes and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
 | Under load | Send `SIGTERM` to serve main process | serve stops accepting new requests and executes shutdown flow; in-flight requests fail; no residue remains and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
 | Under load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | `Ctrl+C`-style serve shutdown; in-flight requests fail (5xx/connection interruption); service unavailable; after exit there is no residue and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
-| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly |  |
+| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly | |
 | Under load | Send `SIGTERM` to all related processes | All processes exit gracefully; in-flight requests fail; service unavailable; GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
 | OOM | Occupy all free GPU memory via an extra process | After OOM injection process starts, GPU memory is continuously saturated; service enters unavailable/degraded state and health check drops to 503; different request types (chat/speech, etc.) fail fast within a fixed time and return 500 (no hanging) | [#4285](https://github.com/vllm-project/vllm-omni/issues/4285) |

@@ -0,0 +1,259 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Run Mage-VL offline end-to-end inference through the upstream Mage scripts.
+
+Mage-VL currently ships its custom Mage-ViT and codec processors as
+Transformers remote code in the Microsoft Mage repository. This example keeps
+the vLLM-Omni entrypoint stable while delegating the checkpoint-specific forward
+path to that upstream implementation.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+from importlib.util import find_spec
+from pathlib import Path
+
+DEFAULT_MODEL = "microsoft/Mage-VL"
+DEFAULT_IMAGE_QUESTION = "Describe this image in detail."
+DEFAULT_VIDEO_QUESTION = "Describe this video."
+
+
+def _default_mage_repo() -> Path | None:
+    value = os.getenv("MAGE_REPO_DIR")
+    if value:
+        return Path(value)
+    sibling = Path(__file__).resolve().parents[3].parent / "Mage"
+    if sibling.exists():
+        return sibling
+    return None
+
+
+def _resolve_mage_repo(value: str | None) -> Path:
+    repo = Path(value) if value else _default_mage_repo()
+    if repo is None:
+        raise ValueError("Set --mage-repo or MAGE_REPO_DIR to a checkout of https://github.com/microsoft/Mage.")
+    repo = repo.expanduser().resolve()
+    if not (repo / "mage_vl" / "inference_base.py").is_file():
+        raise FileNotFoundError(f"Missing Mage-VL scripts under {repo}. Expected mage_vl/inference_base.py.")
+    return repo
+
+
+def _default_media_path(mage_repo: Path, task: str) -> Path:
+    if task == "image":
+        return mage_repo / "mage_vl" / "assets" / "examples" / "dog.jpg"
+    return mage_repo / "mage_vl" / "assets" / "examples" / "soccer-broadcast.mp4"
+
+
+def _resolve_media_path(path: str | None, mage_repo: Path, task: str) -> Path:
+    resolved = Path(path).expanduser().resolve() if path else _default_media_path(mage_repo, task)
+    if not resolved.is_file():
+        raise FileNotFoundError(f"Input media file not found: {resolved}")
+    return resolved
+
+
+def build_command(args: argparse.Namespace, mage_repo: Path, media_path: Path) -> list[str]:
+    if args.task == "streaming":
+        return [
+            sys.executable,
+            str(mage_repo / "mage_vl" / "inference_streaming.py"),
+            "--video",
+            str(media_path),
+            "--checkpoint",
+            args.model,
+            "--video_backend",
+            args.video_backend,
+            "--num_frames",
+            str(args.num_frames),
+            "--cur_fps",
+            str(args.cur_fps),
+            "--segment_sec",
+            str(args.segment_sec),
+            "--max_new_tokens",
+            str(args.max_new_tokens),
+            "--max_segments",
+            str(args.max_segments),
+            "--gate_threshold",
+            str(args.gate_threshold),
+            "--device",
+            args.device,
+            "--attn_impl",
+            args.attn_impl,
+        ]
+
+    cmd = [
+        sys.executable,
+        str(mage_repo / "mage_vl" / "inference_base.py"),
+        "--mode",
+        "offline",
+        "--question",
+        args.question or (DEFAULT_IMAGE_QUESTION if args.task == "image" else DEFAULT_VIDEO_QUESTION),
+        "--model",
+        args.model,
+        "--max-new-tokens",
+        str(args.max_new_tokens),
+    ]
+    if args.task == "image":
+        cmd.extend(["--image", str(media_path)])
+        return cmd
+
+    cmd.extend(
+        [
+            "--video",
+            str(media_path),
+            "--video-backend",
+            args.video_backend,
+            "--num-frames",
+            str(args.num_frames),
+            "--max-pixels",
+            str(args.max_pixels),
+        ]
+    )
+    if args.video_backend == "codec":
+        cmd.extend(["--codec-engine", args.codec_engine])
+    return cmd
+
+
+def check_runtime_dependencies(task: str) -> None:
+    required = ["torch"]
+    if task == "streaming":
+        required.append("mamba_ssm")
+    missing = [name for name in required if find_spec(name) is None]
+    if not missing:
+        return
+
+    install_hint = (
+        "Install Mage-VL runtime dependencies in this environment first:\n"
+        "  pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n"
+        "  grep -v '^mamba-ssm' /path/to/Mage/mage_vl/requirements.txt > /tmp/mage_vl_requirements_no_mamba.txt\n"
+        "  pip install -r /tmp/mage_vl_requirements_no_mamba.txt\n"
+        "  pip install 'causal-conv1d>=1.4.0' --no-build-isolation\n"
+        "  pip install 'mamba-ssm>=2.2' --no-build-isolation"
+    )
+    raise RuntimeError(f"Missing Python packages: {', '.join(missing)}.\n{install_hint}")
+
+
+def run_command(cmd: list[str], cwd: Path, timeout: int) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        cmd,
+        cwd=str(cwd),
+        input="y\n",
+        text=True,
+        capture_output=True,
+        timeout=timeout,
+        check=False,
+    )
+
+
+def write_outputs(
+    *,
+    args: argparse.Namespace,
+    mage_repo: Path,
+    media_path: Path,
+    cmd: list[str],
+    result: subprocess.CompletedProcess[str] | None,
+    elapsed_s: float | None,
+) -> Path:
+    output_dir = Path(args.output_dir).expanduser().resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    summary = {
+        "task": args.task,
+        "model": args.model,
+        "mage_repo": str(mage_repo),
+        "media_path": str(media_path),
+        "command": cmd,
+        "dry_run": args.dry_run,
+        "elapsed_s": elapsed_s,
+        "returncode": None if result is None else result.returncode,
+    }
+    if result is not None:
+        stdout = result.stdout.strip()
+        stderr = result.stderr.strip()
+        summary["stdout"] = stdout
+        summary["stderr"] = stderr
+        (output_dir / "stdout.txt").write_text(result.stdout, encoding="utf-8")
+        (output_dir / "stderr.txt").write_text(result.stderr, encoding="utf-8")
+    summary_path = output_dir / "summary.json"
+    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+    return summary_path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Mage-VL offline end-to-end inference example.")
+    parser.add_argument(
+        "--mage-repo",
+        default=None,
+        help="Path to a Microsoft Mage checkout. Defaults to MAGE_REPO_DIR.",
+    )
+    parser.add_argument("--model", default=DEFAULT_MODEL, help="HF model id or local checkpoint path.")
+    parser.add_argument("--task", choices=("image", "video", "streaming"), default="image")
+    parser.add_argument("--image-path", default=None, help="Image input. Defaults to the Mage sample image.")
+    parser.add_argument("--video-path", default=None, help="Video input. Defaults to the Mage sample video.")
+    parser.add_argument("--question", default=None, help="Question for image/video offline QA.")
+    parser.add_argument("--video-backend", choices=("frames", "codec"), default="frames")
+    parser.add_argument("--codec-engine", choices=("traditional", "neural"), default="traditional")
+    parser.add_argument("--num-frames", type=int, default=32)
+    parser.add_argument("--max-pixels", type=int, default=150000)
+    parser.add_argument("--max-new-tokens", type=int, default=256)
+    parser.add_argument("--cur-fps", type=float, default=2.0)
+    parser.add_argument("--segment-sec", type=float, default=8.0)
+    parser.add_argument("--max-segments", type=int, default=0)
+    parser.add_argument("--gate-threshold", type=float, default=0.5)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--attn-impl", choices=("flash_attention_2", "sdpa", "eager"), default="flash_attention_2")
+    parser.add_argument("--timeout", type=int, default=1800)
+    parser.add_argument("--output-dir", default="outputs/mage_vl")
+    parser.add_argument("--dry-run", action="store_true", help="Build the command and summary without executing it.")
+    args = parser.parse_args()
+    if args.num_frames <= 0:
+        parser.error("--num-frames must be positive")
+    if args.segment_sec <= 0:
+        parser.error("--segment-sec must be positive")
+    if args.task == "image" and args.video_path:
+        parser.error("--video-path cannot be used with --task image")
+    if args.task in {"video", "streaming"} and args.image_path:
+        parser.error("--image-path cannot be used with video tasks")
+    return args
+
+
+def main() -> None:
+    args = parse_args()
+    mage_repo = _resolve_mage_repo(args.mage_repo)
+    media_arg = args.image_path if args.task == "image" else args.video_path
+    media_path = _resolve_media_path(media_arg, mage_repo, args.task)
+    cmd = build_command(args, mage_repo, media_path)
+
+    result = None
+    elapsed_s = None
+    if not args.dry_run:
+        check_runtime_dependencies(args.task)
+        start = time.perf_counter()
+        result = run_command(cmd, cwd=mage_repo, timeout=args.timeout)
+        elapsed_s = time.perf_counter() - start
+
+    summary_path = write_outputs(
+        args=args,
+        mage_repo=mage_repo,
+        media_path=media_path,
+        cmd=cmd,
+        result=result,
+        elapsed_s=elapsed_s,
+    )
+    print(f"summary saved to {summary_path}")
+    if result is not None:
+        if result.stdout:
+            print(result.stdout, end="")
+        if result.returncode != 0:
+            if result.stderr:
+                print(result.stderr, file=sys.stderr, end="")
+            raise SystemExit(result.returncode)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import json
+import time
+from pathlib import Path
+
+import websockets
+
+
+async def run(args: argparse.Namespace) -> None:
+    headers = {"Authorization": f"Bearer {args.auth_token}"} if args.auth_token else None
+    started = time.perf_counter()
+    async with websockets.connect(
+        f"{args.url}?session_id={args.session_id}",
+        additional_headers=headers,
+        max_size=args.max_message_bytes,
+        open_timeout=args.timeout,
+    ) as websocket:
+        created = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout))
+        print(json.dumps(created, ensure_ascii=False))
+        await websocket.send(
+            json.dumps(
+                {
+                    "type": "input.append",
+                    "modality": "text",
+                    "data": args.prompt,
+                }
+            )
+        )
+        encoded = base64.b64encode(Path(args.video).read_bytes()).decode("ascii")
+        await websocket.send(
+            json.dumps(
+                {
+                    "type": "input.append",
+                    "modality": "video",
+                    "data": {
+                        "kind": "video/mp4",
+                        "segment_id": args.segment_id,
+                        "pts_ms": args.pts_ms,
+                        "duration_ms": args.duration_ms,
+                        "data": {"video_base64": encoded},
+                    },
+                }
+            )
+        )
+        while True:
+            event = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout))
+            print(json.dumps(event, ensure_ascii=False))
+            if event.get("type") in {"response.done", "error"}:
+                break
+        await websocket.send(json.dumps({"type": "close"}))
+    print(f"elapsed_s={time.perf_counter() - started:.3f}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Send one Mage-VL video window over the duplex WebSocket.")
+    parser.add_argument("--url", default="ws://127.0.0.1:8090/v1/mage-vl/duplex")
+    parser.add_argument("--video", required=True)
+    parser.add_argument("--prompt", default="Describe this video segment in detail.")
+    parser.add_argument("--session-id", default="mage-vl-example")
+    parser.add_argument("--segment-id", default="segment-0")
+    parser.add_argument("--pts-ms", type=int, default=0)
+    parser.add_argument("--duration-ms", type=int, default=4000)
+    parser.add_argument("--auth-token")
+    parser.add_argument("--timeout", type=float, default=120.0)
+    parser.add_argument("--max-message-bytes", type=int, default=16 * 1024 * 1024)
+    args = parser.parse_args()
+    asyncio.run(run(args))
+
+
+if __name__ == "__main__":
+    main()
@@ -13,7 +13,7 @@ For the full list of supported architectures across all modalities, see
 ## Supported Models

 | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
-|---|---|---|---|---|---|
+| --- | --- | --- | --- | --- | --- |
 | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
 | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
 | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
@@ -104,6 +104,7 @@ For full request-shape documentation (all parameters, response formats, error co
 2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.

 ### Launch
+
 ```bash
 vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
 # or:
@@ -111,6 +112,7 @@ bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-T
 ```

 ### Sending requests
+
 ```bash
 # Voice cloning (required)
 python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
@@ -127,11 +129,13 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
 ```

 ### Gradio demo
+
 ```bash
 bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
 ```

 ### Notes
+
 - Output: 24 kHz mono WAV via HiFT vocoder.
 - `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
 - Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
@@ -151,6 +155,7 @@ pip install 'vllm-omni[indextts2]'
 ```

 ### Launch
+
 ```bash
 vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
 # or, to pass the bundled deploy config explicitly:
@@ -164,6 +169,7 @@ bash examples/online_serving/text_to_speech/indextts2/run_server.sh
 ```

 ### Sending requests
+
 ```bash
 # Voice cloning (ref_audio required)
 python examples/online_serving/text_to_speech/indextts2/speech_client.py \
@@ -194,6 +200,7 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
 ```

 ### Notes
+
 - Output: 22.05 kHz mono WAV.
 - Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; neither version provides a built-in text-only preset voice.
 - IndexTTS-2.5 request controls `lang` and `text_normalization` are carried in `extra_params`; the client does this for `--model-version 2.5`.
@@ -212,6 +219,7 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.

 ### Prerequisites
+
 ```bash
 pip install fish-speech
 ```
@@ -237,14 +245,17 @@ export VLLM_OMNI_FISH_KVCACHE_ATTN=0
 ```

 ### Launch
+
 ```bash
 vllm serve fishaudio/s2-pro --omni --port 8091
 # or:
 ./fish_speech/run_server.sh
 ```
+
 The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`).

 ### Voice cloning
+
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
@@ -257,6 +268,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 ```

 ### CLI client
+
 ```bash
 cd examples/online_serving/text_to_speech/fish_speech
 python speech_client.py --text "Hello, how are you?"
@@ -264,12 +276,14 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
 ```

 ### Gradio demo
+
 ```bash
 ./fish_speech/run_gradio_demo.sh             # launches server + Gradio
 python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server already running
 ```

 ### Notes
+
 - Output: 44.1 kHz mono.
 - Streaming PCM player command must use `-r 44100`.

@@ -280,10 +294,13 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091  # if server
 Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`.

 ### Launch
+
 ```bash
 bash examples/online_serving/text_to_speech/ming_tts/run_server.sh
 ```
+
 Equivalent manual command:
+
 ```bash
 vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
    --deploy-config vllm_omni/deploy/ming_tts.yaml \
@@ -292,12 +309,14 @@ vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
 ```

 ### Sending requests
+
 ```bash
 python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
    --text "你好，这是 Ming 在线语音合成测试。"
 ```

 Structured dialect control:
+
 ```bash
 python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
    --text "我觉得社会企业同个人都有责任" \
@@ -306,6 +325,7 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
 ```

 Zero-shot cloning:
+
 ```bash
 python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
    --text "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。" \
@@ -314,6 +334,7 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
 ```

 ### Notes
+
 - `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook.
 - Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples.
 - Full request details live in [`ming_tts/README.md`](ming_tts/README.md).
@@ -325,11 +346,14 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
 Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`.

 ### Launch
+
 ```bash
 # from repo root
 bash examples/online_serving/text_to_speech/ming_flash_omni_tts/run_server.sh
 ```
+
 Equivalent manual command:
+
 ```bash
 vllm serve Jonathan1909/Ming-flash-omni-2.0 \
    --deploy-config vllm_omni/deploy/ming_flash_omni_tts.yaml \
@@ -338,6 +362,7 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 \
 ```

 ### Sending requests
+
 ```bash
 python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.py \
    --text "我们当迎着阳光辛勤耕作，去摘取，去制作，去品尝，去馈赠。" \
@@ -345,6 +370,7 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
 ```

 ASMR-style caption via `instructions`:
+
 ```bash
 python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.py \
    --text "我会一直在这里陪着你，直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗？" \
@@ -353,6 +379,7 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
 ```

 ### Notes
+
 - Server uses `use_zero_spk_emb=True` and the cookbook decode defaults (`max_decode_steps=200`, `cfg=2.0`, `sigma=0.25`, `temperature=0.0`). For other caption fields (`语速`, `基频`, `IP`, BGM, etc.) or overriding decode args, use the offline example where `additional_information` is set explicitly.
 - This is the online counterpart of [`examples/offline_inference/text_to_speech/ming_flash_omni_tts/`](../../offline_inference/text_to_speech/ming_flash_omni_tts/).
 - For multimodal Ming-flash-omni online serving, see [`examples/online_serving/ming_flash_omni/`](../../ming_flash_omni/).
@@ -366,14 +393,17 @@ Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every
 > The OpenAI-schema `voice` and `ref_text` fields are accepted but ignored — `voice_clone` does not consume a transcript, and upstream's `continuation` mode (the only path that accepts `prompt_text`) emits near-silent output, so it is not exposed here. Sample reference clips ship in the upstream repo under [`assets/audio/`](https://github.com/OpenMOSS/MOSS-TTS-Nano/tree/main/assets/audio).

 ### Launch
+
 ```bash
 vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
 # or:
 ./moss_tts_nano/run_server.sh
 ```
+
 The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--deploy-config`, `--trust-remote-code`, or `--enforce-eager` flags are needed.

 ### Sending requests
+
 ```bash
 # One-off fetch of a sample reference clip; cache under XDG_CACHE_HOME.
 REF_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/moss-tts-nano"
@@ -392,6 +422,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 ```

 ### Streaming PCM
+
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
@@ -405,6 +436,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 ```

 ### Gradio demo
+
 ```bash
 # Option 1: launch server + Gradio together
 ./moss_tts_nano/run_gradio_demo.sh
@@ -412,9 +444,11 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 # Option 2: server already running
 python moss_tts_nano/gradio_demo.py --api-base http://localhost:8091
 ```
-Then open http://localhost:7860 in your browser.
+
+Then open <http://localhost:7860> in your browser.

 ### Notes
+
 - Output is 48 kHz mono PCM (the upstream tokenizer is internally stereo at 48 kHz; the wrapper averages to mono before reaching the engine).
 - Standard `/v1/audio/speech` request shape: `input`, `ref_audio` (base64 data URL), `response_format`, `stream`, `max_new_tokens`. The `voice` and `ref_text` fields from the OpenAI schema are accepted but ignored.

@@ -425,12 +459,15 @@ Then open http://localhost:7860 in your browser.
 Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.

 ### Prerequisites
+
 ```bash
 huggingface-cli download k2-fsa/OmniVoice
 ```
+
 Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`.

 ### Launch
+
 ```bash
 vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
 # or:
@@ -438,6 +475,7 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
 ```

 ### CLI client
+
 ```bash
 cd examples/online_serving/text_to_speech/omnivoice
 # Text-only (auto voice)
@@ -463,13 +501,12 @@ python speech_client.py --text "Hello, how are you?" --seed 42

 The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`.

-
 ## Qwen3-TTS

 Three model variants exposed via separate checkpoints:

 | Variant | HF repo | Use |
-|---|---|---|
+| --- | --- | --- |
 | CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions |
 | VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description |
 | Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio |
@@ -477,6 +514,7 @@ Three model variants exposed via separate checkpoints:
 Each variant ships smaller `0.6B` companions where available.

 ### Launch
+
 ```bash
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
 # or:
@@ -486,11 +524,13 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
 ```

 ### Executor backend
+
 Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.

 To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.

 ### Sending requests
+
 ```bash
 # CustomVoice with a predefined speaker
 python qwen3_tts/openai_speech_client.py \
@@ -516,7 +556,9 @@ python qwen3_tts/openai_speech_client.py \
 ```

 ### Voices endpoint
+
 List available voices, or upload a custom one for Base cloning:
+
 ```bash
 # List
 curl http://localhost:8091/v1/audio/voices
@@ -529,10 +571,18 @@ curl -X POST http://localhost:8091/v1/audio/voices \
    -F "ref_text=The exact transcript of the audio sample." \
    -F "speaker_description=warm narrator"
 ```
-Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
+
+For Qwen3-TTS, uploaded voices are Base voice-cloning inputs and require a Base
+checkpoint. When a request names an uploaded voice, the server infers
+`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
+CustomVoice speakers and require a CustomVoice checkpoint.

 ### Precomputed custom voices
-For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
+
+For reused Base voice-cloning speakers, precompute the reference artifacts once
+and load them at server startup. Precomputed voices use the same Base task and
+checkpoint-matching rules as uploaded voices:
+
 ```bash
 python qwen3_tts/precompute_custom_voice.py \
    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
@@ -542,11 +592,15 @@ python qwen3_tts/precompute_custom_voice.py \
    --mode icl \
    --output-dir /path/to/custom_voices
 ```
+
 `--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
+
 ```yaml
 custom_voice_dir: /path/to/custom_voices
 ```
+
 Then start the server with that config and call the Speech API with only the voice name:
+
 ```bash
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml

@@ -557,6 +611,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 ```

 ### Streaming PCM
+
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \
    -H "Content-Type: application/json" \
@@ -569,16 +624,20 @@ curl -X POST http://localhost:8091/v1/audio/speech \
        "response_format": "pcm"
    }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
 ```
+
 Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.

 ### Streaming WebSocket
+
 The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
+
 ```bash
 python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
 python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
 ```

 `input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
+
 ```bash
 python qwen3_tts/streaming_speech_client.py \
    --text "First utterance." \
@@ -586,6 +645,7 @@ python qwen3_tts/streaming_speech_client.py \
 ```

 To receive word-level timestamps, launch the server with a forced aligner:
+
 ```bash
 vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
    --omni \
@@ -593,7 +653,9 @@ vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
    --trust-remote-code \
    --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
 ```
+
 Then request PCM JSON sidecar chunks:
+
 ```bash
 python qwen3_tts/streaming_speech_client.py \
    --text "Hello world. How are you?" \
@@ -601,6 +663,7 @@ python qwen3_tts/streaming_speech_client.py \
    --response-format pcm \
    --word-timestamps
 ```
+
 The client writes one PCM file per sentence and a matching
 `sentence_XXX_timestamps.json` sidecar.

@@ -613,15 +676,18 @@ still returns — use the WebSocket path for long transcripts.

 To *see* the alignment instead of reading a JSON sidecar, run the
 word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
+
 ```bash
 python qwen3_tts/word_timestamps_demo.py --api-base http://localhost:8091
 ```
+
 Each sentence's audio plays in an `<audio>` element while its text is rendered
 as inline word spans; the current word highlights as `audio.currentTime`
 crosses each `start_ms`. The **Stop (barge-in)** button cuts playback and
 reports the last-spoken word, useful for the voice-agent barge-in case.

 ### Gradio demos
+
 ```bash
 ./qwen3_tts/run_gradio_demo.sh                              # CustomVoice (default)
 ./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign
@@ -629,12 +695,15 @@ reports the last-spoken word, useful for the voice-agent barge-in case.
 ```

 ### Speaker embedding interpolation
+
 `qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage.

 ### Batch client
+
 `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.

 ### Notes
+
 - Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
 - With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
 - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
@@ -646,12 +715,15 @@ reports the last-spoken word, useful for the voice-agent barge-in case.
 Single-stage native AR TTS at 48 kHz.

 ### Launch
+
 ```bash
 vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000
 ```
+
 Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks.

 ### Sending requests
+
 ```bash
 # Zero-shot synthesis
 python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2."
@@ -661,10 +733,13 @@ python voxcpm2/openai_speech_client.py \
    --text "This should sound like the reference speaker." \
    --ref-audio /path/to/reference.wav
 ```
+
 The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.

 ### Precomputed custom voices
+
 For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
+
 ```bash
 python voxcpm2/precompute_custom_voice.py \
    --model openbmb/VoxCPM2 \
@@ -674,16 +749,21 @@ python voxcpm2/precompute_custom_voice.py \
    --prompt-text "Original transcript of the reference audio" \
    --output-dir /path/to/custom_voices
 ```
+
 Add the output directory to the deploy config:
+
 ```yaml
 custom_voice_dir: /path/to/custom_voices
 ```
+
 After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.

 ### Gradio demo (gapless streaming via AudioWorklet)
+
 ```bash
 python voxcpm2/gradio_demo.py
 ```
+
 Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`.

 ---
@@ -693,24 +773,31 @@ Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free p
 Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific.

 ### Prerequisites
+
 Latest `mistral_common` with `SpeechRequest` support:
+
 ```bash
 pip install -e /path/to/mistral-common  # or upgrade from PyPI when available
 ```

 ### Launch
+
 ```bash
 vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
 ```
+
 Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`.

 ### Gradio demo
+
 ```bash
 python voxtral_tts/gradio_demo.py
 ```
+
 The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients).

 ### Notes
+
 - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
 - Voice cloning is gated upstream and may require a recent `mistral_common`.
 - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
@@ -39,9 +39,10 @@ rendering. The Edit and Edit-Turbo checkpoints use the same native

 ## Hardware Support

-This recipe documents tested configurations for CUDA GPU serving. The native
-pipeline runs single-GPU; multi-GPU parallelism, CPU offload, and cache
-acceleration are not yet supported for this model (see Notes).
+This recipe documents supported configurations for CUDA GPU serving. The native
+pipeline supports single-GPU inference and classifier-free-guidance (CFG)
+parallelism. CPU offload, cache acceleration, and other multi-GPU dimensions
+remain unsupported (see Notes).

 ## GPU

@@ -108,8 +109,24 @@ curl -s http://localhost:8091/v1/chat/completions \
  The model's maximum native resolution is 2K.
 - **Known limitations (not yet supported):** CPU offload
  (`--enable-cpu-offload` / `--enable-layerwise-offload`), Cache-DiT
-  (`--cache-backend cache_dit`), and multi-GPU parallelism (TP / SP / CFG /
-  HSDP) are planned follow-ups and are not validated for this model yet.
+  (`--cache-backend cache_dit`), and TP / SP / HSDP multi-GPU parallelism.
+
+### 2 x H100 (CFG parallel, Base T2I)
+
+Base T2I has two predictions per guided denoising step, so two CFG ranks map
+one-to-one to the positive and negative-text branches:
+
+```bash
+vllm serve Boogu/Boogu-Image-0.1-Base \
+  --omni \
+  --cfg-parallel-size 2 \
+  --port 8091
+```
+
+Use `guidance_scale > 1.0` to activate the parallel two-branch path. Requests
+with `guidance_scale=1.0` remain valid on the same server: every rank evaluates
+only the positive branch, no negative embeddings are built, and no guidance is
+applied.

 ## Image editing (Boogu-Image-0.1-Edit)

@@ -124,6 +141,33 @@ is sent).
 vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
 ```

+For text-only or image-only guidance, use two CFG ranks:
+
+```bash
+vllm serve Boogu/Boogu-Image-0.1-Edit \
+  --omni \
+  --cfg-parallel-size 2 \
+  --port 8091
+```
+
+Double guidance has three predictions per denoising step. It supports either
+two ranks (round-robin assignment: rank 0 runs branches 0 and 2) or three ranks
+(one branch per rank):
+
+```bash
+# Practical two-GPU configuration
+vllm serve Boogu/Boogu-Image-0.1-Edit \
+  --omni \
+  --cfg-parallel-size 2 \
+  --port 8091
+
+# Maximum branch parallelism
+vllm serve Boogu/Boogu-Image-0.1-Edit \
+  --omni \
+  --cfg-parallel-size 3 \
+  --port 8091
+```
+
 ### Verification

 Edit an image with `/v1/images/edits` (the model-card example — change a photo
@@ -175,11 +219,29 @@ curl -s http://localhost:8091/v1/chat/completions \
    default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
    (3 model predictions per step), steering more strongly toward the reference
    image.
+- **CFG-size recommendations:**
+
+  | Request mode | Scales | Predictions per step | Recommended CFG size |
+  | --- | --- | ---: | ---: |
+  | Base T2I text guidance | `guidance_scale > 1`, no reference | 2 | 2 |
+  | Edit text-only guidance | `guidance_scale > 1`, `guidance_scale_2 = 1` | 2 | 2 |
+  | Edit image-only guidance | `guidance_scale = 1`, `guidance_scale_2 > 1` | 2 | 2 |
+  | Edit double guidance | both scales `> 1` | 3 | 2 (practical) or 3 (maximum parallelism) |
+  | CFG off | both scales `= 1` | 1 | Reuse a size-2 server safely; no speedup expected |
+
+  Double guidance preserves Boogu's original combination:
+
+  ```text
+  pred = positive_with_reference
+       + (text_scale - 1) * (positive_with_reference - negative_with_reference)
+       + (image_scale - 1) * (negative_with_reference - uncond)
+  ```
 - **Output resolution:** the output size follows the reference image (upstream
  `align_res`, on by default for a single reference), so `height`/`width` are
  derived from the input and requested sizes are not applied for edits.
- **Same limitations** as the Base checkpoint apply (no CPU offload, Cache-DiT,
-  or multi-GPU parallelism yet).
+- **Same limitations** as the Base checkpoint apply except that CFG parallelism
+  is supported as described above; CPU offload, Cache-DiT, TP, SP, and HSDP are
+  still unsupported.
 - **Offline editing:** the shared offline example
  [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_to_image.md)
  supports Boogu-Image-Edit directly
@@ -251,3 +313,21 @@ curl -s http://localhost:8091/v1/chat/completions \
  the regular Edit checkpoint applies.
 - **Known limitations:** the same single-GPU limitations as Base and Edit apply;
  CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
+
+## Performance validation
+
+The checked-in single-device A40 measurements remain the current Boogu
+baseline: Base T2I is about 16.37 s request latency at concurrency 1 and Edit
+text-only guidance is about 29.64 s for 512x512, 28-step requests, with roughly
+36.8 GB peak memory. CFG-parallel A/B numbers must be collected on the same
+hardware and software revision before reporting a speedup; do not compare a
+multi-GPU result against these A40 numbers as if it were a controlled A/B.
+
+Run the dedicated perf rows in
+`tests/dfx/perf/tests/test_boogu_image_vllm_omni.json` and
+`tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json` on 2x/3x H100, then
+record latency, throughput, and per-GPU peak memory beside the single-device
+run. The expected shape is close to the general CFG-parallel range (roughly
+1.5x-1.8x for balanced two-branch workloads); the three-branch Edit result will
+depend on whether two-rank round-robin or full three-rank execution is used and
+must be reported from measurement rather than assumed.
@@ -0,0 +1,114 @@
+# Mage-VL
+
+> Codec-native image/video understanding plus proactive full-duplex streaming.
+
+## Summary
+
+- Vendor: Microsoft
+- Model: [`microsoft/Mage-VL`](https://huggingface.co/microsoft/Mage-VL)
+- Task: Image understanding, video understanding, codec-native video understanding,
+  and event-gated streaming commentary
+- Mode: Full-duplex WebSocket server and session adapter in
+  `vllm_omni/experimental/fullduplex/mage_vl`
+- Maintainer: Community
+
+## Status
+
+This experimental integration provides the Mage-VL adapter contract: bounded causal
+visual windows, codec-window input, rolling gate evaluation, explicit query handling,
+per-session state isolation, and stale-output cancellation through the shared
+full-duplex runtime.
+
+The production transport exposes `WebSocket /v1/mage-vl/duplex`, with bearer-token
+authentication, bounded concurrent sessions and message sizes, idle timeouts,
+disconnect cancellation, per-session state cleanup, and serialized access to the
+shared Transformers checkpoint. Native vLLM model execution remains follow-up work.
+
+The default `frames` backend uses decoded frames and has no codec preprocessing
+requirement. The optional `codec` backend requires `ffmpeg` and `ffprobe`, plus the
+`cv-preinfer` executable provided by `codec-video-prep>=0.2.5`, on `PATH`. Server
+startup fails with an actionable error when `cv-preinfer` is unavailable; use
+`--video-backend frames` when codec preprocessing is not installed.
+
+## Serving
+
+```bash
+python -m vllm_omni.experimental.fullduplex.mage_vl.serving.server \
+  --model microsoft/Mage-VL --host 0.0.0.0 --port 8090 \
+  --auth-token "$MAGE_VL_API_KEY" --max-sessions 32 \
+  --video-backend frames --num-frames 8 --target-fps 1 \
+  --gate-threshold 0.5 --attn-impl sdpa
+
+python examples/online_serving/mage_vl/duplex_client.py \
+  --url ws://127.0.0.1:8090/v1/mage-vl/duplex \
+  --auth-token "$MAGE_VL_API_KEY" \
+  --video /path/to/segment.mp4 \
+  --prompt "Describe this video segment in detail."
+```
+
+Each input window is a JSON `input.append` event whose `data.data.video_base64`
+contains an MP4 segment. A text `input.append` before the video requests an explicit
+answer; without text, the StreamMind gate decides whether to respond proactively.
+The server emits `session.created`, `response.created`, one or more
+`response.delta`, and `response.done` events. Clients can send `response.cancel` for
+barge-in and `close` for graceful shutdown.
+
+The adapter can still be wired to another backend through `--adapter module:factory`,
+including:
+
+- a Transformers `trust_remote_code=True` backend for offline/local testing;
+- the model author's SGLang-compatible serving path;
+- a future native vLLM `mage_vl` backend.
+
+## Adapter Shape
+
+```python
+from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig
+from vllm_omni.experimental.fullduplex.mage_vl import MageVLDuplexAdapter, MageVLDuplexRuntime
+
+
+adapter = MageVLDuplexAdapter(gate=run_mage_gate, generate=run_mage_decoder)
+session = DuplexSession(
+    "stream-1",
+    DuplexSessionConfig(
+        input_modalities=("codec_window", "video", "image", "text"),
+        output_modalities=("text",),
+        proactive=True,
+    ),
+)
+runtime = MageVLDuplexRuntime(session, adapter)
+await runtime.run(input_events, emit_event)
+```
+
+Input events use the shared full-duplex protocol. For decoded video/image, pass the frame
+payload directly. For codec-native streaming, pass `modality="codec_window"` and include
+fields such as `segment_id`, `pts_ms`, `duration_ms`, `codec`, and `metadata`.
+
+```python
+{
+    "type": "input.append",
+    "modality": "codec_window",
+    "data": {
+        "kind": "h264",
+        "segment_id": "seg-42",
+        "pts_ms": 42000,
+        "duration_ms": 1000,
+        "codec": {
+            "motion_vectors": ...,
+            "residual_energy": ...,
+        },
+        "metadata": {"gop": "P"},
+    },
+}
+```
+
+The gate callable receives the current rolling window and returns either a
+`MageVLGateDecision`, a mapping with `should_respond` / `text` / `event_id`, a bool, or
+a text string. Repeated `event_id` values are de-duplicated at the session adapter.
+
+## Verification
+
+```bash
+pytest tests/e2e/features/fullduplex/test_mage_vl_adapter.py \
+  tests/examples/offline_inference/test_mage_vl.py
+```
@@ -278,7 +278,7 @@ steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain

 ### Attention Backends

-On datacenter Blackwell GPUs, MiniMax H3 defaults to dense BF16
+On supported datacenter Blackwell systems, MiniMax H3 defaults to dense BF16
 `TRTLLM_ATTN`; no attention backend flag is required. To select it explicitly,
 use:

@@ -298,17 +298,33 @@ FA4 remains available by explicitly selecting the `FLASH_ATTN` backend:
 --diffusion-attention-backend FLASH_ATTN
 ```

-On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
-`Using CuTe FlashAttention-4 on Blackwell` before recording FA4 measurements.
+With the optional `[fa4]` dependency installed, `FLASH_ATTN` prefers FA4 on
+Blackwell. Confirm the server log contains `Using CuTe FlashAttention-4 on
+Blackwell` before recording FA4 measurements.

-`TRTLLM_ATTN` additionally supports two **lossy** optimizations for the long main
-DiT attention sequence: SAGE attention quantization and Skip-Softmax Sparse
-Attention. SAGE quantizes Q/K to the configured dtype and V to FP8. This example uses
-`fp8_e4m3` for Q/K; B200 also supports `int8` Q/K. The TRTLLM SAGE path fixes V
-to FP8, so vLLM-Omni only exposes the Q/K dtype. The token refiner is a short
-attention path, so the `per_role` override leaves SAGE and Skip-Softmax disabled
-for it. The example enables the calibration-free Skip-Softmax path with
-`threshold=0.05`, after the normalized timestep reaches `0.97`:
+`TRTLLM_ATTN` additionally offers two **lossy** optimizations for the long main
+DiT attention sequence: SAGE attention quantization and Skip-Softmax sparse
+attention. Both work under the pure Ulysses parallelism of the profile above
+(`--usp 4 --ring 1`). The example below enables both:
+
+- SAGE with `fp8_e4m3` Q/K; P and V are always FP8 in this kernel. B200
+  additionally supports `int8` Q/K, which preserves accuracy better than FP8.
+- Skip-Softmax with a direct `threshold=0.05` (the calibrated
+  `target_sparsity` control needs ModelOpt metadata that the official H3
+  checkpoint does not include). Together with the cutoff below this is a
+  **conservative** setting: a low threshold skips only clearly negligible tiles,
+  and a cutoff close to `1.0` leaves a substantial dense prefix. Raise
+  `threshold` or lower `disabled_until_timestep` for more speedup once quality
+  is verified.
+- `disabled_until_timestep=0.97` keeps the early high-noise steps dense. The
+  gate compares against the video sigma, which H3's default flow shift of 12
+  keeps high for much of the run: at 50 steps, `0.99`, `0.97`, and `0.95`
+  leave the first 6, 14, and 19 of 49 denoiser forwards dense. See the
+  [Skip-Softmax design](https://github.com/vllm-project/vllm-omni/blob/main/docs/design/feature/skip_softmax.md#timestep-gating)
+  for how the cutoff maps to steps.
+- A `per_role` entry that keeps the token refiner, a short attention path,
+  dense. A per-role spec does not inherit `quant` or `skip_softmax` from
+  `default`.

 ```bash
 --diffusion-attention-config '{
@@ -332,10 +348,13 @@ for it. The example enables the calibration-free Skip-Softmax path with
 }'
 ```

-For configuration details, see
-[TRTLLM_ATTN Backend and Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-backend-and-skip-softmax)
+Both optimizations trade fidelity for speed and their effects compound.
+Compare against dense output on the same prompt and seed before adopting them.
+For the full key reference, see
+[Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends/trtllm.md#skip-softmax)
 and
-[TRTLLM_ATTN SAGE Quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-sage-quantization).
+[SAGE quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends/trtllm.md#sage-quantization)
+in the TRTLLM attention guide.

 ### Text encoder tensor parallelism

@@ -962,13 +981,27 @@ mostly base H3 weights on a four-step schedule. Offload is refused for the same
 reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and
 `--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast.

-> [!NOTE]
-> Only the **Dense / Data-Free** variant is supported today. The three VSA
-> variants carry `.set_weight` compression gates for modules that do not exist
-> in vLLM-Omni's dense H3 attention, and the release states dense attention is
-> not a supported substitute for them, so they are refused at startup rather
-> than silently run wrong. Video Sparse Attention for H3's packed
-> `[text | cond | audio | video]` sequence is follow-up work.
+The VSA variants are supported through FastVideo's external kernel. Install a
+`fastvideo-kernel` build that provides the `fastvideo_kernel` Python module,
+then add the following flags to the same command:
+
+```bash
+--diffusion-attention-backend FASTVIDEO_VSA \
+--fastvideo-vsa-topk 64
+```
+
+FastH3 VSA applies its learned `.set_weight` compression gates to the complete
+packed `[text | cond | audio | video]` document using the official H3 geometry:
+text/condition/audio prefix tiles never cross segment boundaries, and target
+video rows use `(4, 4, 4)` 3-D tiles (64 tokens). Prefix queries remain dense;
+video queries select all prefix tiles plus the configured top-k video tiles.
+Pure Ulysses sequence parallelism is supported: the learned gate follows the
+same sequence-to-head all-to-all as Q/K/V before VSA runs. Ring and all-gather
+SP remain unsupported because they do not present a complete packed sequence
+to each block-sparse kernel rank.
+
+The Dense / Data-Free variant does not require `fastvideo-kernel` and should be
+served with a dense attention backend.

 Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at
 1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded:
@@ -119,6 +119,30 @@ MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
 bash examples/online_serving/text_to_video/run_curl_sana_video.sh
 ```

+##### Parallel native serving
+
+The native pipelines support tensor parallelism (up to 2 GPUs) and CFG
+parallelism. Tensor parallelism on 2 GPUs:
+
+```bash
+vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
+  --omni \
+  --model-class-name SanaVideoPipeline \
+  --tensor-parallel-size 2 \
+  --dtype bfloat16 \
+  --port 8091
+```
+
+For CFG parallelism on 2 GPUs, replace `--tensor-parallel-size 2` with
+`--cfg-parallel-size 2`; it splits the guided and unguided branches across the
+two GPUs and only helps when `guidance_scale` is above 1. Passing both flags
+combines the two on 4 GPUs.
+
+Whether tensor parallelism lowers latency depends on the interconnect. Each
+transformer block adds an all-reduce, so it speeds up generation only on fast
+GPU links such as NVLink and can be slower than a single GPU on PCIe-only
+systems. Measure on your hardware before enabling it.
+
 To run the black-box compatibility backend for T2V, replace the server script
 with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request
 works; `num_frames` is adapted to Diffusers' `frames` argument. The script
@@ -178,6 +202,75 @@ Use the native `SanaVideoPipeline` and `SanaImageToVideoPipeline` for the
 primary SANA execution paths. The Diffusers adapter is retained as a
 validated compatibility/reference backend.

+#### Native Cache-DiT and CPU offload
+
+The native T2V and I2V pipelines support Cache-DiT, model-level CPU offload,
+and layerwise CPU offload through the common diffusion flags. Layerwise mode
+keeps non-block transformer modules, the text encoder, and the VAE on the
+runtime device while prefetching DiT blocks in order.
+
+Add one of the following flag sets to either native offline command above:
+
+```bash
+# Cache-DiT
+--cache-backend cache_dit
+
+# Model-level CPU offload
+--enable-cpu-offload
+
+# Layerwise CPU offload
+--enable-layerwise-offload
+
+# Cache-DiT plus model-level offload
+--cache-backend cache_dit --enable-cpu-offload
+
+# Cache-DiT plus layerwise offload
+--cache-backend cache_dit --enable-layerwise-offload
+
+# Distributed layerwise offload (weights sharded across the DP group)
+--enable-distributed-layerwise-offload
+```
+
+If both CPU offload flags are supplied, the common offloader keeps its existing
+layerwise precedence. Cache-DiT refreshes each request from the explicit
+`num_inference_steps`; when that field is omitted, both native SANA pipelines
+default to 50 inference steps.
+
+Cache-DiT and CPU offload require TP1, CFG1, and SP1. Combining them with
+tensor, CFG, or sequence parallelism raises before checkpoint components are
+loaded, as do other cache backends such as TeaCache. Cache-DiT cannot be
+combined with distributed layerwise offload: per-rank cache skips would
+desynchronize the weight AllGather.
+
+Measured on one A800-SXM4-80GB (driver 580.126.09, CUDA 13.0, PyTorch
+2.11.0+cu130, Diffusers 0.38.0) with `DIFFUSION_ATTENTION_BACKEND=CUDNN_ATTN`,
+81 frames, 50 steps, seed 42, each configuration in its own process:
+
+| Configuration | 480p T2V latency | 480p T2V generation peak |
+|---|---:|---:|
+| Baseline | 120.19 s | 24.07 GiB |
+| Cache-DiT | 77.11 s (1.56x) | 24.07 GiB |
+| Model CPU offload | 121.15 s | 16.81 GiB (-7.26 GiB) |
+| Layerwise offload | 118.63 s | 20.75 GiB (-3.32 GiB) |
+
+Cache-DiT gives 1.54x on 720p I2V (37.0 s to 24.1 s). Latency is the median of
+three measured runs after one warmup. Offload trades startup peak too: 14.28
+GiB baseline against 6.14 GiB for model-level and 6.80 GiB for layerwise.
+
+Cache-DiT is approximate: it skips block computations, so its output differs
+from an uncached run. Similarity against a reference run is not used as a gate
+here because this pipeline does not reproduce its own trajectory on this
+hardware -- two runs of identical code differ by SSIM 0.9712 / rel_l2 8.7e-2 at
+50 steps. Cache-DiT output was checked frame by frame against the uncached run
+on 480p T2V and 720p I2V instead, and matches it in subject, composition,
+sharpness, and lighting.
+
+The numbers above use the offline example's default cache configuration.
+Lowering `residual_diff_threshold` in `cache_config` caches fewer steps,
+trading speed for a trajectory closer to the uncached run; see the
+[Cache-DiT guide](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md)
+for the tuning knobs.
+
 #### Notes

 - Key flags: select I2V explicitly with `--model-class-name
@@ -194,8 +287,14 @@ validated compatibility/reference backend.
  The denoising loop intentionally retains the checkpoint-compatible
  Diffusers `DPMSolverMultistepScheduler`.
 - Known limitations:
-    - Sequence/tensor/CFG parallelism, Cache-DiT, TeaCache, and step execution
-    are not validated for the native pipeline.
+    - Tensor parallelism (up to 2 GPUs) and CFG parallelism are supported for
+    the native pipeline. Sequence parallelism, TeaCache, and step execution
+    are not validated.
+    - Cache-DiT and CPU offload are limited to TP1/CFG1/SP1. Cache-DiT with
+    distributed layerwise offload is not supported by the native pipeline.
+    - Cache-DiT speedup and offload memory numbers above are single-GPU A800
+    measurements; other hardware will differ. Distributed layerwise offload is
+    only exercised single-rank here.
    - The Diffusers backend is a compatibility path and does not provide native
    vLLM-Omni parallelism or continuous batching.
    - Native describes pipeline and Transformer ownership, not a zero-Diffusers
@@ -166,6 +166,7 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
 - Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
 - Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`.
 - Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
+- Stored voices: Uploaded and precomputed voices use the Base task and require a Base checkpoint. Built-in preset speakers remain CustomVoice voices.
 - Base codec termination: Base requests without an explicit `max_new_tokens`
  use a text-scaled safety ceiling (at least 192 codec frames and no more than
  the configured model limit). If the Talker reaches that ceiling without
@@ -0,0 +1 @@
+首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，<recaption>这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>
@@ -0,0 +1 @@
+首先，我分析所有输入图像：图像1是一个圆形logo，背景为黑色，主体由深蓝、中蓝、浅蓝和白色组成的抽象曲线构成，整体呈现出一种动态的、类似旋涡的视觉效果。图像2展示了四个并排摆放在灰色织物表面上的冰箱贴。这些冰箱贴的造型分别是长颈鹿、小鸟、蜜蜂和海豚，其关键特征是材质：它们具有光滑、高反光的珐琅质感，边缘有明显的银色金属包边，呈现出立体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）转化为一个具有特定三维物理属性（冰箱贴材质）的实体物品。核心在于保留logo的设计精髓，同时精确复刻参考图像中冰箱贴的材质和工艺感。为了构建最终的答案图像，我会将图一的圆形logo作为新冰箱贴的主体图案。这个圆形冰箱贴将拥有与图二冰箱贴一致的厚度和立体感。其表面会覆盖一层高光泽的透明釉质，完美再现珐琅的质感。logo内部的蓝色和白色区域将保持其原有的色彩，但会因为表面的光泽而出现明亮的高光反射。整个圆形的边缘会包裹一圈纤细而有光泽的银色金属边框，这不仅增加了产品的精致感，也与图二中的金属包边工艺相呼应。最后，将这个新制作的圆形冰箱贴放置在图二所示的灰色织物背景上，确保其光照、阴影和反射效果与周围环境和其他冰箱贴完全一致，使其看起来就像是和它们一起被拍摄的，从而创造出一个逼真且风格统一的视觉效果。</think><recaption>这幅图像以产品摄影的精准风格，展示了一枚根据`image_1`的logo设计定制的圆形冰箱贴。最终图像使用`image_2`的分辨率。该冰箱贴居中放置于`image_2`的灰色织物背景之上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝及白色构成的动态旋涡状抽象图案。整个冰箱贴被赋予了与`image_2`中冰箱贴完全一致的高级珐琅质感，表面覆盖着一层晶莹剔透的高光釉面，在柔和的顶光照射下，呈现出清晰而温润的高光反射。一圈纤薄而精致的银色金属边框紧密包裹着圆形主体，其金属光泽与釉面的质感形成精妙对比。冰箱贴在织物背景上投下了一层极其自然的柔和阴影，证明了其与环境光影的无缝融合，整体视觉效果真实可信，宛如一件精心制作的实体工艺品。<relation_1>最终图像完整保留了`image_1`中logo的全部核心设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色四种颜色构成的抽象旋涡图案的精确形态、色彩和布局。</relation_1><relation_2>最终图像的分辨率、背景和材质风格完全来自`image_2`。这包括其作为背景的灰色织物表面的独特纹理，以及应用于新冰箱贴上的、与`image_2`中冰箱贴完全相同的珐琅材质（高光釉面质感）和银色金属边框的物理属性。画面的整体光照环境，包括柔和的光源方向和质感，也与`image_2`保持一致。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep><guidance><timestep_r>[<img>]{3600}<eoi></answer>
@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 from __future__ import annotations

 import sys
@@ -13,6 +13,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "com
 from skip_ci import resolve_ci_decision  # noqa: E402
 from upload_pipeline import (  # noqa: E402
    _expand_mirror_hardwares,
+    _get_mirror_hw_selector,
    _load_bootstrap_steps,
    _render_bootstrap_pipeline,
    _render_test_pipeline,
@@ -116,7 +117,8 @@ def test_yaml_gated_l2_still_enables_image_via_ready_base() -> None:
    assert "if: true" not in rendered


-def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins() -> None:
+def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "")
    doc = {
        "steps": [
            {
@@ -129,8 +131,29 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins() -> None:
    rendered = _render_test_pipeline(doc, changed_files=None)
    step = rendered["steps"][0]
    assert "mirror_hardwares" not in step
-    assert step["agents"]["queue"] == "gpu_1_queue"
-    assert step["plugins"][0]["docker#v5.2.0"]["image"].endswith("$BUILDKITE_COMMIT")
+    assert step["agents"]["queue"] == "l4-k8s"
+    assert step["retry"] == {
+        "automatic": [
+            {"exit_status": -1, "limit": 1},
+            {"exit_status": 128, "limit": 1},
+            {"signal_reason": "agent_stop", "limit": 1},
+            {"signal_reason": "agent_refused", "limit": 1},
+        ],
+    }
+    container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0]
+    assert container["image"].endswith("$BUILDKITE_COMMIT")
+    assert container["resources"]["limits"]["nvidia.com/gpu"] == 1
+
+
+def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
+    step = _expand_mirror_hardwares(
+        {
+            "label": "opt-out",
+            "mirror_hardwares": "l4_1",
+            "retry": {"automatic": [{"exit_status": 255, "limit": 2}]},
+        },
+    )
+    assert step["retry"] == {"automatic": [{"exit_status": 255, "limit": 2}]}


 def test_mirror_hardwares_conflicts_with_explicit_agents() -> None:
@@ -161,6 +184,156 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
    ]


+def _gpu_limit(step: dict) -> int:
+    return step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0]["resources"]["limits"]["nvidia.com/gpu"]
+
+
+def test_mirror_hardwares_mapping_is_rejected() -> None:
+    with pytest.raises(ValueError, match="must be a preset string"):
+        _expand_mirror_hardwares(
+            {"label": "bad", "mirror_hardwares": {"default": "h100_2", "b200": "b200_2"}},
+        )
+
+
+@pytest.mark.parametrize("hardware", [2, "2", "h100_99", "not_a_preset"])
+def test_mirror_hardwares_unknown_name_is_rejected(hardware: int | str) -> None:
+    with pytest.raises(ValueError, match="unknown mirror_hardwares"):
+        _expand_mirror_hardwares({"label": "unknown preset", "mirror_hardwares": hardware})
+
+
+def test_mirror_hardwares_string_ignores_pytest_cards_mark(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "")
+    step = _expand_mirror_hardwares(
+        {
+            "label": "forced",
+            "commands": ['pytest -sv tests/e2e -m "full_model and L4 and cards_4"'],
+            "mirror_hardwares": "h100_1",
+        },
+    )
+    assert step is not None
+    assert step["agents"]["queue"] == "mithril-h100-pool"
+    assert _gpu_limit(step) == 1
+
+
+def test_mirror_hardwares_b200_omits_cuda_strings_and_remaps_inferred(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "b200")
+    assert _expand_mirror_hardwares({"label": "h100", "mirror_hardwares": "h100_4"}) is None
+    assert _expand_mirror_hardwares({"label": "l4", "mirror_hardwares": "l4_1"}) is None
+    b200 = _expand_mirror_hardwares({"label": "b200", "mirror_hardwares": "b200_2"})
+    assert b200 is not None and b200["agents"]["queue"] == "b200-k8s"
+    npu = _expand_mirror_hardwares({"label": "npu", "mirror_hardwares": "a2b3_npu_4"})
+    assert npu is not None and npu["agents"]["queue"] == "ascend-a2b3"
+
+    rendered = _render_test_pipeline(
+        {
+            "steps": [
+                {
+                    "group": ":card_index_dividers: Mixed",
+                    "steps": [
+                        {"label": "H100 only", "mirror_hardwares": "h100_4"},
+                        {
+                            "label": "Count remap",
+                            "commands": [
+                                'pytest -sv tests/e2e -m "full_model and H100 and B200 and omni and cards_2"',
+                            ],
+                        },
+                    ],
+                },
+                {
+                    "group": ":card_index_dividers: H100 only group",
+                    "steps": [{"label": "Skip me", "mirror_hardwares": "h100_1"}],
+                },
+            ],
+        },
+        changed_files=None,
+    )
+    groups = [step.get("group") for step in rendered["steps"]]
+    assert ":card_index_dividers: H100 only group" not in groups
+    mixed = next(step for step in rendered["steps"] if step.get("group") == ":card_index_dividers: Mixed")
+    assert [child["label"] for child in mixed["steps"]] == ["Count remap"]
+    assert mixed["steps"][0]["agents"]["queue"] == "b200-k8s"
+
+
+@pytest.mark.parametrize(
+    ("selector", "expr", "queue", "gpus"),
+    [
+        ("", "H100 and B200 and cards_2", "mithril-h100-pool", 2),
+        ("b200", "H100 and B200 and cards_2", "b200-k8s", 2),
+        ("", "H100 or L4 and cards_4", "mithril-h100-pool", 4),
+        ("", "L4 and B200 and cards_4", "l4-k8s", 4),
+        ("", "H100 and cards_2 and cards_3", "mithril-h100-pool", 3),
+        ("", "H100 and not cards_1", "mithril-h100-pool", 4),
+        ("b200", "H100 and B200 and not cards_1", "b200-k8s", 4),
+    ],
+)
+def test_mirror_hardwares_inferred_from_marks(
+    monkeypatch: pytest.MonkeyPatch,
+    selector: str,
+    expr: str,
+    queue: str,
+    gpus: int,
+) -> None:
+    monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: selector)
+    step = _expand_mirror_hardwares({"label": "job", "commands": [f'pytest -sv tests/e2e -m "{expr}"']})
+    assert step is not None
+    assert step["agents"]["queue"] == queue
+    assert _gpu_limit(step) == gpus
+
+
+@pytest.mark.parametrize(
+    ("selector", "expr"),
+    [
+        ("b200", "H100 and cards_2"),
+        ("", "B200 and cards_2"),
+        ("", "full_model and cards_2"),
+    ],
+)
+def test_mirror_hardwares_inferred_skips_unmatched_chip(
+    monkeypatch: pytest.MonkeyPatch,
+    selector: str,
+    expr: str,
+) -> None:
+    monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: selector)
+    assert _expand_mirror_hardwares({"label": "job", "commands": [f'pytest -sv tests/e2e -m "{expr}"']}) is None
+
+
+def test_mirror_hardwares_inferred_missing_preset_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "")
+    with pytest.raises(ValueError, match="has no preset 'h100_8'"):
+        _expand_mirror_hardwares(
+            {"label": "H100 8-gpu", "commands": ['pytest -sv tests/e2e -m "H100 and cards_8"']},
+        )
+
+
+def test_cpu_step_without_mirror_hardwares_is_unchanged() -> None:
+    step = {"label": "CPU report", "commands": ["echo ok"], "agents": {"queue": "cpu_queue_premerge"}}
+    assert _expand_mirror_hardwares(step) is step
+
+
+@pytest.mark.parametrize(("raw", "expected"), [("", ""), ("  ", ""), ("b200", "b200"), ("B200", "b200")])
+def test_mirror_hw_selector_empty_or_b200(monkeypatch: pytest.MonkeyPatch, raw: str, expected: str) -> None:
+    monkeypatch.setenv("MIRROR_HW", raw)
+    assert _get_mirror_hw_selector() == expected
+
+
+def test_mirror_hw_typo_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("MIRROR_HW", "b20o")
+    with pytest.raises(ValueError, match=r"unsupported MIRROR_HW='b20o'"):
+        _get_mirror_hw_selector()
+    with pytest.raises(ValueError, match=r"unsupported MIRROR_HW='b20o'"):
+        _render_test_pipeline(
+            {
+                "steps": [
+                    {"label": "CPU report", "commands": ["echo ok"]},
+                    {"label": "H100 string", "mirror_hardwares": "h100_4"},
+                ],
+            },
+            changed_files=None,
+        )
+
+
 def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
    rendered = _render_test_pipeline(doc, changed_files=changed_files)
    labels: set[str] = set()
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
        EnvironmentVariableCategory.PUBLIC_OMNI: 24,
        EnvironmentVariableCategory.INHERITED_VLLM: 20,
        EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 58,
        EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
        EnvironmentVariableCategory.INTERNAL: 2,
    }
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
        if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
    )
    assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 34,
+        ModelEnvironmentVariableDisposition.PROMOTE: 35,
        ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
        ModelEnvironmentVariableDisposition.EXTERNAL: 0,
        ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
@@ -50,9 +50,6 @@ from tests.dfx.conftest import (
 )
 from tests.helpers.runtime import get_open_port

-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
-

 # ---------------------------------------------------------------------------
 # Inline field processing
@@ -837,7 +834,7 @@ def run_benchmark(
        "Quantization": _to_quantization_value(server_type, serve_args_dict),
        "offload": _to_offload_string(server_type, serve_args_dict),
        "compile": _to_compile_value(server_type, serve_args_dict),
-        "Attn_backend": os.environ.get("DIFFUSION_ATTENTION_BACKEND", ""),
+        "Attn_backend": str(serve_args_dict.get("diffusion-attention-backend") or ""),
        "num_inference_steps": params.get("num-inference-steps", ""),
        "completed": completed,
        "failed": failed,
@@ -1,7 +1,7 @@
 [
    {
        "test_name": "test_boogu_image_edit_single_device",
-        "description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. Text-only edit path (2 predictions/step); a double-guidance (guidance_scale_2) case is a follow-up since the random dataset cannot inject extra_body params. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.",
+        "description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. This row measures text-only edit guidance (2 predictions/step); dedicated rows below use extra-body to measure double guidance. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.",
        "server_type": "vllm-omni",
        "server_params": {
            "model": "Boogu/Boogu-Image-0.1-Edit",
@@ -33,5 +33,184 @@
                }
            }
        ]
+    },
+    {
+        "test_name": "test_boogu_image_edit_double_guidance_single_device",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 1
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Single-device sequential three-branch double-guidance control row for direct A/B comparison with cfg_parallel_size=2 and 3 on identical H100 hardware.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Edit",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_i2i_double_cfg1",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 5.0,
+                    "guidance_scale_2": 2.0,
+                    "seed": 42
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_boogu_image_edit_text_cfg_parallel_2",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 2
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Two-device CFG-parallel Edit text-only guidance measurement row. Compare with the single-device row on identical H100 hardware and software.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Edit",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_i2i_text_cfg2",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 5.0,
+                    "guidance_scale_2": 1.0,
+                    "seed": 42
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_boogu_image_edit_double_cfg_parallel_2",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 2
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Two-device round-robin execution of the three Edit double-guidance branches. Populate A/B latency, throughput, and per-GPU memory after GPU validation.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Edit",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_i2i_double_cfg2",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 5.0,
+                    "guidance_scale_2": 2.0,
+                    "seed": 42
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_boogu_image_edit_double_cfg_parallel_3",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 3
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Full three-device execution of the three Edit double-guidance branches. Compare directly with sequential and cfg_parallel_size=2 on identical H100 hardware.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Edit",
+            "serve_args": {
+                "cfg-parallel-size": 3,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_i2i_double_cfg3",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 5.0,
+                    "guidance_scale_2": 2.0,
+                    "seed": 42
+                }
+            }
+        ]
    }
 ]
@@ -33,5 +33,49 @@
                }
            }
        ]
+    },
+    {
+        "test_name": "test_boogu_image_cfg_parallel_2",
+        "mark": [
+            {
+                "hardware_marks": {
+                    "res": {
+                        "cuda": "H100"
+                    },
+                    "num_cards": 2
+                }
+            },
+            "full_model",
+            "diffusion"
+        ],
+        "description": "Two-device CFG-parallel Base T2I measurement row. Run beside test_boogu_image_single_device on the same H100 revision and populate latency, throughput, speedup, and per-GPU peak-memory baselines from the resulting report.",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Boogu/Boogu-Image-0.1-Base",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps28_cfg2",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 28,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "warmup-requests": 4,
+                "warmup-concurrency": 1,
+                "warmup-num-inference-steps": 28,
+                "enable-negative-prompt": true,
+                "extra-body": {
+                    "guidance_scale": 4.0,
+                    "seed": 42
+                }
+            }
+        ]
    }
 ]
@@ -30,7 +30,6 @@
        "vae-patch-parallel-size": 4,
        "vae-parallel-mode": "tile",
        "vae-use-tiling": true,
-        "diffusion-attention-backend": "FLASH_ATTN",
        "stage-init-timeout": 1800,
        "init-timeout": 1800,
        "enable-diffusion-pipeline-profiler": true
@@ -102,7 +101,6 @@
        "vae-patch-parallel-size": 4,
        "vae-parallel-mode": "tile",
        "vae-use-tiling": true,
-        "diffusion-attention-backend": "FLASH_ATTN",
        "stage-init-timeout": 1800,
        "init-timeout": 1800,
        "enable-diffusion-pipeline-profiler": true
@@ -175,7 +173,6 @@
        "vae-patch-parallel-size": 4,
        "vae-parallel-mode": "tile",
        "vae-use-tiling": true,
-        "diffusion-attention-backend": "FLASH_ATTN",
        "stage-init-timeout": 1800,
        "init-timeout": 1800,
        "enable-diffusion-pipeline-profiler": true
@@ -247,7 +244,6 @@
        "vae-patch-parallel-size": 4,
        "vae-parallel-mode": "tile",
        "vae-use-tiling": true,
-        "diffusion-attention-backend": "FLASH_ATTN",
        "stage-init-timeout": 1800,
        "init-timeout": 1800,
        "enable-diffusion-pipeline-profiler": true,
@@ -197,6 +197,7 @@ def test_create_unique_server_pytest_params_applies_marks(tmp_path):
    assert len(by_id["test_with_mark"].values) == 1
    assert isinstance(by_id["test_with_mark"].values[0], tuple)
    assert any(m.name == "H100" for m in by_id["test_with_mark"].marks)
+    assert not any(m.name == "B200" for m in by_id["test_with_mark"].marks)
    assert not any(m.name == "H100" for m in by_id["test_without_mark"].marks)


@@ -18,6 +18,7 @@ import torch

 import vllm_omni.diffusion.attention.layer as layer_mod
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.config import (
    get_current_diffusion_config,
@@ -73,6 +74,12 @@ class TestAttentionSpec:
        assert bk["dtype_vo"] == "fp8_e4m3"
        assert bk["flashinfer_backend"] == "trtllm-gen"

+    def test_quant_flashinfer_backend_serialized_without_dtype(self):
+        spec = AttentionSpec(backend="FLASHINFER_ATTN", quant={"flashinfer_backend": "fa2"})
+        assert spec.quant is not None
+        assert spec.quant.enabled is True
+        assert spec.backend_kwargs() == {"quant": {"flashinfer_backend": "fa2"}}
+
    def test_quant_and_skip_softmax_coexist(self):
        bk = AttentionSpec(
            backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5}
@@ -150,8 +157,8 @@ class TestAttentionSpec:
    def test_block_sparse_precision_enum_input_normalized(self):
        """RainFusionPrecision enum members must normalize to their string value.

-        Regression: the str,Enum members passed the value check but str() later
-        produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it).
+        Regression: ``str(str, Enum)`` later produced ``RainFusionPrecision.MIX``
+        instead of ``mix``. ``BlockSparseSpec`` stores ``member.value``.
        """
        from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig
        from vllm_omni.diffusion.data import RainFusionPrecision
@@ -509,6 +516,49 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
            dense_flash.assert_not_called()
            sdpa.assert_called_once()

+    @pytest.mark.parametrize("fa4_available", [False, True])
+    def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch, fa4_available):
+        attention = Attention.__new__(Attention)
+        attention._scheduler_paged_kv = True
+        attention.paged_kv_cache_role = "primary"
+        attention.backend_pref = "FLASH_ATTN"
+        attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
+        flash_output = torch.ones(1)
+        sdpa_output = torch.zeros(1)
+        dense_flash = Mock(return_value=flash_output)
+        sdpa = Mock(return_value=sdpa_output)
+        attention.attention = SimpleNamespace(forward=dense_flash)
+        attention.sdpa_fallback = SimpleNamespace(forward=sdpa)
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
+        monkeypatch.setattr(
+            layer_mod,
+            "get_forward_context",
+            lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True),
+        )
+        monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: False)
+        monkeypatch.setattr(layer_mod.current_omni_platform, "is_cuda", lambda: True)
+        monkeypatch.setattr(
+            layer_mod.current_omni_platform,
+            "supports_diffusion_dense_flash_attention",
+            lambda: fa4_available,
+        )
+
+        result = attention._run_local_attention(
+            torch.zeros(1, dtype=torch.bfloat16),
+            torch.zeros(1, dtype=torch.bfloat16),
+            torch.zeros(1, dtype=torch.bfloat16),
+            AttentionMetadata(),
+        )
+
+        if fa4_available:
+            assert result is flash_output
+            dense_flash.assert_called_once()
+            sdpa.assert_not_called()
+        else:
+            assert result is sdpa_output
+            dense_flash.assert_not_called()
+            sdpa.assert_called_once()
+
    def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
        attention = Attention.__new__(Attention)
        attention._scheduler_paged_kv = False
@@ -579,7 +629,106 @@ class TestAttentionInitUsesCurrentDiffusionConfig:

        assert len(calls) == 1
        assert attention.attn_backend is _DenseBackend
-        assert attention.backend_pref is None
+        # Auto path records the resolved backend name for Ring/SP; it is not explicit.
+        assert attention.backend_pref == _DenseBackend.get_name()
+        assert attention.backend_explicit is False
+
+    def test_internal_paged_selector_does_not_require_dense_fa4(self, monkeypatch):
+        """Blackwell CUDNN default + marked paged layer must not act as explicit FLASH_ATTN."""
+
+        class _FakeAttentionImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+            def forward(self, query, key, value, attn_metadata=None):
+                return query
+
+        class _CuDNNBackend:
+            supports_paged_kv = False
+
+            @staticmethod
+            def get_name() -> str:
+                return "CUDNN_ATTN"
+
+            @staticmethod
+            def get_impl_cls():
+                return _FakeAttentionImpl
+
+        def _fake_get_attn_backend_for_role(*args, attention_config=None, **kwargs):
+            if attention_config is not None:
+                spec, _ = attention_config.resolve_with_source(role="self", role_category=None)
+                if spec is not None and spec.backend.upper() == "FLASH_ATTN":
+                    raise ValueError(
+                        "FLASH_ATTN was explicitly selected but is unsupported "
+                        "(Blackwell requires CuTe FlashAttention-4 "
+                        "(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only). "
+                        "Select a compatible backend."
+                    )
+            return _CuDNNBackend, None
+
+        monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(FlashAttentionBackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.PAGED_SCHEDULER,
+            parallel_config=SimpleNamespace(ring_degree=1),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            attention = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+                paged_kv_cache_role="primary",
+            )
+
+        assert attention.attn_backend is FlashAttentionBackend
+        assert attention.backend_pref == "FLASH_ATTN"
+        assert attention.backend_explicit is False
+        assert attention.attn_spec is None
+
+    def test_user_explicit_flash_attn_still_requires_dense_fa4(self, monkeypatch):
+        def _fake_get_attn_backend_for_role(*args, attention_config=None, **kwargs):
+            if attention_config is not None:
+                spec, _ = attention_config.resolve_with_source(role="self", role_category=None)
+                if spec is not None and spec.backend.upper() == "FLASH_ATTN":
+                    raise ValueError(
+                        "FLASH_ATTN was explicitly selected but is unsupported "
+                        "(Blackwell requires CuTe FlashAttention-4 "
+                        "(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only). "
+                        "Select a compatible backend."
+                    )
+            raise AssertionError("explicit FLASH_ATTN must not fall through to the platform default")
+
+        monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")),
+            diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.PAGED_SCHEDULER,
+            parallel_config=SimpleNamespace(ring_degree=1),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config), pytest.raises(ValueError, match="explicitly selected"):
+            Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+                paged_kv_cache_role="primary",
+            )

    def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
        class _FakeAttentionImpl:
@@ -614,9 +763,10 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
            return _FakeBackend, AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})

        class _FakeRingParallelAttention:
-            def __init__(self, sp_group, attn_backend_pref=None):
+            def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False):
                self.sp_group = sp_group
                self.attn_backend_pref = attn_backend_pref
+                self.attn_backend_explicit = attn_backend_explicit

        monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
@@ -662,11 +812,134 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
        assert attn.use_ring is True
        assert attn.ring_runner is not None
        assert attn.ring_runner.attn_backend_pref == "TRTLLM_ATTN"
+        assert attn.ring_runner.attn_backend_explicit is True
+
+    def test_attention_init_forwards_flashinfer_backend_without_dtype(self, monkeypatch):
+        class _FakeAttentionImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+        class _FakeBackend:
+            @staticmethod
+            def get_name() -> str:
+                return "FLASHINFER_ATTN"
+
+            @staticmethod
+            def get_impl_cls():
+                return _FakeAttentionImpl
+
+        spec = AttentionSpec(backend="FLASHINFER_ATTN", quant={"flashinfer_backend": "fa2"})
+
+        monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", lambda **kwargs: (_FakeBackend, spec))
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+        monkeypatch.setattr(layer_mod, "get_sp_group", lambda: None)
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(default=spec),
+            parallel_config=SimpleNamespace(ring_degree=1),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            attn = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+                role="self",
+            )
+
+        assert attn.backend_pref == "FLASHINFER_ATTN"
+        assert attn.attention.kwargs["backend_kwargs"] == {"quant": {"flashinfer_backend": "fa2"}}
+        assert attn.attention.kwargs["backend_explicit"] is True
+
+    def test_attention_init_propagates_ring_setup_failure(self, monkeypatch):
+        class _FakeAttentionImpl:
+            def __init__(self, **kwargs):
+                pass
+
+        class _FakeBackend:
+            @staticmethod
+            def get_name() -> str:
+                return "TORCH_SDPA"
+
+            @staticmethod
+            def get_impl_cls():
+                return _FakeAttentionImpl
+
+        monkeypatch.setattr(
+            layer_mod,
+            "get_attn_backend_for_role",
+            lambda **kwargs: (_FakeBackend, AttentionSpec(backend="TORCH_SDPA")),
+        )
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(
+            layer_mod,
+            "get_sp_group",
+            lambda: (_ for _ in ()).throw(RuntimeError("SP group is not initialized")),
+        )
+
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
+            parallel_config=SimpleNamespace(ring_degree=2),
+        )
+        with set_current_diffusion_config(od_config):
+            with pytest.raises(RuntimeError, match="SP group is not initialized"):
+                Attention(num_heads=4, head_size=64, causal=False, softmax_scale=1.0)
+
+    def test_float32_main_dispatch_does_not_override_selected_backend(self):
+        sentinel = torch.randn(1, 2, 4, 8)
+        selected_calls = []
+
+        def _selected_forward(*args):
+            selected_calls.append(args)
+            return sentinel
+
+        fake_attention = SimpleNamespace(
+            attention=SimpleNamespace(
+                forward=_selected_forward,
+            ),
+            sdpa_fallback=SimpleNamespace(
+                forward=lambda *args: pytest.fail("unexpected SDPA fallback"),
+            ),
+            _assert_metadata_compatible=lambda metadata: None,
+            _scheduler_paged_kv=False,
+            paged_kv_cache_role=None,
+        )
+        query = torch.randn(1, 2, 4, 8, dtype=torch.float32)
+
+        output = Attention._run_local_attention(fake_attention, query, query, query, None)
+
+        assert output is sentinel
+        assert len(selected_calls) == 1
+
+    def test_unsupported_attention_mask_is_rejected_before_backend_call(self):
+        fake_backend = SimpleNamespace(
+            get_name=lambda: "NO_MASK_BACKEND",
+            supports_attention_mask=lambda *args, **kwargs: False,
+        )
+        fake_attention = SimpleNamespace(attn_backend=fake_backend)
+        metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool))
+
+        with pytest.raises(ValueError, match="does not support attn_mask"):
+            Attention._assert_metadata_compatible(fake_attention, metadata)
+
+    def test_ring_attention_rejects_mask_instead_of_ignoring_it(self):
+        fake_attention = SimpleNamespace(attention=SimpleNamespace(skip=None), ring_runner=None)
+        query = torch.randn(1, 2, 4, 8)
+        metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool))
+
+        with pytest.raises(ValueError, match="Ring attention does not support attn_mask"):
+            Attention._run_ring_attention(fake_attention, query, query, query, metadata)


 class TestDiffusionKvCacheQuantization:
    @staticmethod
-    def _install_attention_init_stubs(monkeypatch):
+    def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True):
        class _FakeAttentionImpl:
            def __init__(self, **kwargs):
                self.kwargs = kwargs
@@ -674,6 +947,10 @@ class TestDiffusionKvCacheQuantization:
            def forward(self, query, key, value, attn_metadata=None):
                return query

+            @staticmethod
+            def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
+                return supports_kv_cache_dtype
+
        class _FakeBackend:
            @staticmethod
            def get_name() -> str:
@@ -683,14 +960,11 @@ class TestDiffusionKvCacheQuantization:
            def get_impl_cls():
                return _FakeAttentionImpl

-            @staticmethod
-            def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
-                return True
-
        class _FakeRingParallelAttention:
-            def __init__(self, sp_group, attn_backend_pref=None):
+            def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False):
                self.sp_group = sp_group
                self.attn_backend_pref = attn_backend_pref
+                self.attn_backend_explicit = attn_backend_explicit

        monkeypatch.setattr(
            layer_mod,
@@ -764,3 +1038,22 @@ class TestDiffusionKvCacheQuantization:
                    causal=False,
                    softmax_scale=1.0,
                )
+
+    def test_explicit_unsupported_kv_cache_dtype_raises(self, monkeypatch):
+        self._install_attention_init_stubs(monkeypatch, supports_kv_cache_dtype=False)
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            parallel_config=SimpleNamespace(ring_degree=1),
+            diffusion_kv_cache_dtype="fp8",
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            with pytest.raises(ValueError, match="does not support kv_cache_dtype='fp8'"):
+                Attention(
+                    num_heads=4,
+                    head_size=64,
+                    causal=False,
+                    softmax_scale=1.0,
+                )
@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import sys
+import types
+from types import SimpleNamespace
+
+import pytest
+import torch
+from vllm.platforms import current_platform
+from vllm.platforms.interface import DeviceCapability
+
+from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.attention.selector import _cached_get_backend_cls
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import AttentionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+
+# Importing CudaOmniPlatform pulls in vllm.platforms.cuda, which imports the
+# CUDA-only ``vllm._C_stable_libtorch`` extension at module top level. That
+# extension is absent on non-CUDA builds (e.g. Intel/XPU), so skip the whole
+# module there before the import can crash collection. ``is_cuda()`` resolves
+# the platform without importing cuda.py.
+if not current_platform.is_cuda():
+    pytest.skip("CUDA-only diffusion backend selection tests", allow_module_level=True)
+
+from vllm_omni.platforms.cuda.platform import CudaOmniPlatform  # noqa: E402
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _blackwell_sm120(monkeypatch: pytest.MonkeyPatch, *, cudnn_version: int = 90500) -> None:
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(12, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+    monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: False))
+    monkeypatch.setattr(torch.backends.cudnn, "version", lambda: cudnn_version)
+
+
+def _install_dummy_flashinfer(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Top-level ``flashinfer`` only — no ``prefill`` wrapper (partial install)."""
+    monkeypatch.setitem(sys.modules, "flashinfer", types.ModuleType("flashinfer"))
+    # A previously imported real submodule must not satisfy the wrapper probe.
+    monkeypatch.delitem(sys.modules, "flashinfer.prefill", raising=False)
+
+
+def _install_dummy_flashinfer_prefill(monkeypatch: pytest.MonkeyPatch) -> None:
+    """FlashInfer stub that exposes the prefill wrapper FLASHINFER_ATTN needs."""
+    prefill = types.ModuleType("flashinfer.prefill")
+    setattr(prefill, "BatchPrefillWithRaggedKVCacheWrapper", type("BatchPrefillWithRaggedKVCacheWrapper", (), {}))
+    pkg = types.ModuleType("flashinfer")
+    setattr(pkg, "prefill", prefill)
+    monkeypatch.setitem(sys.modules, "flashinfer", pkg)
+    monkeypatch.setitem(sys.modules, "flashinfer.prefill", prefill)
+
+
+def test_auto_selects_cudnn_for_supported_blackwell_head_size(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=128)
+
+    assert path == DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
+
+
+def test_auto_routes_incompatible_head_size_to_sdpa(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+    _install_dummy_flashinfer(monkeypatch)
+
+    # 320 is outside cuDNN FMHA (<=256, multiple of 8) and FlashInfer {64,128,256}.
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=320)
+
+    assert path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+
+def test_auto_routes_non_multiple_head_size_to_sdpa_without_flashinfer(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+    monkeypatch.setitem(sys.modules, "flashinfer", None)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=12)
+
+    assert path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+
+def test_auto_selects_flashinfer_when_cudnn_too_old(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch, cudnn_version=90499)
+    _install_dummy_flashinfer_prefill(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=128)
+
+    assert path == DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
+
+
+def test_auto_skips_flashinfer_when_only_toplevel_package(monkeypatch: pytest.MonkeyPatch):
+    """A stub that only imports ``flashinfer`` must not be auto-selected."""
+    _blackwell_sm120(monkeypatch, cudnn_version=90499)
+    _install_dummy_flashinfer(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=128)
+
+    assert path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+
+def test_explicit_flashinfer_raises_without_prefill_wrapper(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+    _install_dummy_flashinfer(monkeypatch)
+
+    with pytest.raises(ValueError, match="BatchPrefillWithRaggedKVCacheWrapper"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASHINFER_ATTN", head_size=128)
+
+
+def test_auto_skips_flashinfer_for_unsupported_head_size_when_cudnn_too_old(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch, cudnn_version=90499)
+    _install_dummy_flashinfer(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=72)
+
+    assert path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+
+def test_explicit_cudnn_raises_for_incompatible_head_size(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+
+    with pytest.raises(ValueError, match="head_size=12 is unsupported"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("CUDNN_ATTN", head_size=12)
+
+
+def test_explicit_cudnn_accepts_supported_head_size(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls("CUDNN_ATTN", head_size=72)
+
+    assert path == DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
+
+
+def test_explicit_cudnn_skips_head_size_check_for_capability_probe(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls("CUDNN_ATTN", head_size=-1)
+
+    assert path == DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
+
+
+@pytest.mark.parametrize("head_size", [32, 320])
+def test_explicit_flashinfer_raises_for_unsupported_head_size(monkeypatch: pytest.MonkeyPatch, head_size: int):
+    _blackwell_sm120(monkeypatch)
+    _install_dummy_flashinfer_prefill(monkeypatch)
+
+    with pytest.raises(ValueError, match=f"head_size={head_size} is unsupported"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASHINFER_ATTN", head_size=head_size)
+
+
+def test_explicit_flashinfer_accepts_supported_head_size(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+    _install_dummy_flashinfer_prefill(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASHINFER_ATTN", head_size=128)
+
+    assert path == DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
+
+
+def test_explicit_flashinfer_skips_head_size_check_for_capability_probe(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+    _install_dummy_flashinfer_prefill(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASHINFER_ATTN", head_size=-1)
+
+    assert path == DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
+
+
+def test_auto_selects_cudnn_for_unknown_head_size_probe(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+    _install_dummy_flashinfer_prefill(monkeypatch)
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=-1)
+
+    assert path == DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
+
+
+def test_explicit_flash_attn_rejects_blackwell_without_fa4(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
+
+    with pytest.raises(ValueError, match="requires CuTe FlashAttention-4"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
+
+
+def test_explicit_flash_attn_accepts_blackwell_when_fa4_available(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+    monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: True))
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
+
+    assert path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+
+def test_blackwell_dense_flash_capability_tracks_fa4(monkeypatch: pytest.MonkeyPatch):
+    _blackwell_sm120(monkeypatch)
+
+    assert CudaOmniPlatform.supports_diffusion_dense_flash_attention() is False
+
+    monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: True))
+
+    assert CudaOmniPlatform.supports_diffusion_dense_flash_attention() is True
+
+
+def test_hopper_dense_flash_capability_tracks_fa_package(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
+    )
+    monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_package", classmethod(lambda cls: False))
+
+    assert CudaOmniPlatform.supports_diffusion_dense_flash_attention() is False
+
+    monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_package", classmethod(lambda cls: True))
+
+    assert CudaOmniPlatform.supports_diffusion_dense_flash_attention() is True
+
+
+def test_auto_selects_flash_attn_on_blackwell_when_cudnn_old_and_fa4_available(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    _blackwell_sm120(monkeypatch, cudnn_version=90499)
+    monkeypatch.setitem(sys.modules, "flashinfer", None)
+    monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: True))
+
+    path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=128)
+
+    assert path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
+
+def test_marked_paged_layer_constructs_on_blackwell_cudnn_without_fa4(monkeypatch: pytest.MonkeyPatch):
+    """Scheduler-paged remap must not require dense FA4 when CUDNN is the default."""
+    _blackwell_sm120(monkeypatch)
+    _cached_get_backend_cls.cache_clear()
+
+    od_config = SimpleNamespace(
+        diffusion_attention_config=AttentionConfig(),
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        parallel_config=SimpleNamespace(ring_degree=1, allgather_degree=1),
+        diffusion_kv_cache_dtype=None,
+        diffusion_kv_cache_skip_step_indices=None,
+        diffusion_kv_cache_skip_layer_indices=None,
+        model_class_name=None,
+    )
+
+    with set_current_diffusion_config(od_config):
+        attention = Attention(
+            num_heads=4,
+            head_size=64,
+            causal=False,
+            softmax_scale=1.0,
+            paged_kv_cache_role="primary",
+        )
+
+    assert attention.attn_backend.get_name() == "FLASH_ATTN"
+    assert attention.backend_pref == "FLASH_ATTN"
+    assert attention.backend_explicit is False
+    assert attention.attn_spec is None
@@ -1,62 +1,84 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
+from typing import Any

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

+import vllm_omni.diffusion.attention.backends.cudnn_attn as cudnn_backend
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
+from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend, CuDNNAttentionImpl

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


-@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
-):
+def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
+    """Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V."""
    selected_backends = []

+    @contextmanager
    def fake_sdpa_kernel(backends):
-        selected_backends.append(backends)
-        return nullcontext()
+        selected_backends.append(tuple(backends))
+        yield

-    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)
+    def fake_sdpa(query, key, value, **kwargs):
+        return query

-    output = impl.forward_cuda(query, key, value)
+    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel)
+    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
+
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    query = torch.randn(1, 2, 2, 8)
+    singleton_kv = torch.randn(1, 1, 2, 8)
+
+    output = impl.forward_cuda(query, singleton_kv, singleton_kv)

    assert output.shape == query.shape
-    assert selected_backends == [[expected_backend]]
+    assert selected_backends == [(SDPBackend.MATH,)]
+
+
+def test_explicit_cudnn_rejects_kv_seq_len_one():
+    impl = CuDNNAttentionImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=0.5,
+        backend_explicit=True,
+    )
+    query = torch.randn(1, 2, 2, 8)
+    singleton_kv = torch.randn(1, 1, 2, 8)
+
+    with pytest.raises(ValueError, match="explicitly selected.*sequence length 1"):
+        impl.forward_cuda(query, singleton_kv, singleton_kv)
+
+
+def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
+    selected_backends = []
+
+    @contextmanager
+    def fake_sdpa_kernel(backends):
+        selected_backends.append(tuple(backends))
+        yield
+
+    def reject_shape(*args, **kwargs):
+        raise RuntimeError("No available kernel. Aborting execution.")
+
+    monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel)
+    monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", reject_shape)
+
+    impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
+    tensors = torch.randn(1, 2, 2, 8)
+
+    with pytest.raises(RuntimeError, match="No available kernel"):
+        impl.forward_cuda(tensors, tensors, tensors)
+
+    assert selected_backends == [(SDPBackend.CUDNN_ATTENTION,)]


 def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
-    observed = {}
+    observed: dict[str, Any] = {}

    def fake_sdpa(query, key, value, **kwargs):
        observed.update(query=query, key=key, value=value, kwargs=kwargs)
@@ -105,3 +127,15 @@ def test_cudnn_rejects_invalid_valid_kv_length():
            query,
            AttentionMetadata(extra={"valid_kv_length": 9}),
        )
+
+
+@pytest.mark.parametrize("head_size", [8, 64, 128, 256])
+def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size):
+    assert CuDNNAttentionBackend.supports_head_size(head_size)
+    assert head_size in CuDNNAttentionBackend.get_supported_head_sizes()
+
+
+@pytest.mark.parametrize("head_size", [0, 7, 12, 320])
+def test_cudnn_backend_rejects_incompatible_head_sizes(head_size):
+    assert not CuDNNAttentionBackend.supports_head_size(head_size)
+    assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes()
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

+import math
 import sys
 import types
 from typing import Any
@@ -8,9 +9,17 @@ from typing import Any
 import pytest
 import torch

-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+    VideoTokenLayout,
+    VideoTokenSpan,
+)
 from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
+    FastVideoVSABackend,
    FastVideoVSAImpl,
+    _build_h3_block_map,
+    _get_h3_tile_metadata,
 )
 from vllm_omni.diffusion.attention.backends.registry import (
    DiffusionAttentionBackendEnum,
@@ -23,6 +32,12 @@ def test_fastvideo_vsa_backend_is_registered():
    assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")


+def test_fastvideo_vsa_reports_missing_optional_kernel(monkeypatch):
+    monkeypatch.setattr("importlib.util.find_spec", lambda name: None)
+    with pytest.raises(ImportError, match="fastvideo-kernel"):
+        FastVideoVSABackend.validate_available()
+
+
 def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
    calls: dict[str, Any] = {}
    fake_module = types.ModuleType("fastvideo_kernel")
@@ -201,3 +216,190 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
    # CPU/float32 is rejected later, but k=N itself must not trigger fallback.
    assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
    assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
+
+
+def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
+    extra: dict[str, Any] = {
+        "vsa_h3_prefix_segments": prefix_segments,
+    }
+    if gate is not None:
+        extra["gate_compress"] = gate
+    return AttentionMetadata(
+        packed_padding=packed_padding,
+        video_layout=VideoTokenLayout(
+            used_len=sum(prefix_segments) + math.prod(video_shape),
+            video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
+        ),
+        extra=extra,
+    )
+
+
+def _fake_block_sparse_kernel(monkeypatch, calls):
+    """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
+
+    def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
+        calls["block_map"] = block_map.clone()
+        calls["variable_block_sizes"] = variable_block_sizes.tolist()
+        calls["q_shape"] = tuple(q.shape)
+        blocks = block_map.shape[-1]
+        assert blocks * 64 == q.shape[2]
+        # Expand the per-block map to a per-token mask so the stand-in honours
+        # the sparsity contract the real kernel implements, and drop the rows
+        # each partial block pads with, which the real kernel skips through
+        # variable_block_sizes.
+        token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
+        within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
+        token_mask = token_mask & within.reshape(1, 1, 1, -1)
+        out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
+        return out.to(q.dtype), None
+
+    module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
+    setattr(module, "block_sparse_attn", block_sparse_attn)
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
+    monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
+
+
+def _h3_impl(**backend_kwargs):
+    kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
+    kwargs.update(backend_kwargs)
+    return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
+
+
+def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5, 70), (2, 4, 4)
+    seq_len = sum(prefix_segments) + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    assert output.shape == query.shape
+    # topk covers every video block, so the block map is dense and the tiled
+    # kernel must reproduce plain attention over the same rows.
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
+    # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
+    assert calls["variable_block_sizes"] == [5, 64, 6, 32]
+
+
+def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
+    calls: dict[str, Any] = {}
+    _fake_block_sparse_kernel(monkeypatch, calls)
+    impl = _h3_impl(topk=1)
+    prefix_segments, video_shape = (5,), (8, 4, 4)
+    seq_len = 5 + 8 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+
+    # 1 prefix block + 2 video blocks. The odd count is padded to an even one
+    # for the paired-CTA kernel contract, and the Triton route slices that
+    # transport-only partner off again before launching.
+    block_map = calls["block_map"]
+    assert block_map.shape[-1] == 3
+    assert calls["variable_block_sizes"] == [5, 64, 64]
+    assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
+    assert block_map[..., :1].all(), "prefix keys are exempt from selection"
+    assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
+
+    # The same odd geometry must still untile back onto the original rows.
+    dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        *(x.float().transpose(1, 2) for x in (query, query, query))
+    ).transpose(1, 2)
+    torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
+
+
+def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    seq_len = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
+
+    without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
+    zero_gate = torch.zeros_like(query)
+    with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
+    live_gate = torch.full_like(query, 0.5)
+    with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
+
+    torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
+    assert not torch.allclose(with_live_gate.float(), without_gate.float())
+
+
+def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
+    _fake_block_sparse_kernel(monkeypatch, {})
+    impl = _h3_impl(topk=64)
+    prefix_segments, video_shape = (5,), (2, 4, 4)
+    valid = 5 + 2 * 4 * 4
+    torch.manual_seed(0)
+    query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
+    padding = PackedPaddingMetadata(
+        q_length=valid,
+        kv_length=valid,
+        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
+        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
+    )
+
+    output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
+
+    assert output.shape == query.shape
+    assert torch.count_nonzero(output[:, valid:]) == 0
+
+
+def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
+    # The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
+    # building the padding mask. Every fallback out of the VSA path must honour
+    # that contract itself; SDPA reads attn_mask and nothing else.
+    impl = FastVideoVSAImpl(
+        num_heads=2,
+        head_size=8,
+        softmax_scale=8**-0.5,
+        backend_kwargs={"topk": 1, "min_seq_len": 4096},
+    )
+    valid = 40
+    torch.manual_seed(0)
+    query = torch.randn(1, 64, 2, 8)
+    padding = PackedPaddingMetadata(
+        q_length=valid,
+        kv_length=valid,
+        cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
+        cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
+    )
+
+    baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
+    perturbed_input = query.clone()
+    perturbed_input[:, valid:] += 100.0
+    perturbed = impl.forward_cuda(
+        perturbed_input, perturbed_input, perturbed_input, AttentionMetadata(packed_padding=padding, extra={})
+    )
+
+    assert baseline.shape == query.shape
+    torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
+    assert torch.count_nonzero(baseline[:, valid:]) == 0
+
+
+def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
+    partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
+        (5, 70, 9), (5, 6, 6), torch.device("cpu")
+    )
+    assert sizes[:4].tolist() == [5, 64, 6, 9]
+    assert prefix_blocks == 4
+    assert video_blocks == 8
+    assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
+    assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
+
+
+def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
+    scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
+    block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
+    assert block_map[:, :, :2].all()
+    assert block_map[..., :2].all()
+    assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends import flashinfer_attn
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.flashinfer_attn import (
+    FlashInferAttentionBackend,
+    FlashInferAttentionImpl,
+)
+from vllm_omni.diffusion.data import AttentionSpec, AttnQuantSpec
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _impl(*, causal: bool = False, backend_explicit: bool = False):
+    # Avoid CUDA/wrapper init; these tests only cover mask validation.
+    obj = FlashInferAttentionImpl.__new__(FlashInferAttentionImpl)
+    obj.causal = causal
+    obj.softmax_scale = 0.5
+    obj.flashinfer_backend = "fa2"
+    obj.backend_explicit = backend_explicit
+    obj._sdpa_fallback = None
+    return obj
+
+
+def test_flashinfer_rejects_float_mask_instead_of_falling_back(monkeypatch):
+    monkeypatch.setattr(flashinfer_attn, "HAS_FLASHINFER", True)
+    query = torch.randn(1, 2, 2, 8)
+    metadata = AttentionMetadata(attn_mask=torch.zeros(2, 2))
+
+    with pytest.raises(ValueError, match="boolean-only"):
+        _impl(backend_explicit=True).forward_cuda(query, query, query, metadata)
+
+
+def test_flashinfer_rejects_causal_custom_mask_instead_of_falling_back(monkeypatch):
+    monkeypatch.setattr(flashinfer_attn, "HAS_FLASHINFER", True)
+    query = torch.randn(1, 2, 2, 8)
+    metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False], [True, True]]))
+
+    with pytest.raises(ValueError, match="causal=True"):
+        _impl(causal=True, backend_explicit=True).forward_cuda(query, query, query, metadata)
+
+
+def test_explicit_cute_dsl_rejects_custom_mask_instead_of_falling_back(monkeypatch):
+    monkeypatch.setattr(flashinfer_attn, "HAS_FLASHINFER", True)
+    impl = _impl(backend_explicit=True)
+    impl.flashinfer_backend = "cute-dsl"
+    query = torch.randn(1, 2, 2, 8)
+    metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False], [True, True]]))
+
+    with pytest.raises(ValueError, match="cute-dsl"):
+        impl.forward_cuda(query, query, query, metadata)
+
+
+def test_auto_cute_dsl_falls_back_to_sdpa_for_custom_mask(monkeypatch):
+    monkeypatch.setattr(flashinfer_attn, "HAS_FLASHINFER", True)
+    impl = _impl(backend_explicit=False)
+    impl.flashinfer_backend = "cute-dsl"
+    impl._run_batch_prefill = lambda *args, **kwargs: pytest.fail("unexpected cute-dsl prefill")
+    query = torch.randn(1, 2, 2, 8)
+    metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False], [True, True]]))
+
+    output = impl.forward_cuda(query, query, query, metadata)
+
+    assert output.shape == query.shape
+
+
+def test_explicit_cute_dsl_is_not_mask_capable(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
+    spec = AttentionSpec(backend="FLASHINFER_ATTN")
+
+    assert FlashInferAttentionBackend.supports_attention_mask() is True
+    assert FlashInferAttentionBackend.supports_attention_mask(spec) is False
+
+
+def test_explicit_fa2_flashinfer_is_mask_capable(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
+    spec = AttentionSpec(backend="FLASHINFER_ATTN", quant=AttnQuantSpec(flashinfer_backend="fa2"))
+
+    assert FlashInferAttentionBackend.supports_attention_mask(spec) is True
+
+
+def test_explicit_fa2_backend_kwargs_select_fa2_on_blackwell(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
+    spec = AttentionSpec(backend="FLASHINFER_ATTN", quant=AttnQuantSpec(flashinfer_backend="fa2"))
+    kwargs = spec.backend_kwargs() or {}
+    requested = (kwargs.get("quant") or {}).get("flashinfer_backend", "auto")
+
+    assert kwargs == {"quant": {"flashinfer_backend": "fa2"}}
+    assert FlashInferAttentionImpl._select_backend(requested) == "fa2"
+    assert FlashInferAttentionBackend.supports_attention_mask(spec) is True
+
+
+def test_auto_flashinfer_backend_kwargs_select_cute_dsl_on_blackwell(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
+    spec = AttentionSpec(backend="FLASHINFER_ATTN")
+    kwargs = spec.backend_kwargs() or {}
+    requested = (kwargs.get("quant") or {}).get("flashinfer_backend", "auto")
+
+    assert "quant" not in kwargs
+    assert FlashInferAttentionImpl._select_backend(requested) == "cute-dsl"
+    assert FlashInferAttentionBackend.supports_attention_mask(spec) is False
+
+
+def test_hopper_explicit_flashinfer_is_mask_capable(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (9, 0))
+    spec = AttentionSpec(backend="FLASHINFER_ATTN")
+
+    assert FlashInferAttentionBackend.supports_attention_mask(spec) is True
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

 import sys
 import types
@@ -9,18 +9,13 @@ import torch

 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.platforms import current_omni_platform


@pytest.mark.core_model
@pytest.mark.cpu
-def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
-    """Test that when kernels package is not available, the platform fallback logic
-
-    routes to local FLASH_ATTN.
-    """
+def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform

    # Temporarily hide/remove kernels module if present
@@ -38,19 +33,73 @@ def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
    )
    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})

-    # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+    with pytest.raises(ImportError, match="explicitly selected"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
+    with pytest.raises(ImportError, match="explicitly selected"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)

-    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
-
-    # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs
    kernels_module = types.ModuleType("kernels")
    monkeypatch.setitem(sys.modules, "kernels", kernels_module)
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
+    with pytest.raises(ValueError, match="require.*Hopper GPU"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
+    )
+    with pytest.raises(ValueError, match="current FA2 kernels require"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
+    with pytest.raises(ValueError, match="require.*Hopper GPU"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_explicit_flash_attention_unavailable_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+
+    with pytest.raises(ValueError, match="explicitly selected but is unsupported"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
+    monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: False))
+    with pytest.raises(ValueError, match="requires CuTe FlashAttention-4"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_explicit_sage_attention_unsupported_arch_does_not_reach_kernel(monkeypatch: pytest.MonkeyPatch):
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+
+    with pytest.raises(ValueError, match="does not provide a kernel for sm_103"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN", head_size=64)


@pytest.mark.core_model
@@ -0,0 +1,199 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import importlib
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.ring import ring_globals, ring_kernels
+from vllm_omni.diffusion.attention.parallel import factory, ring
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_source_build_fa3_is_hopper_only(monkeypatch):
+    fake_interface = SimpleNamespace(
+        _flash_attn_forward=lambda *args, **kwargs: None,
+        flash_attn_func=lambda *args, **kwargs: None,
+    )
+    try:
+        with monkeypatch.context() as context:
+            context.setitem(sys.modules, "flash_attn_interface", fake_interface)
+            reloaded = importlib.reload(ring_globals)
+            assert reloaded.HAS_FA3
+            assert reloaded.FA3_SUPPORTED_CUDA_MAJORS == frozenset({9})
+    finally:
+        importlib.reload(ring_globals)
+
+
+@pytest.mark.parametrize(
+    ("capability", "expected"),
+    [
+        ((8, 0), True),
+        ((8, 9), True),
+        ((9, 0), True),
+        ((10, 0), False),
+        ((10, 3), False),
+        ((12, 0), False),
+    ],
+)
+def test_can_use_fa3_checks_device_arch(monkeypatch, capability, expected):
+    monkeypatch.setattr(ring, "HAS_FA3", True)
+    monkeypatch.setattr(ring, "FA3_SUPPORTED_CUDA_MAJORS", frozenset({8, 9}))
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: capability)
+
+    assert ring._can_use_fa3(torch.device("cuda")) is expected
+
+
+def test_can_use_fa3_preserves_unknown_extension_contract(monkeypatch):
+    monkeypatch.setattr(ring, "HAS_FA3", True)
+    monkeypatch.setattr(ring, "FA3_SUPPORTED_CUDA_MAJORS", None)
+
+    assert ring._can_use_fa3(torch.device("cuda"))
+
+
+@pytest.mark.parametrize(
+    ("capability", "expected"),
+    [
+        ((8, 0), True),
+        ((9, 0), True),
+        ((10, 0), False),
+        ((10, 3), False),
+        ((12, 0), False),
+    ],
+)
+def test_can_use_fa2_checks_device_arch(monkeypatch, capability, expected):
+    monkeypatch.setattr(ring, "HAS_FLASH_ATTN", True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: capability)
+
+    assert ring._can_use_fa2(torch.device("cuda")) is expected
+
+
+def test_ring_falls_back_to_sdpa_when_fa3_extension_has_no_device_kernel(monkeypatch):
+    monkeypatch.setattr(ring, "HAS_FA3", True)
+    monkeypatch.setattr(ring, "HAS_FA4", False)
+    monkeypatch.setattr(ring, "HAS_FLASH_ATTN", False)
+    monkeypatch.setattr(ring, "HAS_AITER", False)
+    monkeypatch.setattr(ring, "_can_use_fa3", lambda _device: False)
+
+    expected = torch.randn(1, 2, 1, 8)
+
+    def fake_ring_sdpa(*args, **kwargs):
+        return expected
+
+    module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
+    monkeypatch.setitem(sys.modules, module_name, SimpleNamespace(ring_pytorch_attn_func=fake_ring_sdpa))
+
+    strategy = ring.RingParallelAttention(SimpleNamespace(ring_group=object()))
+    query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
+    actual = strategy.run_attention(query, query, query, attn_metadata=None)
+
+    assert actual is expected
+
+
+def test_explicit_ring_backend_does_not_fallback_to_sdpa(monkeypatch):
+    monkeypatch.setattr(ring, "HAS_FA3", True)
+    monkeypatch.setattr(ring, "HAS_FA4", False)
+    monkeypatch.setattr(ring, "HAS_FLASH_ATTN", False)
+    monkeypatch.setattr(ring, "HAS_AITER", False)
+    monkeypatch.setattr(ring, "_can_use_fa3", lambda _device: False)
+
+    strategy = ring.RingParallelAttention(
+        SimpleNamespace(ring_group=object()),
+        attn_backend_pref="FLASH_ATTN",
+        attn_backend_explicit=True,
+    )
+    query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
+
+    with pytest.raises(RuntimeError, match="explicitly selected"):
+        strategy.run_attention(query, query, query, attn_metadata=None)
+
+
+@pytest.mark.parametrize(
+    "backend_pref",
+    [
+        "CUDNN_ATTN",
+        "FLASH_ATTN_HUB",
+        "FLASH_ATTN_3_HUB",
+    ],
+)
+def test_explicit_non_ring_backend_is_rejected(monkeypatch, backend_pref):
+    # Local FA is available so a missing rejection would silently run FA4/FA3/FA2
+    # instead of the requested Hub/cuDNN backend.
+    monkeypatch.setattr(ring, "HAS_FA4", True)
+    monkeypatch.setattr(ring, "_can_use_fa4", lambda _device: True)
+    module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
+    monkeypatch.setitem(
+        sys.modules,
+        module_name,
+        SimpleNamespace(ring_pytorch_attn_func=lambda *args, **kwargs: pytest.fail("unexpected SDPA fallback")),
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.diffusion.attention.backends.ring_flash_attn",
+        SimpleNamespace(ring_flash_attn_func=lambda *args, **kwargs: pytest.fail("unexpected local FA ring")),
+    )
+    strategy = ring.RingParallelAttention(
+        SimpleNamespace(ring_group=object()),
+        attn_backend_pref=backend_pref,
+        attn_backend_explicit=True,
+    )
+    query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
+
+    with pytest.raises(ValueError, match="ring sequence parallelism has no implementation"):
+        strategy.run_attention(query, query, query, attn_metadata=None)
+
+
+def test_explicit_flash_ring_rejects_float32_instead_of_falling_back(monkeypatch):
+    module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
+    monkeypatch.setitem(
+        sys.modules,
+        module_name,
+        SimpleNamespace(ring_pytorch_attn_func=lambda *args, **kwargs: pytest.fail("unexpected SDPA fallback")),
+    )
+    strategy = ring.RingParallelAttention(
+        SimpleNamespace(ring_group=object()),
+        attn_backend_pref="FLASH_ATTN",
+        attn_backend_explicit=True,
+    )
+    query = torch.randn(1, 2, 1, 8, dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="does not support float32"):
+        strategy.run_attention(query, query, query, attn_metadata=None)
+
+
+def test_pytorch_ring_flash_op_rejects_float32():
+    query = torch.randn(1, 2, 1, 8, dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="does not support float32"):
+        ring_kernels.pytorch_attn_forward(query, query, query, op_type="flash")
+
+
+def test_configured_sp_does_not_fallback_when_group_is_unavailable(monkeypatch):
+    parallel_config = SimpleNamespace(ulysses_degree=2, ring_degree=1, allgather_degree=1)
+    forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(parallel_config=parallel_config))
+    monkeypatch.setattr(factory, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(factory, "get_forward_context", lambda: forward_context)
+    monkeypatch.setattr(
+        factory,
+        "get_sp_group",
+        lambda: (_ for _ in ()).throw(RuntimeError("SP group is not initialized")),
+    )
+
+    with pytest.raises(RuntimeError, match="SP is configured"):
+        factory.build_parallel_attention_strategy(scatter_idx=2, gather_idx=1, use_sync=False)
+
+
+def test_configured_sp_rejects_single_rank_group(monkeypatch):
+    parallel_config = SimpleNamespace(ulysses_degree=2, ring_degree=1, allgather_degree=1)
+    forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(parallel_config=parallel_config))
+    monkeypatch.setattr(factory, "is_forward_context_available", lambda: True)
+    monkeypatch.setattr(factory, "get_forward_context", lambda: forward_context)
+    monkeypatch.setattr(factory, "get_sp_group", lambda: SimpleNamespace())
+    monkeypatch.setattr(factory, "get_sequence_parallel_world_size", lambda: 1)
+
+    with pytest.raises(RuntimeError, match="world size is not greater than one"):
+        factory.build_parallel_attention_strategy(scatter_idx=2, gather_idx=1, use_sync=False)
@@ -72,3 +72,25 @@ def test_ring_prefers_fa4_on_blackwell(monkeypatch):

    assert strategy.run_attention(query, "key", "value", None) == "output"
    assert captured["attn_type"] is AttnType.FA4
+
+
+def test_flashinfer_wrapper_canonicalizes_unbatched_lse_to_bhs(monkeypatch):
+    seq_len = num_heads = 4
+    query = torch.randn(1, seq_len, num_heads, 8)
+    unbatched_out = torch.randn(seq_len, num_heads, 8)
+    unbatched_lse = torch.arange(num_heads * seq_len, dtype=torch.float32).reshape(num_heads, seq_len)
+    assert not torch.equal(unbatched_lse, unbatched_lse.transpose(0, 1))
+
+    def fake_prefill(_q, _k, _v, **_kwargs):
+        return unbatched_out.clone(), unbatched_lse.clone()
+
+    monkeypatch.setattr(ring_kernels, "HAS_FLASHINFER", True)
+    # FlashInfer symbols are only bound when HAS_FLASHINFER was true at import.
+    monkeypatch.setattr(ring_kernels, "single_prefill_with_kv_cache", fake_prefill, raising=False)
+    monkeypatch.setattr(ring_kernels, "_LOG2_E", 1.0, raising=False)
+
+    out, lse = ring_kernels.flashinfer_attn_forward(query, query, query, softmax_scale=0.1)
+
+    assert out.shape == query.shape
+    assert lse.shape == (1, num_heads, seq_len)
+    assert torch.equal(lse[0], unbatched_lse)
@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.ring.ring_utils import update_out_and_lse
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("layout", ["bhs", "bsh", "bhs_padded", "bhs1"])
+def test_ring_lse_known_layouts_are_normalized(layout):
+    block_out = torch.randn(2, 5, 3, 4)
+    canonical = torch.randn(2, 5, 3)
+    if layout == "bhs":
+        block_lse = canonical.transpose(1, 2)
+        lse_layout = "bhs"
+    elif layout == "bsh":
+        block_lse = canonical
+        lse_layout = "bsh"
+    elif layout == "bhs_padded":
+        block_lse = torch.cat([canonical.transpose(1, 2), torch.randn(2, 3, 2)], dim=2)
+        lse_layout = "bhs"
+    else:
+        block_lse = canonical.transpose(1, 2).unsqueeze(-1)
+        lse_layout = "bhs"
+
+    out, lse = update_out_and_lse(None, None, block_out, block_lse, lse_layout=lse_layout)
+
+    assert out.shape == block_out.shape
+    assert out.dtype == torch.float32
+    assert lse.shape == (2, 5, 3, 1)
+    assert torch.equal(lse.squeeze(-1), canonical)
+
+
+def test_ring_lse_seq_len_equals_num_heads_keeps_asymmetric_values():
+    batch, seq_len, num_heads, head_dim = 2, 4, 4, 8
+    block_out = torch.randn(batch, seq_len, num_heads, head_dim)
+    canonical = torch.arange(batch * seq_len * num_heads, dtype=torch.float32).reshape(batch, seq_len, num_heads)
+    bhs = canonical.transpose(1, 2)
+
+    assert bhs.shape == canonical.shape
+    assert not torch.equal(bhs, canonical)
+
+    _, lse_from_bhs = update_out_and_lse(None, None, block_out, bhs, lse_layout="bhs")
+    _, lse_from_bsh = update_out_and_lse(None, None, block_out, canonical, lse_layout="bsh")
+
+    assert torch.equal(lse_from_bhs.squeeze(-1), canonical)
+    assert torch.equal(lse_from_bsh.squeeze(-1), canonical)
+
+    _, lse_if_bsh_read_as_bhs = update_out_and_lse(None, None, block_out, canonical, lse_layout="bhs")
+    assert not torch.equal(lse_if_bsh_read_as_bhs.squeeze(-1), canonical)
+    assert torch.equal(lse_if_bsh_read_as_bhs.squeeze(-1), bhs)
+
+
+@pytest.mark.parametrize(
+    "invalid_lse",
+    [
+        torch.randn(2, 5),
+        torch.randn(2, 4, 4),
+        torch.randn(2, 3, 5, 2),
+    ],
+)
+def test_ring_lse_unknown_layout_is_rejected(invalid_lse):
+    block_out = torch.randn(2, 5, 3, 4)
+
+    with pytest.raises(ValueError, match="Ring LSE"):
+        update_out_and_lse(None, None, block_out, invalid_lse, lse_layout="bhs")
+
+
+def test_ring_accumulation_rejects_mismatched_block_output():
+    out = torch.randn(2, 5, 3, 4)
+    lse = torch.randn(2, 5, 3, 1)
+    block_out = torch.randn(2, 6, 3, 4)
+    block_lse = torch.randn(2, 3, 6)
+
+    with pytest.raises(ValueError, match="block output shape"):
+        update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs")
+
+
+def test_ring_accumulation_matches_logsumexp_weighting():
+    first_out = torch.zeros(1, 2, 1, 1)
+    second_out = torch.ones(1, 2, 1, 1)
+    first_lse = torch.zeros(1, 1, 2)
+    second_lse = torch.zeros(1, 1, 2)
+
+    out, lse = update_out_and_lse(None, None, first_out, first_lse, lse_layout="bhs")
+    out, lse = update_out_and_lse(out, lse, second_out, second_lse, lse_layout="bhs")
+
+    assert torch.allclose(out, torch.full_like(out, 0.5))
+    assert torch.allclose(lse, torch.full_like(lse, torch.log(torch.tensor(2.0))))
@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib
+import sys
+import types
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_sage_attention_rejects_mask_instead_of_ignoring_it(monkeypatch):
+    fake_package = types.ModuleType("sageattention")
+    fake_package.sageattn = lambda *args, **kwargs: pytest.fail("unexpected Sage kernel call")
+    monkeypatch.setitem(sys.modules, "sageattention", fake_package)
+    module_name = "vllm_omni.diffusion.attention.backends.sage_attn"
+    sys.modules.pop(module_name, None)
+    try:
+        backend_module = importlib.import_module(module_name)
+        impl = backend_module.SageAttentionImpl(
+            num_heads=4,
+            head_size=64,
+            softmax_scale=1.0 / 8.0,
+            causal=False,
+        )
+        query = torch.randn(1, 2, 4, 64)
+        metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool))
+
+        with pytest.raises(ValueError, match="does not support attn_mask"):
+            impl.forward_cuda(query, query, query, metadata)
+    finally:
+        sys.modules.pop(module_name, None)
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

 import importlib
 import sys
@@ -9,6 +9,8 @@ import pytest
 import torch
 from vllm.platforms import current_platform

+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


@@ -17,7 +19,7 @@ SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"

 def load_sage_attn3_module(monkeypatch: pytest.MonkeyPatch, kernel_impl):
    fake_module = types.ModuleType("sageattn3")
-    fake_module.sageattn3_blackwell = kernel_impl
+    setattr(fake_module, "sageattn3_blackwell", kernel_impl)
    monkeypatch.setitem(sys.modules, "sageattn3", fake_module)
    sys.modules.pop(SAGE_ATTN3_MODULE, None)
    return importlib.import_module(SAGE_ATTN3_MODULE)
@@ -51,21 +53,11 @@ def test_sage_attn3_forward_uses_blackwell_layout(monkeypatch: pytest.MonkeyPatc
    assert torch.allclose(output, expected)


-def test_sage_attn3_falls_back_to_sdpa_for_gqa(monkeypatch: pytest.MonkeyPatch):
+def test_sage_attn3_rejects_gqa_instead_of_falling_back(monkeypatch: pytest.MonkeyPatch):
    def fake_kernel(*args, **kwargs):
        raise AssertionError("sageattn3_blackwell should not be used for GQA")

    backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
-    sdpa_calls = {}
-
-    def fake_sdpa(query, key, value, **kwargs):
-        sdpa_calls["query_shape"] = query.shape
-        sdpa_calls["key_shape"] = key.shape
-        sdpa_calls["enable_gqa"] = kwargs["enable_gqa"]
-        return query + 1
-
-    monkeypatch.setattr(backend_module.F, "scaled_dot_product_attention", fake_sdpa)
-
    impl = backend_module.SageAttention3Impl(
        num_heads=4,
        head_size=64,
@@ -77,13 +69,38 @@ def test_sage_attn3_falls_back_to_sdpa_for_gqa(monkeypatch: pytest.MonkeyPatch):
    key = torch.randn(2, 8, 2, 64)
    value = torch.randn(2, 8, 2, 64)

-    output = impl.forward_cuda(query, key, value)
+    with pytest.raises(NotImplementedError, match="does not support GQA/MQA"):
+        impl.forward_cuda(query, key, value)

-    assert sdpa_calls["query_shape"] == (2, 4, 8, 64)
-    assert sdpa_calls["key_shape"] == (2, 2, 8, 64)
-    assert sdpa_calls["enable_gqa"] is True
-    expected = (query.permute(0, 2, 1, 3) + 1).permute(0, 2, 1, 3)
-    assert torch.allclose(output, expected)
+
+def test_sage_attn3_rejects_mask_instead_of_ignoring_it(monkeypatch: pytest.MonkeyPatch):
+    def fake_kernel(*args, **kwargs):
+        raise AssertionError("sageattn3_blackwell should not run with an unsupported mask")
+
+    backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
+    impl = backend_module.SageAttention3Impl(
+        num_heads=4,
+        head_size=64,
+        softmax_scale=1.0 / 8.0,
+        causal=False,
+    )
+    query = torch.randn(2, 8, 4, 64)
+    metadata = AttentionMetadata(attn_mask=torch.ones(2, 8, dtype=torch.bool))
+
+    with pytest.raises(ValueError, match="does not support attn_mask"):
+        impl.forward_cuda(query, query, query, metadata)
+
+
+def test_sage_attn3_rejects_custom_softmax_scale(monkeypatch: pytest.MonkeyPatch):
+    backend_module = load_sage_attn3_module(monkeypatch, lambda *args, **kwargs: None)
+
+    with pytest.raises(ValueError, match="does not expose a custom softmax scale"):
+        backend_module.SageAttention3Impl(
+            num_heads=4,
+            head_size=64,
+            softmax_scale=1.0,
+            causal=False,
+        )


@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
@@ -115,10 +132,39 @@ def test_cuda_platform_selects_sage_attn3_alias(monkeypatch: pytest.MonkeyPatch)


@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
-def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch: pytest.MonkeyPatch):
+@pytest.mark.parametrize("head_size", [32, 320])
+def test_cuda_platform_rejects_explicit_sage_attn3_unsupported_head_size(
+    monkeypatch: pytest.MonkeyPatch, head_size: int
+):
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+    from vllm_omni.platforms.cuda import platform as cuda_platform_module
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    original_import_module = importlib.import_module
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+    monkeypatch.setattr(
+        cuda_platform_module.importlib,
+        "import_module",
+        lambda module_name: object() if module_name == "sageattn3" else original_import_module(module_name),
+    )
+    load_sage_attn3_module(monkeypatch, lambda *args, **kwargs: None)
+
+    with pytest.raises(ValueError, match=f"head_size={head_size} is unsupported"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=head_size)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
+def test_cuda_platform_rejects_explicit_sage_attn3_on_unsupported_gpu(monkeypatch: pytest.MonkeyPatch):
    from vllm.platforms.interface import DeviceCapability

-    from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform

@@ -129,6 +175,33 @@ def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch
    )
    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})

-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
+    with pytest.raises(ValueError, match="explicitly selected"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)

-    assert backend_path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
+def test_cuda_platform_rejects_missing_explicit_sage_attn3(monkeypatch: pytest.MonkeyPatch):
+    from vllm.platforms.interface import DeviceCapability
+
+    from vllm_omni.diffusion.envs import PACKAGES_CHECKER
+    from vllm_omni.platforms.cuda import platform as cuda_platform_module
+    from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
+
+    original_import_module = importlib.import_module
+
+    monkeypatch.setattr(
+        CudaOmniPlatform,
+        "get_device_capability",
+        classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)),
+    )
+    monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
+
+    def _missing_sageattn3(module_name):
+        if module_name == "sageattn3":
+            raise ImportError("sageattn3 not installed")
+        return original_import_module(module_name)
+
+    monkeypatch.setattr(cuda_platform_module.importlib, "import_module", _missing_sageattn3)
+
+    with pytest.raises(ImportError, match="explicitly selected"):
+        CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for diffusion attention backend selection."""

 from types import SimpleNamespace
@@ -100,9 +100,12 @@ def test_platform_default_used_without_resolved_spec(monkeypatch, attention_conf
 def test_backend_class_resolution_is_cached(monkeypatch):
    platform_calls = []
    load_calls = []
-    fake_platform = SimpleNamespace(
-        get_diffusion_attn_backend_cls=lambda **kwargs: (platform_calls.append(kwargs) or "fake.module.Backend")
-    )
+
+    def fake_get_cls(**kwargs):
+        platform_calls.append(kwargs)
+        return "fake.module.Backend"
+
+    fake_platform = SimpleNamespace(get_diffusion_attn_backend_cls=fake_get_cls)

    monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform)

@@ -126,6 +129,39 @@ def test_backend_class_resolution_is_cached(monkeypatch):
    assert load_calls == ["fake.module.Backend"]


+def test_capability_query_loads_explicit_cudnn_without_platform(monkeypatch):
+    platform_calls = []
+
+    def fake_get_cls(**kwargs):
+        platform_calls.append(kwargs)
+        return "fake.module.MustNotBeCalled"
+
+    fake_platform = SimpleNamespace(get_diffusion_attn_backend_cls=fake_get_cls)
+    monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform)
+
+    config = AttentionConfig(default=AttentionSpec(backend="CUDNN_ATTN"))
+    backend = selector.get_attn_backend_for_capability(role="self", attention_config=config)
+
+    assert backend.get_name() == "CUDNN_ATTN"
+    assert backend.supports_attention_mask()
+    assert platform_calls == []
+
+
+def test_capability_query_uses_unknown_head_size_for_platform_default(monkeypatch):
+    calls = []
+
+    def fake_get_backend(backend_name, head_size, allow_trtllm_default=True):
+        calls.append((backend_name, head_size, allow_trtllm_default))
+        return _PlatformBackend
+
+    monkeypatch.setattr(selector, "_cached_get_backend_cls", fake_get_backend)
+
+    backend = selector.get_attn_backend_for_capability(role="self", attention_config=AttentionConfig())
+
+    assert backend is _PlatformBackend
+    assert calls == [(None, selector.HEAD_SIZE_UNKNOWN, True)]
+
+
 def test_load_backend_cls_reports_missing_module():
    with pytest.raises(ImportError, match="Failed to import module missing_attention_backend"):
        selector._load_backend_cls("missing_attention_backend.Backend")
@@ -29,31 +29,343 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

@pytest.mark.core_model
@pytest.mark.cpu
-def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
-    monkeypatch.setattr(
-        "vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
-        lambda default: "advanced_uaa",
+def test_uaa_gqa_head_padding_preserves_the_query_to_kv_ratio(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Q must be padded to a multiple of the *KV* padding, not of world_size.
+
+    Rounding each up independently is the obvious implementation and is wrong:
+    28Q/7KV at SP2 would become 28/8, i.e. 14/4 per rank, whose 3.5 ratio is not
+    a valid GQA shape. Deriving Q from the padded KV count gives 32/8 -> 16/4.
+    """
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    class FakeGroup:
+        ulysses_group = object()
+        ulysses_world_size = 2
+        ulysses_rank = 0
+        ring_world_size = 1
+
+    observed = []
+
+    def fake_all_to_all(pg, tensor, **kwargs):
+        observed.append(kwargs["padded_head_cnt"])
+        return tensor, tensor.shape[2]
+
+    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
+    strategy = ulysses.UlyssesParallelAttention(
+        FakeGroup(),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
    )
+
+    with set_forward_context():
+        strategy.pre_attention(
+            torch.zeros(1, 3, 28, 4),
+            torch.zeros(1, 3, 7, 4),
+            torch.zeros(1, 3, 7, 4),
+            None,
+        )
+
+    q_padded, k_padded, v_padded = observed
+    assert (q_padded, k_padded, v_padded) == (32, 8, 8)
+    # The per-rank shape must still be a valid GQA shape.
+    assert (q_padded // 2) % (k_padded // 2) == 0
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_uaa_rejects_head_counts_that_are_not_a_gqa_shape(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    class FakeGroup:
+        ulysses_group = object()
+        ulysses_world_size = 2
+        ulysses_rank = 0
+        ring_world_size = 1
+
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
+    strategy = ulysses.UlyssesParallelAttention(
+        FakeGroup(),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+
+    with set_forward_context(), pytest.raises(ValueError, match="multiple of KV heads"):
+        strategy.pre_attention(
+            torch.zeros(1, 3, 10, 4),
+            torch.zeros(1, 3, 4, 4),
+            torch.zeros(1, 3, 4, 4),
+            None,
+        )
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "q_heads, kv_heads, world_size, expected_warn",
+    [
+        # 28Q/7KV @ U=2 -> pad to 32Q, 32/28 = 1.14x, below threshold, no warn.
+        (28, 7, 2, False),
+        # 32Q/1KV @ U=8 -> pad KV to 8 -> pad Q to 256, 256/32 = 8x, warn.
+        (32, 1, 8, True),
+        # 16Q/2KV @ U=4 -> pad KV to 4 -> pad Q to 32, 32/16 = 2x, warn.
+        (16, 2, 4, True),
+    ],
+)
+def test_uaa_padding_ratio_warning_fires_when_blowup_is_large(
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: pytest.LogCaptureFixture,
+    q_heads: int,
+    kv_heads: int,
+    world_size: int,
+    expected_warn: bool,
+) -> None:
+    """advanced_uaa must warn once when GQA padding materially inflates Q heads.
+
+    MQA-style shapes at high ulysses_degree pad Q by U x, which grows attention
+    FLOPs and temporary VRAM by the same factor. We do not fall back (there is
+    no cheap alternative -- Ulysses needs KV divisible by U for the head-dim
+    split), but we surface a one-shot warning so the user can pick a
+    ulysses_degree that divides KV_heads if the overhead is unacceptable.
+    """
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    # Reset the warn-once dedup set so parametrized cases do not shadow each other.
+    ulysses._uaa_pad_ratio_warned.clear()
+
+    class FakeGroup:
+        ulysses_group = object()
+        ulysses_world_size = world_size
+        ulysses_rank = 0
+        ring_world_size = 1
+
+    def fake_all_to_all(pg, tensor, **kwargs):
+        return tensor, tensor.shape[2]
+
+    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    monkeypatch.setattr(ulysses, "_all_gather_int", lambda *a, **kw: [3] * world_size)
+    strategy = ulysses.UlyssesParallelAttention(
+        FakeGroup(),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+
+    caplog.set_level("WARNING", logger="vllm_omni.diffusion.attention.parallel.ulysses")
+    with set_forward_context():
+        strategy.pre_attention(
+            torch.zeros(1, 3, q_heads, 4),
+            torch.zeros(1, 3, kv_heads, 4),
+            torch.zeros(1, 3, kv_heads, 4),
+            None,
+        )
+
+    warn_msgs = [r.message for r in caplog.records if "GQA padding inflates" in r.message]
+    if expected_warn:
+        assert len(warn_msgs) == 1, warn_msgs
+        # Second call with the same shape must not re-warn (warn-once dedup).
+        with set_forward_context():
+            strategy.pre_attention(
+                torch.zeros(1, 3, q_heads, 4),
+                torch.zeros(1, 3, kv_heads, 4),
+                torch.zeros(1, 3, kv_heads, 4),
+                None,
+            )
+        warn_msgs = [r.message for r in caplog.records if "GQA padding inflates" in r.message]
+        assert len(warn_msgs) == 1, "warn-once dedup should suppress the second call"
+    else:
+        assert warn_msgs == [], f"unexpected warning for {q_heads}Q/{kv_heads}KV @ U={world_size}: {warn_msgs}"
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_uaa_joint_gqa_head_padding_preserves_the_query_to_kv_ratio(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Joint Q/K/V must also preserve the GQA ratio through Ulysses padding.
+
+    MMDiT-style models (SD3, FLUX, Qwen-Image, Hunyuan-Video, ...) pass a
+    separate text stream via ``AttentionMetadata.joint_*``. Before the fix the
+    joint block padded K/V by Q's amount and integer-divided K/V by world_size,
+    which for 28Q/7KV at SP=2 produces mismatched joint vs main head counts and
+    the downstream ring-attention concat blows up. Post-fix the joint block
+    mirrors the main-path rule: pad KV to a world_size multiple, derive Q as
+    ``padded_kv * (Q // KV)``.
+    """
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    class FakeGroup:
+        ulysses_group = object()
+        ulysses_world_size = 2
+        ulysses_rank = 0
+        ring_world_size = 1
+
+    def fake_all_to_all(pg, tensor, **kwargs):
+        # Simulate the head-dim sharding so the post-a2a joint concat has
+        # matching per-rank head counts on Q and joint_Q.
+        heads_per_rank = kwargs["padded_head_cnt"] // 2
+        return tensor[:, :, :heads_per_rank, :].contiguous(), tensor.shape[2]
+
+    monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
+    strategy = ulysses.UlyssesParallelAttention(
+        FakeGroup(),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+
+    attn_metadata = AttentionMetadata(
+        joint_query=torch.zeros(1, 5, 28, 4),
+        joint_key=torch.zeros(1, 5, 7, 4),
+        joint_value=torch.zeros(1, 5, 7, 4),
+    )
+
+    with set_forward_context():
+        strategy.pre_attention(
+            torch.zeros(1, 3, 28, 4),
+            torch.zeros(1, 3, 7, 4),
+            torch.zeros(1, 3, 7, 4),
+            attn_metadata,
+        )
+
+    # 28Q/7KV at SP=2 -> pad to 32/8 -> per-rank slice is 16Q/4KV.
+    assert attn_metadata.joint_key.shape[-2] == 4
+    assert attn_metadata.joint_value.shape[-2] == 4
+    # Per-rank joint shape must remain a valid GQA shape (Q multiple of KV).
+    per_rank_joint_q = 32 // 2
+    per_rank_joint_kv = attn_metadata.joint_key.shape[-2]
+    assert per_rank_joint_q % per_rank_joint_kv == 0
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_uaa_rejects_joint_head_counts_that_are_not_a_gqa_shape(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
+    class FakeGroup:
+        ulysses_group = object()
+        ulysses_world_size = 2
+        ulysses_rank = 0
+        ring_world_size = 1
+
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
+    strategy = ulysses.UlyssesParallelAttention(
+        FakeGroup(),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+
+    attn_metadata = AttentionMetadata(
+        joint_query=torch.zeros(1, 5, 10, 4),
+        joint_key=torch.zeros(1, 5, 4, 4),
+        joint_value=torch.zeros(1, 5, 4, 4),
+    )
+
+    with set_forward_context(), pytest.raises(ValueError, match="multiple of joint KV heads"):
+        strategy.pre_attention(
+            torch.zeros(1, 3, 8, 4),
+            torch.zeros(1, 3, 4, 4),
+            torch.zeros(1, 3, 4, 4),
+            attn_metadata,
+        )
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_advanced_uaa_hybrid_rejects_non_gqa_shapes(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Hybrid Ulysses+Ring still rejects shapes that are not a GQA shape.
+
+    Valid GQA shapes are safe under advanced_uaa head padding (the derived-Q
+    rule keeps real Q heads paired with real K/V heads after the split, see
+    pre_attention), but a query count that is not a multiple of the K/V count
+    cannot be padded while preserving the ratio, so it must be rejected.
+    """
+    from vllm_omni.diffusion.attention.parallel import ulysses
+
    sp_group = SimpleNamespace(
        ulysses_group=None,
        ulysses_world_size=2,
        ulysses_rank=0,
        ring_world_size=2,
+        ring_group=object(),
    )
+    monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
+    monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
    strategy = UlyssesParallelAttention(
        sp_group=sp_group,
        scatter_idx=2,
        gather_idx=1,
        use_sync=False,
    )
-    query = torch.zeros(1, 2, 8, 4)
-    key = torch.zeros(1, 2, 1, 4)
+    query = torch.zeros(1, 2, 10, 4)
+    key = torch.zeros(1, 2, 4, 4)
    value = torch.zeros_like(key)

-    with pytest.raises(ValueError, match="does not support K/V head padding"):
+    with pytest.raises(ValueError, match="multiple of KV heads"):
        strategy.pre_attention(query, key, value, None)


+def test_strict_ulysses_reshards_vsa_gate_with_qkv(monkeypatch) -> None:
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
+        lambda default: "strict",
+    )
+    calls: list[torch.Tensor] = []
+
+    def fake_all_to_all(pg, tensor, scatter_idx, gather_idx, use_sync):
+        del pg, scatter_idx, gather_idx, use_sync
+        calls.append(tensor)
+        return tensor.reshape(tensor.shape[0], tensor.shape[1] * 2, tensor.shape[2] // 2, tensor.shape[3])
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.parallel.ulysses.SeqAllToAll4D.apply",
+        fake_all_to_all,
+    )
+    strategy = UlyssesParallelAttention(
+        sp_group=SimpleNamespace(
+            ulysses_group=object(),
+            ulysses_world_size=2,
+            ulysses_rank=0,
+            ring_world_size=1,
+        ),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+    query = torch.arange(24).reshape(1, 3, 4, 2)
+    key = query + 10
+    value = query + 20
+    gate = query + 100
+    metadata = AttentionMetadata(extra={"gate_compress": gate})
+
+    query_out, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
+
+    assert all(actual is expected for actual, expected in zip(calls, (query, key, value, gate), strict=True))
+    assert query_out.shape == (1, 6, 2, 2)
+    assert metadata_out is metadata
+    assert metadata.extra["gate_compress"].shape == query_out.shape
+    torch.testing.assert_close(metadata.extra["gate_compress"], gate.reshape_as(query_out))
+
+
 def _run_attention_case(
    local_rank: int,
    world_size: int,
@@ -67,6 +379,9 @@ def _run_attention_case(
    ring_degree: int = 1,
    split_sizes: list[int] | None = None,
    sdp_kernel_mode: str = "math",
+    num_kv_heads: int | None = None,
+    has_joint: bool = False,
+    joint_strategy: str = "front",
 ) -> None:
    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
    current_omni_platform.set_device(device)
@@ -102,15 +417,33 @@ def _run_attention_case(
            head_size=head_size,
            causal=False,
            softmax_scale=1.0 / (head_size**0.5),
+            num_kv_heads=num_kv_heads,
        ).to(device=device, dtype=torch.float32)

        with np.load(input_file, allow_pickle=False) as payload:
            q_full = torch.from_numpy(payload["q"]).to(device=device)
            k_full = torch.from_numpy(payload["k"]).to(device=device)
            v_full = torch.from_numpy(payload["v"]).to(device=device)
+            if has_joint:
+                joint_q = torch.from_numpy(payload["joint_q"]).to(device=device)
+                joint_k = torch.from_numpy(payload["joint_k"]).to(device=device)
+                joint_v = torch.from_numpy(payload["joint_v"]).to(device=device)

+        attn_metadata = None
        if world_size == 1:
-            q, k, v = q_full, k_full, v_full
+            if has_joint:
+                # Baseline: SP is inactive so pre_attention won't concat joint for us.
+                # Build the full sequence directly and skip the joint metadata.
+                if joint_strategy == "front":
+                    q = torch.cat([joint_q, q_full], dim=1)
+                    k = torch.cat([joint_k, k_full], dim=1)
+                    v = torch.cat([joint_v, v_full], dim=1)
+                else:
+                    q = torch.cat([q_full, joint_q], dim=1)
+                    k = torch.cat([k_full, joint_k], dim=1)
+                    v = torch.cat([v_full, joint_v], dim=1)
+            else:
+                q, k, v = q_full, k_full, v_full
        else:
            if split_sizes is None:
                # NOTE: torch.chunk may return fewer than `world_size` chunks for some
@@ -135,6 +468,16 @@ def _run_attention_case(
            # shard manually for testing.
            get_forward_context()._sp_shard_depth = 1

+            if has_joint:
+                from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+                attn_metadata = AttentionMetadata(
+                    joint_query=joint_q,
+                    joint_key=joint_k,
+                    joint_value=joint_v,
+                    joint_strategy=joint_strategy,
+                )
+
        if sdp_kernel_mode == "math":
            sdp_ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
        elif sdp_kernel_mode == "mem_efficient":
@@ -143,25 +486,47 @@ def _run_attention_case(
            raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}")

        with sdp_ctx:
-            out_local = attn(q, k, v).contiguous()
+            out_local = attn(q, k, v, attn_metadata=attn_metadata).contiguous()

        if world_size == 1:
            out_full = out_local
        else:
-            local_len = torch.tensor([out_local.shape[1]], device=device, dtype=torch.int64)
+            if has_joint:
+                joint_len = int(joint_q.shape[1])
+                if joint_strategy == "front":
+                    joint_out_local = out_local[:, :joint_len]
+                    img_out_local = out_local[:, joint_len:]
+                else:
+                    img_out_local = out_local[:, : out_local.shape[1] - joint_len]
+                    joint_out_local = out_local[:, -joint_len:]
+            else:
+                joint_out_local = None
+                img_out_local = out_local
+
+            local_len = torch.tensor([img_out_local.shape[1]], device=device, dtype=torch.int64)
            gathered_lens = [torch.empty_like(local_len) for _ in range(world_size)]
            torch.distributed.all_gather(gathered_lens, local_len)
            lens = [int(t.item()) for t in gathered_lens]
            max_len = max(lens)

-            if out_local.shape[1] < max_len:
-                pad = max_len - out_local.shape[1]
-                out_local = torch.nn.functional.pad(out_local, (0, 0, 0, 0, 0, pad)).contiguous()
+            if img_out_local.shape[1] < max_len:
+                pad = max_len - img_out_local.shape[1]
+                img_out_local = torch.nn.functional.pad(img_out_local, (0, 0, 0, 0, 0, pad)).contiguous()
+            else:
+                img_out_local = img_out_local.contiguous()

-            gathered = [torch.empty_like(out_local) for _ in range(world_size)]
-            torch.distributed.all_gather(gathered, out_local)
+            gathered = [torch.empty_like(img_out_local) for _ in range(world_size)]
+            torch.distributed.all_gather(gathered, img_out_local)
            if local_rank == 0:
-                out_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1)
+                img_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1)
+                if has_joint:
+                    out_full = (
+                        torch.cat([joint_out_local, img_full], dim=1)
+                        if joint_strategy == "front"
+                        else torch.cat([img_full, joint_out_local], dim=1)
+                    )
+                else:
+                    out_full = img_full
            else:
                out_full = None

@@ -172,20 +537,43 @@ def _run_attention_case(


@pytest.mark.parametrize(
-    "sp_world_size,seq_len,num_heads",
+    "sp_world_size,seq_len,joint_len,num_heads,num_kv_heads",
    [
-        (2, 6, 3),  # head_cnt not divisible by P=2
-        (2, 5, 4),  # seq_len not divisible by P=2
-        (4, 9, 30),  # Z-Image-like: head_cnt not divisible by P=4
-        (4, 10, 8),  # seq_len not divisible by P=4
+        # MHA (num_kv_heads=None, joint_len=None)
+        (2, 6, None, 3, None),  # head_cnt not divisible by P=2
+        (2, 5, None, 4, None),  # seq_len not divisible by P=2
+        (4, 9, None, 30, None),  # Z-Image-like: head_cnt not divisible by P=4
+        (4, 10, None, 8, None),  # seq_len not divisible by P=4
+        # GQA (joint_len=None)
+        (2, 6, None, 28, 7),  # BOOGU-like GQA: neither Q nor KV divisible by P=2
+        (2, 5, None, 6, 3),  # KV divisible by P, Q is not
+        (4, 8, None, 12, 3),  # KV not divisible by P=4
+        # Joint GQA (MMDiT-style text stream via AttentionMetadata.joint_*)
+        (2, 6, 5, 28, 7),  # BOOGU-like GQA joint: KV=7 not divisible by P=2
+        (2, 5, 4, 6, 3),  # KV divisible by P, Q is not
    ],
 )
-def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_heads: int) -> None:
+def test_ulysses_uaa_matches_baseline(
+    sp_world_size: int,
+    seq_len: int,
+    joint_len: int | None,
+    num_heads: int,
+    num_kv_heads: int | None,
+) -> None:
+    """End-to-end SP-vs-single-GPU parity under UAA.
+
+    Covers MHA, GQA and joint (MMDiT-style) shapes. The joint variant catches
+    the pre-fix bug where the joint block padded K/V by Q's head count and
+    integer-divided K/V by world_size, giving per-rank shapes that mismatched
+    the main-tensor slice (e.g. 14Q/3KV vs. 16Q/4KV for 28Q/7KV at SP=2).
+    """
    if current_omni_platform.get_device_count() < sp_world_size:
        pytest.skip(f"Test requires {sp_world_size} GPUs")

    batch_size = 2
    head_size = 8
+    kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+    has_joint = joint_len is not None

    base_init_method = get_distributed_init_method()
    sp_init_method = get_distributed_init_method()
@@ -200,17 +588,23 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
    try:
        torch.manual_seed(0)
        q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
-        k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
-        v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
-        np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy())
+        k = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32)
+        v = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32)
+        payload = {"q": q.numpy(), "k": k.numpy(), "v": v.numpy()}
+        if has_joint:
+            payload["joint_q"] = torch.randn(batch_size, joint_len, num_heads, head_size, dtype=torch.float32).numpy()
+            payload["joint_k"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
+            payload["joint_v"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
+        np.savez(input_file, **payload)
+
+        common_tail = (1, None, "math", num_kv_heads, has_joint)

        # Baseline (no SP)
        torch.multiprocessing.spawn(
            _run_attention_case,
-            args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"),
+            args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict", *common_tail),
            nprocs=1,
        )
-
        # SP (Ulysses-P with UAA)
        torch.multiprocessing.spawn(
            _run_attention_case,
@@ -223,15 +617,13 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
                head_size,
                sp_world_size,
                "advanced_uaa",
+                *common_tail,
            ),
            nprocs=sp_world_size,
        )

-        baseline = np.load(baseline_file, allow_pickle=False)
-        sp = np.load(sp_file, allow_pickle=False)
-
-        baseline_t = torch.from_numpy(baseline)
-        sp_t = torch.from_numpy(sp)
+        baseline_t = torch.from_numpy(np.load(baseline_file, allow_pickle=False))
+        sp_t = torch.from_numpy(np.load(sp_file, allow_pickle=False))
        assert baseline_t.shape == sp_t.shape
        torch.testing.assert_close(sp_t, baseline_t, atol=1e-5, rtol=1e-5)
    finally:
@@ -242,7 +634,17 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
                pass


-def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
+@pytest.mark.parametrize(
+    "num_heads,num_kv_heads,joint_len",
+    [
+        (3, None, None),  # MHA baseline (head_cnt not divisible by ulysses_degree=2)
+        (28, 7, None),  # GQA: exercises Ring pytorch_attn_forward KV expansion (Q!=K heads)
+        (28, 7, 5),  # GQA + MMDiT joint stream: joint K/V ride the padded Ulysses layout and Ring SDPA
+    ],
+)
+def test_ulysses_uaa_hybrid_ring_matches_baseline(
+    num_heads: int, num_kv_heads: int | None, joint_len: int | None
+) -> None:
    sp_world_size = 4
    ulysses_degree = 2
    ring_degree = 2
@@ -253,9 +655,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
    batch_size = 2
    head_size = 8
    seq_len = 10
-    # This is the positive Hybrid case: K/V heads must not require padding,
-    # which is unsupported by advanced_uaa when Ring is also enabled.
-    num_heads = 4
+    kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+    has_joint = joint_len is not None

    # Ensure ring ranks see equal post-Ulysses seq_len:
    # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
@@ -274,9 +675,14 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
    try:
        torch.manual_seed(0)
        q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
-        k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
-        v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
-        np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy())
+        k = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32)
+        v = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32)
+        payload = {"q": q.numpy(), "k": k.numpy(), "v": v.numpy()}
+        if has_joint:
+            payload["joint_q"] = torch.randn(batch_size, joint_len, num_heads, head_size, dtype=torch.float32).numpy()
+            payload["joint_k"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
+            payload["joint_v"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
+        np.savez(input_file, **payload)

        # Baseline (no SP)
        torch.multiprocessing.spawn(
@@ -293,6 +699,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
                1,
                None,
                "mem_efficient",
+                num_kv_heads,
+                has_joint,
            ),
            nprocs=1,
        )
@@ -312,6 +720,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
                ring_degree,
                split_sizes,
                "mem_efficient",
+                num_kv_heads,
+                has_joint,
            ),
            nprocs=sp_world_size,
        )
@@ -6,7 +6,9 @@ Model specific tests for CacheDiT enablement.
 """

 import ast
+from contextlib import contextmanager
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import Mock, patch

 import pytest
@@ -17,7 +19,8 @@ from vllm.distributed import parallel_state
 import vllm_omni.diffusion.cache.cachedit as cd_backend
 import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
-from vllm_omni.diffusion.data import DiffusionCacheConfig
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig
 from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
 from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
@@ -35,6 +38,17 @@ SEPARATE_CFG_TRANSFORMERS = [
 SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()


+@contextmanager
+def _force_torch_sdpa():
+    """Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3)."""
+    od_config = SimpleNamespace(
+        diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    with set_current_diffusion_config(od_config):
+        yield
+
+
 def test_custom_cache_dit_enablers_are_registered_explicitly():
    expected_enablers = {
        "Wan22Pipeline": cd_model_specific.enable_cache_for_wan22,
@@ -213,22 +227,23 @@ def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
    text_in = torch.full((1, seq_len, 16), 3.0)
    audio_text_in = torch.full((1, seq_len, 16), 4.0)

-    model = LTX2VideoTransformer3DModel(
-        in_channels=16,
-        out_channels=16,
-        patch_size=1,
-        patch_size_t=1,
-        num_attention_heads=2,
-        attention_head_dim=8,
-        cross_attention_dim=16,
-        audio_in_channels=16,
-        audio_out_channels=16,
-        audio_num_attention_heads=2,
-        audio_attention_head_dim=8,
-        audio_cross_attention_dim=16,
-        num_layers=2,
-        caption_channels=16,
-    )
+    with _force_torch_sdpa():
+        model = LTX2VideoTransformer3DModel(
+            in_channels=16,
+            out_channels=16,
+            patch_size=1,
+            patch_size_t=1,
+            num_attention_heads=2,
+            attention_head_dim=8,
+            cross_attention_dim=16,
+            audio_in_channels=16,
+            audio_out_channels=16,
+            audio_num_attention_heads=2,
+            audio_attention_head_dim=8,
+            audio_cross_attention_dim=16,
+            num_layers=2,
+            caption_channels=16,
+        )

    # NOTE: This is currently using the LTX2 custom enabler, but the custom
    # enablers will be consolidated after
@@ -13,6 +13,8 @@ import vllm_omni.diffusion.data as diffusion_data
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
 from vllm_omni.diffusion.data import (
+    AttentionConfig,
+    AttentionSpec,
    DiffusionOutput,
    DiffusionParallelConfig,
    OmniDiffusionConfig,
@@ -141,9 +143,96 @@ class TestPipelineArgumentsHandling:
        assert isinstance(output.output, MockPipelineOutput)
        assert output.output.image is stub_image

+    @staticmethod
+    def _make_adapter(backend: str, set_attention_backend):
+        adapter = DiffusersAdapterPipeline(
+            od_config=_make_od_config(
+                diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend=backend)),
+            )
+        )
+        adapter._pipeline = SimpleNamespace(
+            transformer=SimpleNamespace(set_attention_backend=set_attention_backend),
+        )
+        return adapter
+
+    def test_explicit_unbound_backend_is_rejected(self):
+        calls: list[str] = []
+        adapter = self._make_adapter("CUDNN_ATTN", calls.append)
+
+        with pytest.raises(ValueError, match="do not provide a binding"):
+            adapter._set_attention_backend()
+
+        assert calls == []
+
+    def test_failed_explicit_backend_does_not_fallback_to_native(self):
+        calls: list[str] = []
+
+        def reject_backend(backend):
+            calls.append(backend)
+            raise RuntimeError(f"{backend} unavailable")
+
+        adapter = self._make_adapter("SAGE_ATTN", reject_backend)
+
+        with pytest.raises(RuntimeError, match="explicitly selected attention backend 'SAGE_ATTN'"):
+            adapter._set_attention_backend()
+
+        assert calls == ["sage_hub", "sage", "sage_varlen"]
+        assert "native" not in calls
+
+    def test_explicit_flash_attn_3_hub_does_not_fallback_to_fa2_or_local(self, monkeypatch):
+        monkeypatch.setattr(DiffusersAdapterPipeline, "_is_blackwell_gpu", lambda self: False)
+        calls: list[str] = []
+
+        def accept_later_backends(backend):
+            calls.append(backend)
+            if backend in ("flash_hub", "flash_varlen_hub", "flash", "flash_varlen", "_native_flash"):
+                return
+            raise RuntimeError(f"{backend} unavailable")
+
+        adapter = self._make_adapter("FLASH_ATTN_3_HUB", accept_later_backends)
+
+        with pytest.raises(RuntimeError, match="explicitly selected attention backend 'FLASH_ATTN_3_HUB'"):
+            adapter._set_attention_backend()
+
+        assert calls == ["_flash_3_hub", "_flash_3_varlen_hub"]
+
+    def test_explicit_flash_attn_hub_does_not_fallback_to_local(self, monkeypatch):
+        monkeypatch.setattr(DiffusersAdapterPipeline, "_is_blackwell_gpu", lambda self: False)
+        calls: list[str] = []
+
+        def accept_local_flash(backend):
+            calls.append(backend)
+            if backend in ("flash", "flash_varlen", "_native_flash"):
+                return
+            raise RuntimeError(f"{backend} unavailable")
+
+        adapter = self._make_adapter("FLASH_ATTN_HUB", accept_local_flash)
+
+        with pytest.raises(RuntimeError, match="explicitly selected attention backend 'FLASH_ATTN_HUB'"):
+            adapter._set_attention_backend()
+
+        assert calls == ["flash_hub", "flash_varlen_hub"]
+
+    def test_explicit_torch_sdpa_maps_to_diffusers_native(self):
+        calls: list[str] = []
+        adapter = self._make_adapter("TORCH_SDPA", calls.append)
+
+        adapter._set_attention_backend()
+
+        assert calls == ["native"]
+
    @pytest.mark.parametrize(
        "feature_id",
-        ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"],
+        [
+            "cfg_parallel",
+            "tensor_parallel",
+            "ulysses",
+            "ring",
+            "teacache",
+            "cache_dit",
+            "enforce_eager",
+            "unsupported_quantization",
+        ],
    )
    def test_adapter_guard_unsupported_feature(self, feature_id):
        if feature_id == "cfg_parallel":
@@ -151,6 +240,11 @@ class TestPipelineArgumentsHandling:
                parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1),
                cache_backend="none",
            )
+        elif feature_id == "tensor_parallel":
+            od_config = _make_od_config(
+                parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, sequence_parallel_size=1),
+                cache_backend="none",
+            )
        elif feature_id == "ulysses":
            od_config = _make_od_config(
                parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ulysses_degree=2),
@@ -17,7 +17,12 @@ import torch

 from tests.helpers.mark import hardware_marks, hardware_test
 from vllm_omni.diffusion.distributed import parallel_state
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _dispatch_branches, _wrap
+from vllm_omni.diffusion.distributed.cfg_parallel import (
+    CFGParallelMixin,
+    _dispatch_branches,
+    _get_cfg_world_size_or_one,
+    _wrap,
+)
 from vllm_omni.diffusion.distributed.parallel_state import (
    destroy_distributed_env,
    get_classifier_free_guidance_rank,
@@ -393,6 +398,68 @@ def _build_pipeline_on_device(
 # ---------------------------------------------------------------------------


+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_world_size_defaults_only_when_group_is_uninitialized(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(parallel_state, "_CFG", None)
+
+    assert not parallel_state.is_cfg_group_initialized()
+    assert _get_cfg_world_size_or_one() == 1
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_world_size_does_not_swallow_initialized_group_errors(monkeypatch: pytest.MonkeyPatch):
+    class BrokenCfgGroup:
+        @property
+        def world_size(self) -> int:
+            raise AssertionError("CFG group is broken")
+
+    monkeypatch.setattr(parallel_state, "_CFG", BrokenCfgGroup())
+
+    assert parallel_state.is_cfg_group_initialized()
+    with pytest.raises(AssertionError, match="CFG group is broken"):
+        _get_cfg_world_size_or_one()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_defaults_to_sequential_without_initialized_group():
+    pipeline = _build_pipeline_on_device(
+        TestCFGPipeline,
+        hidden_dim=32,
+        model_seed=42,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+    positive_kwargs, negative_kwargs = _make_two_branch_inputs(
+        batch_size=1,
+        channels=4,
+        height=4,
+        width=4,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        input_seed=123,
+    )
+
+    with torch.no_grad():
+        positive = pipeline.predict_noise(**positive_kwargs)
+        negative = pipeline.predict_noise(**negative_kwargs)
+        actual = pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=5.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    expected = negative + 5.0 * (positive - negative)
+    torch.testing.assert_close(actual, expected)
+
+
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
@@ -989,6 +989,89 @@ class TestDimensionValidation:
        # Since expected_dims=3 but tensor has 4 dims, should return original
        assert result.shape == tensor_4d.shape

+    def test_auto_pad_calls_attention_mask_capability(self, monkeypatch):
+        from vllm_omni.diffusion.attention import selector
+        from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook
+
+        class NoMaskBackend:
+            @classmethod
+            def supports_attention_mask(cls, attention_spec=None):
+                del attention_spec
+                return False
+
+            @staticmethod
+            def get_name():
+                return "NO_MASK_BACKEND"
+
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: NoMaskBackend)
+
+        hook = SequenceParallelSplitHook(
+            metadata={},
+            config=SequenceParallelConfig(ulysses_degree=2, ring_degree=1),
+        )
+        tensor = torch.randn(1, 3, 8)
+
+        with pytest.raises(ValueError, match="does not support attention_mask"):
+            hook._shard_with_auto_pad(tensor, dim=1)
+
+    def test_auto_pad_allows_explicit_cudnn(self, monkeypatch):
+        from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
+        from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import set_forward_context
+        from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook
+
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 0)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+
+        hook = SequenceParallelSplitHook(
+            metadata={},
+            config=SequenceParallelConfig(ulysses_degree=2, ring_degree=1),
+        )
+        tensor = torch.randn(1, 3, 8)
+        od_config = OmniDiffusionConfig(
+            model="test",
+            dtype=torch.float32,
+            diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="CUDNN_ATTN")),
+        )
+
+        with set_forward_context(omni_diffusion_config=od_config):
+            result = hook._shard_with_auto_pad(tensor, dim=1)
+
+        # seq_len=3 padded to 4, then sharded by world_size=2 → rank 0 gets 2.
+        assert result.shape == (1, 2, 8)
+
+    def test_auto_pad_rejects_explicit_flashinfer_cute_on_blackwell(self, monkeypatch):
+        from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
+        from vllm_omni.diffusion.distributed import parallel_state
+        from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
+        from vllm_omni.diffusion.forward_context import set_forward_context
+        from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook
+
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 0)
+        monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
+
+        hook = SequenceParallelSplitHook(
+            metadata={},
+            config=SequenceParallelConfig(ulysses_degree=2, ring_degree=1),
+        )
+        tensor = torch.randn(1, 3, 8)
+        od_config = OmniDiffusionConfig(
+            model="test",
+            dtype=torch.float32,
+            diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FLASHINFER_ATTN")),
+        )
+
+        with set_forward_context(omni_diffusion_config=od_config):
+            with pytest.raises(ValueError, match="does not support attention_mask"):
+                hook._shard_with_auto_pad(tensor, dim=1)
+

@pytest.mark.cpu
 class TestStrictModeSplitValidation:
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.mot.ops import mot_gemm
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_CONFIG = {
+    "BLOCK_SIZE_M": 16,
+    "BLOCK_SIZE_N": 64,
+    "BLOCK_SIZE_K": 32,
+    "GROUP_SIZE_M": 8,
+    "num_warps": 4,
+    "num_stages": 2,
+}
+
+
+class _RecordingKernel:
+    def __init__(self) -> None:
+        self.calls: list[tuple[tuple[object, ...], dict[str, object]]] = []
+
+    def __getitem__(self, grid):
+        del grid
+        return self._launch
+
+    def _launch(self, *args, **kwargs) -> None:
+        self.calls.append((args, kwargs))
+
+
+@pytest.mark.parametrize(
+    ("bias_text", "bias_vae"),
+    [(None, None), ("text", None), (None, "vae"), ("text", "vae")],
+    ids=["no-bias", "text-only", "vae-only", "both"],
+)
+def test_invoke_mot_gemm_accepts_independent_biases(
+    bias_text: str | None,
+    bias_vae: str | None,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _RecordingKernel()
+    monkeypatch.setattr(mot_gemm, "mot_unified_gemm_kernel", kernel)
+
+    A = torch.ones((3, 4), dtype=torch.float16)
+    B_text = torch.ones((4, 5), dtype=torch.float16)
+    B_vae = torch.ones((4, 5), dtype=torch.float16)
+    C = torch.empty((3, 5), dtype=torch.float16)
+    text_indices = torch.tensor([0, 2], dtype=torch.long)
+    vae_indices = torch.tensor([1], dtype=torch.long)
+    text_bias = torch.ones(5, dtype=torch.float16) if bias_text else None
+    vae_bias = torch.ones(5, dtype=torch.float16) if bias_vae else None
+
+    mot_gemm.invoke_mot_gemm(
+        A=A,
+        B_text=B_text,
+        B_vae=B_vae,
+        C=C,
+        bias_text=text_bias,
+        bias_vae=vae_bias,
+        text_indices=text_indices,
+        vae_indices=vae_indices,
+        A_scale=None,
+        B_text_scale=None,
+        B_vae_scale=None,
+        use_fp8_w8a8=False,
+        use_int8_w8a8=False,
+        use_int8_w8a16=False,
+        use_int4_w4a16=False,
+        A_per_channel_quant=False,
+        B_per_channel_quant=False,
+        config=_CONFIG,
+    )
+
+    assert len(kernel.calls) == 1
+    args, kwargs = kernel.calls[0]
+    assert kwargs["HAS_BIAS_TEXT"] is (text_bias is not None)
+    assert kwargs["HAS_BIAS_VAE"] is (vae_bias is not None)
+    assert "HAS_BIAS" not in kwargs
+
+    launched_text_bias, launched_vae_bias = args[4:6]
+    if text_bias is None and vae_bias is None:
+        assert launched_text_bias == launched_vae_bias == 0
+    elif text_bias is None:
+        assert launched_text_bias is vae_bias
+        assert launched_vae_bias is vae_bias
+    elif vae_bias is None:
+        assert launched_text_bias is text_bias
+        assert launched_vae_bias is text_bias
+    else:
+        assert launched_text_bias is text_bias
+        assert launched_vae_bias is vae_bias
@@ -513,6 +513,57 @@ def test_mot_o_proj(
            )


+@pytest.mark.parametrize(
+    ("bias_text", "bias_vae"),
+    [(True, False), (False, True)],
+    ids=["text-only", "vae-only"],
+)
+def test_mot_row_one_sided_bias(bias_text: bool, bias_vae: bool):
+    """The fused row projection must support a bias on only one expert."""
+    dcfg = _parse_dtype("w16a16_bf16")
+    K = 32
+    N = 64
+    M = 5
+
+    with set_current_vllm_config(VllmConfig()):
+        text_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias_text,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        vae_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias_vae,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTRowParallelLinear(
+            K,
+            N,
+            bias=bias_text,
+            vae_bias=bias_vae,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        _sync_weights(text_linear, vae_linear, mot_linear)
+        text_idx = torch.tensor([0, 2], dtype=torch.long, device="cuda")
+        vae_idx = torch.tensor([1, 3, 4], dtype=torch.long, device="cuda")
+        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
+
+        with torch.no_grad():
+            ref = _reference_forward(x, text_idx, vae_idx, text_linear, vae_linear)
+            mot_out, _ = mot_linear(x, text_idx, vae_idx)
+
+        _check_and_report(ref, mot_out, f"Row one-sided bias text={bias_text} vae={bias_vae}")
+
+
 # =========================================================================
 #  Test: und-mode (text_indices=None) — falls back to standard forward
 # =========================================================================
@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the CustomOp platform-dispatch base class."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _NativeOnlyOp(CustomOp):
+    """An op that implements only the PyTorch-native path."""
+
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        return x * 2
+
+
+def test_forward_xpu_falls_back_to_native():
+    """XPU has no bespoke kernels, so the base class must reuse forward_native.
+
+    Without this fallback every CustomOp subclass has to define an identical
+    forward_xpu, and any that forgets crashes with NotImplementedError on XPU.
+    """
+    op = _NativeOnlyOp()
+    x = torch.randn(4, 8)
+
+    assert torch.equal(op.forward_xpu(x), op.forward_native(x))
+
+
+def test_forward_cuda_still_requires_an_implementation():
+    """The native fallback must not mask a missing CUDA kernel."""
+    op = _NativeOnlyOp()
+
+    with pytest.raises(NotImplementedError):
+        op.forward_cuda(torch.randn(4, 8))
+
+
+def test_mot_rmsnorm_dispatches_on_xpu_without_an_override():
+    """MoTRMSNorm relies on the inherited fallback (regression test)."""
+    from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+
+    assert "forward_xpu" not in MoTRMSNorm.__dict__
+
+    norm = MoTRMSNorm(16)
+    x = torch.randn(8, 16)
+    text_indices = torch.arange(0, 4)
+    vae_indices = torch.arange(4, 8)
+
+    assert torch.equal(norm.forward_xpu(x), norm.forward_native(x))
+    assert torch.equal(
+        norm.forward_xpu(x, text_indices, vae_indices),
+        norm.forward_native(x, text_indices, vae_indices),
+    )
@@ -29,7 +29,7 @@ class _DummyQuantMethod:
        return y


-def test_diffusion_base_linear_apply_multi_slice():
+def test_diffusion_base_linear_apply_multi_slice(monkeypatch):
    # Build a fake diffusion LoRA layer with 2 slices and rank=2.
    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
    layer.tp_size = 1
@@ -72,8 +72,18 @@ def test_diffusion_base_linear_apply_multi_slice():
    layer.lora_b_stacked = (b0, b1)
    layer.output_slices = out_slices

+    addmm_calls = []
+    original_addmm = torch.addmm
+
+    def _record_addmm(input, mat1, mat2, *, out=None):
+        addmm_calls.append((input, out))
+        return original_addmm(input, mat1, mat2, out=out)
+
+    monkeypatch.setattr(torch, "addmm", _record_addmm)
+
    x = torch.tensor([[1.0, 2.0, 3.0]])
-    out = layer.apply(x)
+    with torch.inference_mode():
+        out = layer.apply(x)

    # Base output is identity: [1,2,3]
    base_out = x @ base_weight.t()
@@ -87,6 +97,66 @@ def test_diffusion_base_linear_apply_multi_slice():
    delta1 = torch.tensor([[6.0]])
    expected = torch.cat([base_out[:, :2] + delta0, base_out[:, 2:3] + delta1], dim=-1)
    assert torch.allclose(out, expected)
+    assert len(addmm_calls) == 2
+    assert all(input is out for input, out in addmm_calls)
+
+
+def test_diffusion_base_linear_apply_preserves_autograd(monkeypatch):
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _DummyQuantMethod(torch.eye(2))
+
+    a = torch.tensor([[[[1.0, 2.0]]]], requires_grad=True)
+    b = torch.tensor([[[[3.0], [4.0]]]], requires_grad=True)
+    layer.lora_a_stacked = (a,)
+    layer.lora_b_stacked = (b,)
+    layer.output_slices = (2,)
+    layer._diffusion_lora_active_slices = (True,)
+
+    def _unexpected_addmm(*args, **kwargs):
+        raise AssertionError("the out= fast path must not run with autograd enabled")
+
+    monkeypatch.setattr(torch, "addmm", _unexpected_addmm)
+
+    x = torch.tensor([[2.0, 1.0]], requires_grad=True)
+    out = layer.apply(x)
+    assert torch.allclose(out, torch.tensor([[14.0, 17.0]]))
+
+    out.sum().backward()
+    assert torch.allclose(x.grad, torch.tensor([[8.0, 15.0]]))
+    assert torch.allclose(a.grad, torch.tensor([[[[14.0, 7.0]]]]))
+    assert torch.allclose(b.grad, torch.tensor([[[[4.0], [4.0]]]]))
+
+
+def test_diffusion_base_linear_apply_falls_back_for_mixed_output_dtype(monkeypatch):
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    class _Float64OutputQuantMethod(_DummyQuantMethod):
+        def apply(self, base_layer, x: torch.Tensor, bias: torch.Tensor | None):
+            return super().apply(base_layer, x, bias).to(torch.float64)
+
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _Float64OutputQuantMethod(torch.eye(2))
+    layer.lora_a_stacked = (torch.ones((1, 1, 1, 2)),)
+    layer.lora_b_stacked = (torch.tensor([[[[1.0], [2.0]]]]),)
+    layer.output_slices = (2,)
+    layer._diffusion_lora_active_slices = (True,)
+
+    def _unexpected_addmm(*args, **kwargs):
+        raise AssertionError("mixed output and LoRA dtypes must use the functional fallback")
+
+    monkeypatch.setattr(torch, "addmm", _unexpected_addmm)
+
+    with torch.inference_mode():
+        out = layer.apply(torch.tensor([[1.0, 2.0]]))
+
+    assert out.dtype == torch.float64
+    assert torch.allclose(out, torch.tensor([[4.0, 8.0]], dtype=torch.float64))


 def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

 import logging
 import os
@@ -23,9 +23,13 @@ NORM_EPS = 1e-5


@pytest.fixture(autouse=True)
-def _init_distributed():
+def _init_distributed(request):
    """Initialize the minimal single-rank distributed environment required by
    the vLLM parallel linear layers (tensor-parallel group must exist)."""
+    if request.node.name.startswith("test_real_rotary_"):
+        yield
+        return
+
    from vllm.distributed.parallel_state import (
        cleanup_dist_env_and_memory,
        init_distributed_environment,
@@ -59,6 +63,20 @@ def _force_default_gemm(monkeypatch):
    )


+@pytest.fixture(autouse=True)
+def _force_torch_sdpa():
+    """Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3)."""
+    from vllm_omni.diffusion.config import set_current_diffusion_config
+    from vllm_omni.diffusion.data import AttentionConfig
+
+    od_config = SimpleNamespace(
+        diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    with set_current_diffusion_config(od_config):
+        yield
+
+
 def _randomize_parameters(module: torch.nn.Module) -> None:
    """Fill parameters with small random values.

@@ -70,11 +88,11 @@ def _randomize_parameters(module: torch.nn.Module) -> None:
            param.uniform_(-0.02, 0.02)


-def _identity_rotary_emb(batch_size: int, seq_len: int) -> torch.Tensor:
-    """Complex rotary frequencies encoding a zero rotation."""
-    return torch.polar(
-        torch.ones(batch_size, seq_len, HEAD_DIM // 2),
-        torch.zeros(batch_size, seq_len, HEAD_DIM // 2),
+def _identity_rotary_emb(batch_size: int, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Real-valued rotary frequencies encoding a zero rotation."""
+    return (
+        torch.ones(batch_size, seq_len, HEAD_DIM),
+        torch.zeros(batch_size, seq_len, HEAD_DIM),
    )


@@ -119,6 +137,72 @@ def test_boogu_image_transformer_import():
    assert BooguImageTransformer2DModel is not None


+def test_real_rotary_emb_matches_complex_reference():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import apply_rotary_emb
+
+    torch.manual_seed(0)
+    x = torch.randn(2, 7, NUM_HEADS, HEAD_DIM)
+    angles = torch.randn(2, 7, HEAD_DIM // 2)
+    rotary_emb = (
+        angles.cos().repeat_interleave(2, dim=-1),
+        angles.sin().repeat_interleave(2, dim=-1),
+    )
+
+    actual = apply_rotary_emb(x, rotary_emb)
+    x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], HEAD_DIM // 2, 2))
+    freqs_cis = torch.polar(torch.ones_like(angles), angles).unsqueeze(2)
+    expected = torch.view_as_real(x_complex * freqs_cis).flatten(3).type_as(x)
+
+    torch.testing.assert_close(actual, expected)
+
+
+def test_real_rotary_frequency_tables_are_real_and_repeated():
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageDoubleStreamRotaryPosEmbed,
+    )
+
+    freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(AXES_DIM_ROPE, AXES_LENS, theta=10000)
+
+    assert len(freqs_real) == len(AXES_DIM_ROPE)
+    for (freqs_cos, freqs_sin), axis_dim, axis_len in zip(freqs_real, AXES_DIM_ROPE, AXES_LENS):
+        assert freqs_cos.shape == (axis_len, axis_dim)
+        assert freqs_sin.shape == (axis_len, axis_dim)
+        assert not freqs_cos.is_complex()
+        assert not freqs_sin.is_complex()
+        torch.testing.assert_close(freqs_cos[..., ::2], freqs_cos[..., 1::2])
+        torch.testing.assert_close(freqs_sin[..., ::2], freqs_sin[..., 1::2])
+
+
+def test_real_rotary_frequency_tables_use_float32_on_npu(monkeypatch):
+    from vllm_omni.diffusion.models.boogu_image import boogu_image_transformer
+
+    npu_platform = SimpleNamespace(is_npu=lambda: True, supports_float64=lambda: True)
+    monkeypatch.setattr(boogu_image_transformer, "current_omni_platform", npu_platform)
+
+    freqs_real = boogu_image_transformer.BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
+        AXES_DIM_ROPE, AXES_LENS, theta=10000
+    )
+
+    assert all(freqs.dtype == torch.float32 for pair in freqs_real for freqs in pair)
+
+
+def test_real_rotary_frequency_tables_use_float32_on_mps(monkeypatch):
+    from vllm_omni.diffusion.models.boogu_image import boogu_image_transformer
+
+    platform = SimpleNamespace(
+        is_npu=lambda: False,
+        supports_float64=lambda: True,
+    )
+    monkeypatch.setattr(boogu_image_transformer, "current_omni_platform", platform)
+    monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
+
+    freqs_real = boogu_image_transformer.BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
+        AXES_DIM_ROPE, AXES_LENS, theta=10000
+    )
+
+    assert all(freqs.dtype == torch.float32 for pair in freqs_real for freqs in pair)
+
+
 def test_single_stream_block_shape():
    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
        BooguImageTransformerBlock,
@@ -398,10 +482,10 @@ def test_transformer_forward_t2i_shape():
    timestep = torch.full((batch_size,), 0.5)
    instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
    instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
-    freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
+    freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(model.axes_dim_rope, model.axes_lens, theta=10000)

    with torch.no_grad():
-        out = model(latents, timestep, instruction_hidden_states, freqs_cis, instruction_attention_mask)
+        out = model(latents, timestep, instruction_hidden_states, freqs_real, instruction_attention_mask)

    assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
    assert torch.isfinite(out).all()
@@ -431,7 +515,7 @@ def test_transformer_forward_ti2i_shape():
    timestep = torch.full((batch_size,), 0.5)
    instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
    instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
-    freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
+    freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(model.axes_dim_rope, model.axes_lens, theta=10000)

    # One sample, one reference image (Boogu editing supports a single ref).
    ref_image_hidden_states = [[torch.randn(in_channels, ref_h, ref_w)]]
@@ -441,7 +525,7 @@ def test_transformer_forward_ti2i_shape():
            latents,
            timestep,
            instruction_hidden_states,
-            freqs_cis,
+            freqs_real,
            instruction_attention_mask,
            ref_image_hidden_states=ref_image_hidden_states,
        )
@@ -192,7 +192,6 @@ def test_constructor_weights_sources(boogu_pipeline):
        (DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"),
        (DiffusionParallelConfig(ulysses_degree=2), "none", "Sequence parallelism"),
        (DiffusionParallelConfig(ring_degree=2), "none", "Sequence parallelism"),
-        (DiffusionParallelConfig(cfg_parallel_size=2), "none", "CFG parallelism"),
        (
            DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=2),
            "none",
@@ -227,6 +226,26 @@ def test_constructor_rejects_unsupported_execution_modes(
    mock_dependencies["mllm_wrapper"].model.to.assert_not_called()


+@pytest.mark.parametrize("cfg_parallel_size", [2, 3])
+def test_constructor_accepts_cfg_parallel(mock_dependencies, cfg_parallel_size):
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    od_config = OmniDiffusionConfig(
+        model="dummy-boogu",
+        tf_model_config=TransformerConfig(params={}),
+        dtype=torch.float32,
+        parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size),
+    )
+
+    pipeline = BooguImagePipeline(od_config=od_config)
+
+    assert pipeline.od_config.parallel_config.cfg_parallel_size == cfg_parallel_size
+    assert hasattr(pipeline, "predict_noise_maybe_with_cfg")
+    assert hasattr(pipeline, "predict_noise_with_multi_branch_cfg")
+
+
 # ---------------------------------------------------------------------------
 # Prompt-encoding tests (deterministic fakes, no constructor)
 # ---------------------------------------------------------------------------
@@ -467,7 +486,8 @@ class _FakeTransformer:
        "reduce_type": "mean",
    }

-    def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
+    def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+        assert all(not freqs.is_complex() for pair in freqs_real for freqs in pair)
        return torch.zeros_like(latents)


@@ -525,6 +545,14 @@ def test_forward_returns_diffusion_output():

 def test_forward_cfg_off_when_guidance_one():
    pipeline = _make_forward_pipeline()
+    cfg_calls = []
+    original_predict_noise_maybe_with_cfg = pipeline.predict_noise_maybe_with_cfg
+
+    def record_cfg_call(**kwargs):
+        cfg_calls.append(kwargs)
+        return original_predict_noise_maybe_with_cfg(**kwargs)
+
+    pipeline.predict_noise_maybe_with_cfg = record_cfg_call
    req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2, guidance_scale=1.0)
    # guidance_scale=1.0 is falsy-adjacent but explicitly disables CFG; the
    # request layer would set guidance_scale_provided, so emulate that here.
@@ -534,6 +562,30 @@ def test_forward_cfg_off_when_guidance_one():

    # Only the positive prompt is encoded when CFG is off.
    assert len(pipeline.processor.calls) == 1
+    assert len(cfg_calls) == 2
+    assert all(call["do_true_cfg"] is False for call in cfg_calls)
+    assert all(call["negative_kwargs"] is None for call in cfg_calls)
+
+
+def test_double_guidance_combine_matches_legacy_formula():
+    pipeline = _make_forward_pipeline()
+    positive_with_reference = torch.randn(2, 4, 8, 8)
+    negative_with_reference = torch.randn(2, 4, 8, 8)
+    uncond = torch.randn(2, 4, 8, 8)
+    text_scale = 5.0
+    image_scale = 2.0
+
+    actual = pipeline.combine_multi_branch_cfg_noise(
+        [positive_with_reference, negative_with_reference, uncond],
+        {"text": text_scale, "image": image_scale},
+    )
+    legacy = (
+        positive_with_reference
+        + (text_scale - 1.0) * (positive_with_reference - negative_with_reference)
+        + (image_scale - 1.0) * (negative_with_reference - uncond)
+    )
+
+    torch.testing.assert_close(actual, legacy, rtol=0, atol=1e-5)


 # ---------------------------------------------------------------------------
@@ -801,6 +853,39 @@ def test_build_ref_latents_expands_per_output_and_handles_none():
    assert ref_latents[2] is None and ref_latents[3] is None


+def test_vae_attention_context_is_noop_on_cpu():
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
+        BooguImagePipeline,
+    )
+
+    with BooguImagePipeline._vae_attention_context(torch.device("cpu")):
+        pass
+
+
+def test_vae_attention_context_prioritizes_efficient_cuda(monkeypatch):
+    from vllm_omni.diffusion.models.boogu_image import pipeline_boogu_image
+
+    recorded = {}
+
+    def fake_sdpa_kernel(backends, *, set_priority):
+        recorded["backends"] = backends
+        recorded["set_priority"] = set_priority
+        return pipeline_boogu_image.nullcontext()
+
+    monkeypatch.setattr(pipeline_boogu_image, "sdpa_kernel", fake_sdpa_kernel)
+
+    with pipeline_boogu_image.BooguImagePipeline._vae_attention_context(torch.device("cuda")):
+        pass
+
+    assert recorded == {
+        "backends": [
+            pipeline_boogu_image.SDPBackend.EFFICIENT_ATTENTION,
+            pipeline_boogu_image.SDPBackend.MATH,
+        ],
+        "set_priority": True,
+    }
+
+
 # ---------------------------------------------------------------------------
 # Editing / TI2I: forward CFG branch selection
 # ---------------------------------------------------------------------------
@@ -811,9 +896,11 @@ class _RecordingRefTransformer(_FakeTransformer):

    def __init__(self):
        self.calls = []
+        self.instruction_embed_calls = []

-    def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
+    def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
        self.calls.append(kwargs.get("ref_image_hidden_states"))
+        self.instruction_embed_calls.append(instruction_embeds)
        return torch.zeros_like(latents)


@@ -875,6 +962,9 @@ def test_forward_ti2i_text_only_two_predictions_with_ref():
    # Text-only ti2i: cond+ref and neg+ref -> 2 predictions/step, both carry ref.
    assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
    assert all(ref is not None for ref in pipeline.transformer.calls)
+    for i in range(0, len(pipeline.transformer.instruction_embed_calls), 2):
+        positive, negative = pipeline.transformer.instruction_embed_calls[i : i + 2]
+        assert positive is not negative


 def test_forward_ti2i_double_guidance_three_predictions():
@@ -894,6 +984,11 @@ def test_forward_ti2i_double_guidance_three_predictions():
    per_step = [pipeline.transformer.calls[i : i + 3] for i in range(0, len(pipeline.transformer.calls), 3)]
    for step_calls in per_step:
        assert sum(ref is None for ref in step_calls) == 1
+    embed_calls = pipeline.transformer.instruction_embed_calls
+    for i in range(0, len(embed_calls), 3):
+        positive, negative_with_reference, uncond = embed_calls[i : i + 3]
+        assert positive is not negative_with_reference
+        assert negative_with_reference is uncond


 def test_forward_ti2i_image_only_two_predictions_drop_ref():
@@ -913,6 +1008,10 @@ def test_forward_ti2i_image_only_two_predictions_drop_ref():
    per_step = [pipeline.transformer.calls[i : i + 2] for i in range(0, len(pipeline.transformer.calls), 2)]
    for step_calls in per_step:
        assert sum(ref is None for ref in step_calls) == 1
+    embed_calls = pipeline.transformer.instruction_embed_calls
+    for i in range(0, len(embed_calls), 2):
+        positive_with_reference, positive_without_reference = embed_calls[i : i + 2]
+        assert positive_with_reference is positive_without_reference


 def test_forward_ti2i_no_guidance_single_prediction_with_ref():
@@ -1150,6 +1150,72 @@ def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
    assert cfg.fps == sp.resolved_frame_rate == 12.0


+def test_transfer_vae_executor_requires_distributed_vae() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+    executor = object()
+    pipeline.vae = SimpleNamespace(distributed_executor=executor, is_distributed_enabled=lambda: True)
+    assert pipeline._transfer_vae_executor() is executor
+
+    pipeline.vae = SimpleNamespace(distributed_executor=executor, is_distributed_enabled=lambda: False)
+    assert pipeline._transfer_vae_executor() is None
+
+
+def test_sync_transfer_overlap_slices_output_rank_and_broadcasts() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    class RecordingExecutor:
+        rank = 0
+
+        def __init__(self) -> None:
+            self.value = None
+
+        def broadcast_tensor(self, value):
+            self.value = value
+            return value
+
+    pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
+    output = torch.arange(1 * 3 * 5 * 2 * 2).reshape(1, 3, 5, 2, 2)
+    executor = RecordingExecutor()
+
+    overlap = pipeline._sync_transfer_overlap(
+        output,
+        overlap_frames=2,
+        reference_video=output,
+        vae_executor=executor,
+    )
+
+    assert overlap is not None
+    torch.testing.assert_close(overlap, output[:, :, -2:])
+    assert executor.value is overlap
+
+
+def test_sync_transfer_overlap_allocates_receive_buffer_on_non_output_rank() -> None:
+    from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
+
+    class ReceivingExecutor:
+        rank = 1
+
+        def broadcast_tensor(self, value):
+            assert value.shape == (1, 3, 2, 4, 5)
+            return torch.ones_like(value)
+
+    pipeline = SimpleNamespace(device=torch.device("cpu"), vae=SimpleNamespace(dtype=torch.float32))
+    reference = torch.zeros(1, 3, 6, 4, 5)
+
+    overlap = Cosmos3OmniDiffusersPipeline._sync_transfer_overlap(
+        pipeline,
+        torch.empty(0),
+        overlap_frames=2,
+        reference_video=reference,
+        vae_executor=ReceivingExecutor(),
+    )
+
+    assert overlap is not None
+    assert torch.equal(overlap, torch.ones_like(overlap))
+
+
 def test_transfer_edge_uses_rgb_canny(monkeypatch: pytest.MonkeyPatch) -> None:
    from vllm_omni.diffusion.models.cosmos3 import transfer

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
    build_hunyuan_diffusion_kv_requests,
    extract_hunyuan_prompt_inputs,
    hunyuan_num_image_tokens,
+    hunyuan_num_special_tokens,
    normalize_hunyuan_cot_text,
    prepare_hunyuan_layout,
 )
@@ -114,7 +115,7 @@ class _FakeImageProcessor:
        self.image_sizes: list[tuple[int, int]] = []
        self.vision_encoder_processor = SimpleNamespace(patch_size=1)

-    def build_image_info(self, image_size):
+    def build_image_info(self, image_size, **kwargs):
        self.image_sizes.append(image_size)
        return ImageInfo(
            image_type="gen_image",
@@ -125,6 +126,7 @@ class _FakeImageProcessor:
            image_token_length=self.image_token_length,
            base_size=1024,
            ratio_index=0,
+            **kwargs,
        )


@@ -166,6 +168,29 @@ def _request(
    )


+def test_distilled_layout_uses_embedded_cfg_and_meanflow_tokens() -> None:
+    tokenizer, image_processor = _components([12])
+    request = _request(guidance_scale=2.5)
+
+    prepared_layout = prepare_hunyuan_layout(
+        request,
+        tokenizer_wrapper=tokenizer,
+        image_processor=image_processor,
+        generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False),
+        image_base_size=1024,
+        cfg_distilled=True,
+        use_meanflow=True,
+    )
+
+    image_info = prepared_layout.generated_image_info
+    assert image_info.add_guidance_token
+    assert image_info.add_timestep_r_token
+    assert tokenizer.calls[0]["cfg_factor"] == 1
+    assert hunyuan_num_special_tokens(image_info) == 3
+    assert hunyuan_num_image_tokens(image_info) == 19
+    assert len(build_hunyuan_diffusion_kv_requests(request, prepared_layout)) == 1
+
+
 def test_builds_kv_request_lengths_without_model_execution() -> None:
    tokenizer, image_processor = _components([12])
    request = _request(guidance_scale=1.0)
@@ -385,9 +410,13 @@ def test_prepared_model_inputs_match_local_tokenization(
    torch.testing.assert_close(prepared_mask, local_mask)
    assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"]

-    num_image_tokens = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
-    local_inputs.update(attention_mask=local_mask, num_image_tokens=num_image_tokens)
-    prepared_inputs.update(attention_mask=prepared_mask, num_image_tokens=num_image_tokens)
+    image_info = prepared_layout.generated_image_info
+    generation_token_counts = {
+        "num_image_tokens": hunyuan_num_image_tokens(image_info),
+        "num_special_tokens": hunyuan_num_special_tokens(image_info),
+    }
+    local_inputs.update(attention_mask=local_mask, **generation_token_counts)
+    prepared_inputs.update(attention_mask=prepared_mask, **generation_token_counts)
    local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs)
    prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs)

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HunyuanImage3 prepare_inputs_for_generation to ensure
+distilled model parameters (guidance, timesteps_r) are properly passed."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+    HunyuanImage3Text2ImagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class MockHunyuanImage3Pipeline:
+    """Minimal mock that replicates prepare_inputs_for_generation method."""
+
+    def __init__(self):
+        # Mock config for distilled models
+        self.config = type("Config", (), {})()
+        self.hf_config = type("HFConfig", (), {"cfg_distilled": True, "use_meanflow": True})()
+
+    # Bind the real method from the pipeline class
+    from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
+        HunyuanImage3Pipeline as _Real,
+    )
+
+    prepare_inputs_for_generation = _Real.prepare_inputs_for_generation
+    _update_model_kwargs_for_generation = _Real._update_model_kwargs_for_generation
+
+
+class TestPrepareInputsForGenerationDistilled:
+    """Tests that distilled model parameters are passed through
+    prepare_inputs_for_generation."""
+
+    def setup_method(self):
+        self.pipeline = MockHunyuanImage3Pipeline()
+
+    def _create_mock_kwargs(self, include_distilled_params: bool = True) -> dict:
+        """Create mock kwargs for prepare_inputs_for_generation."""
+        kwargs = {
+            "position_ids": torch.tensor([[0, 1, 2]]),
+            "custom_pos_emb": (torch.tensor([[0.0]]), torch.tensor([[1.0]])),
+            "mode": "gen_image",
+            "images": torch.randn(1, 4, 64, 64),
+            "image_mask": torch.ones(1, 64, 64),
+            "timestep": torch.tensor([0.5]),
+            "gen_timestep_scatter_index": torch.tensor([[5]]),
+            "cond_vae_images": None,
+            "cond_timestep": None,
+            "cond_vae_image_mask": None,
+            "cond_vit_images": None,
+            "cond_vit_image_mask": None,
+            "vit_kwargs": None,
+            "cond_timestep_scatter_index": None,
+            "query_lens": [10],
+            "seq_lens": [100],
+            "num_image_tokens": 4096,
+            "num_special_tokens": 3,
+            "ar_kv_reuse_len": 0,
+            "full_attn_spans": None,
+            "use_cache": True,
+        }
+        if include_distilled_params:
+            # CFG distilled model parameters
+            kwargs["guidance"] = torch.tensor([2500.0])
+            kwargs["guidance_scatter_index"] = torch.tensor([[4]])
+            # MeanFlow distilled model parameters
+            kwargs["timesteps_r"] = torch.tensor([0.3])
+            kwargs["timesteps_r_scatter_index"] = torch.tensor([[3]])
+        else:
+            kwargs["guidance"] = None
+            kwargs["guidance_scatter_index"] = None
+            kwargs["timesteps_r"] = None
+            kwargs["timesteps_r_scatter_index"] = None
+        return kwargs
+
+    def test_guidance_passed_to_model_inputs(self):
+        """Test that guidance is passed through prepare_inputs_for_generation."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "guidance" in model_inputs, "guidance should be in model_inputs"
+        assert model_inputs["guidance"] is not None, "guidance should not be None"
+        assert torch.equal(model_inputs["guidance"], kwargs["guidance"])
+
+    def test_guidance_scatter_index_passed_to_model_inputs(self):
+        """Test that guidance_scatter_index is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "guidance_scatter_index" in model_inputs
+        assert model_inputs["guidance_scatter_index"] is not None
+        assert torch.equal(model_inputs["guidance_scatter_index"], kwargs["guidance_scatter_index"])
+
+    def test_timesteps_r_passed_to_model_inputs(self):
+        """Test that timesteps_r is passed through for MeanFlow models."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "timesteps_r" in model_inputs, "timesteps_r should be in model_inputs"
+        assert model_inputs["timesteps_r"] is not None
+        assert torch.equal(model_inputs["timesteps_r"], kwargs["timesteps_r"])
+
+    def test_timesteps_r_scatter_index_passed_to_model_inputs(self):
+        """Test that timesteps_r_scatter_index is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "timesteps_r_scatter_index" in model_inputs
+        assert model_inputs["timesteps_r_scatter_index"] is not None
+        assert torch.equal(model_inputs["timesteps_r_scatter_index"], kwargs["timesteps_r_scatter_index"])
+
+    def test_num_special_tokens_passed_to_model_inputs(self):
+        """Test that num_special_tokens is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "num_special_tokens" in model_inputs
+        assert model_inputs["num_special_tokens"] == kwargs["num_special_tokens"]
+
+    def test_decode_update_keeps_all_special_token_positions(self):
+        """Test decode step uses timestep, guidance, and timestep_r positions."""
+        model_kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        model_kwargs.update(
+            {
+                "tokenizer_output": object(),
+                "attention_mask": torch.ones(1, 1, 7, 7, dtype=torch.bool),
+                "image_mask": torch.tensor([[False, False, False, False, True, True, True]]),
+                "gen_timestep_scatter_index": torch.tensor([[1]]),
+                "guidance_scatter_index": torch.tensor([[2]]),
+                "timesteps_r_scatter_index": torch.tensor([[3]]),
+            }
+        )
+
+        updated = self.pipeline._update_model_kwargs_for_generation({}, model_kwargs)
+
+        assert torch.equal(updated["position_ids"], torch.tensor([[1, 2, 3, 4, 5, 6]]))
+        assert updated["num_special_tokens"] == 3
+        assert updated["guidance_scatter_index"] is model_kwargs["guidance_scatter_index"]
+        assert updated["timesteps_r_scatter_index"] is model_kwargs["timesteps_r_scatter_index"]
+        assert updated["attention_mask"].shape == (1, 1, 6, 7)
+
+    def test_distilled_params_none_when_not_provided(self):
+        """Test that distilled params are None when not provided."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=False)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        # All distilled params should be present but None
+        assert "guidance" in model_inputs
+        assert model_inputs["guidance"] is None
+        assert "guidance_scatter_index" in model_inputs
+        assert model_inputs["guidance_scatter_index"] is None
+        assert "timesteps_r" in model_inputs
+        assert model_inputs["timesteps_r"] is None
+        assert "timesteps_r_scatter_index" in model_inputs
+        assert model_inputs["timesteps_r_scatter_index"] is None
+
+    def test_all_distilled_keys_present_in_model_inputs(self):
+        """Test that all 4 distilled keys are always present (even if None)."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=False)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        required_distilled_keys = [
+            "guidance",
+            "guidance_scatter_index",
+            "timesteps_r",
+            "timesteps_r_scatter_index",
+        ]
+        for key in required_distilled_keys:
+            assert key in model_inputs, f"{key} must be in model_inputs dict"
+
+
+def test_ar_kv_reuse_rebases_distilled_scatter_indices():
+    pipeline = object.__new__(HunyuanImage3Text2ImagePipeline)
+    model_kwargs = {
+        "query_lens": [8],
+        "attention_mask": torch.ones(1, 1, 8, 12, dtype=torch.bool),
+        "position_ids": torch.arange(8).unsqueeze(0),
+        "image_mask": torch.ones(1, 8, dtype=torch.bool),
+        "gen_timestep_scatter_index": torch.tensor([[5]]),
+        "guidance_scatter_index": torch.tensor([[6]]),
+        "timesteps_r_scatter_index": torch.tensor([[7]]),
+    }
+
+    truncated = pipeline._truncate_reused_prefix(
+        input_ids=torch.arange(8).unsqueeze(0),
+        model_kwargs=model_kwargs,
+        positive_reuse_len=3,
+    )
+
+    assert truncated.tolist() == [[3, 4, 5, 6, 7]]
+    assert model_kwargs["gen_timestep_scatter_index"].tolist() == [[2]]
+    assert model_kwargs["guidance_scatter_index"].tolist() == [[3]]
+    assert model_kwargs["timesteps_r_scatter_index"].tolist() == [[4]]
+
+
+def test_cfg_parallel_splits_meanflow_scatter_indices():
+    model_kwargs = {
+        "timesteps_r_scatter_index": torch.tensor([[1], [2], [3], [4]]),
+    }
+
+    HunyuanImage3Text2ImagePipeline._split_model_kwargs_for_cfg_parallel(
+        model_kwargs,
+        batch_size=2,
+        cfg_rank=1,
+    )
+
+    assert model_kwargs["timesteps_r_scatter_index"].tolist() == [[3], [4]]
@@ -38,6 +38,7 @@ def _pipeline():
        cache_backend=None,
        diffusion_kv_cache_skip_step_indices=None,
    )
+    pipeline.hf_config = SimpleNamespace(cfg_distilled=False, use_meanflow=False)
    pipeline._pipeline = SimpleNamespace()
    return pipeline

@@ -451,3 +452,44 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
    assert states[1].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (2, 1, 2, 6)
    assert states[0].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(2, 4)], [(2, 4)]]
    assert states[1].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(4, 6)], [(4, 6)]]
+
+
+def test_distilled_step_supplies_guidance_and_meanflow_timestep(monkeypatch):
+    pipeline = _pipeline()
+    pipeline.hf_config = SimpleNamespace(cfg_distilled=True, use_meanflow=True)
+    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
+    state = _state("distilled", 1)
+    state.latents = torch.zeros(1, 1)
+    state.extra[_STEP_GUIDANCE_SCALE] = 2.5
+    state.extra[_STEP_MODEL_KWARGS].update(
+        {
+            "attention_mask": torch.ones(1, 1, 2, 4, dtype=torch.bool),
+            "full_attn_spans": [[(2, 4)]],
+        }
+    )
+    state.extra[_STEP_PROMPT_KV] = [
+        {
+            "key": torch.zeros(1, 2, 1, 1),
+            "value": torch.zeros(1, 2, 1, 1),
+            "lens": torch.tensor([2]),
+        }
+    ]
+    state.scheduler = SimpleNamespace(get_timestep_r=lambda _timestep: torch.tensor(0.25))
+    captured = {}
+
+    pipeline._restore_prompt_kv_cache = lambda *_args: None
+
+    def fake_prepare_inputs(input_ids, images, timestep, **model_kwargs):
+        del input_ids, images, timestep
+        captured.update(model_kwargs)
+        return {"model_kwargs": model_kwargs}
+
+    pipeline.prepare_inputs_for_generation = fake_prepare_inputs
+    pipeline.forward_call = lambda **_kwargs: {"diffusion_prediction": torch.tensor([[1.0]])}
+    pipeline._update_model_kwargs_for_generation = lambda _output, model_kwargs: model_kwargs
+
+    output = pipeline.denoise_step(InputBatch.make_batch([state]))
+
+    torch.testing.assert_close(output, torch.tensor([[1.0]]))
+    torch.testing.assert_close(captured["guidance"], torch.tensor([2500.0], dtype=torch.bfloat16))
+    torch.testing.assert_close(captured["timesteps_r"], torch.tensor([0.25]))
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

 from types import SimpleNamespace

@@ -182,7 +182,8 @@ def _fake_ltx_parallel_attention():
            return "FLASH_ATTN"

        @staticmethod
-        def supports_attention_mask():
+        def supports_attention_mask(attention_spec=None):
+            del attention_spec
            return True

    def run(name):
@@ -795,6 +795,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():


 def test_packed_attention_skips_mask_for_packed_mask_free_backend():
+    from vllm_omni.diffusion.attention.backends.abstract import (
+        VideoTokenLayout,
+        VideoTokenSpan,
+    )
    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
        MiniMaxH3Attention,
    )
@@ -821,6 +825,10 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
    torch.nn.Module.__init__(attention)
    attention.attention = FakeAttention()
    q = torch.randn(8, 2, 4)
+    video_layout = VideoTokenLayout(
+        used_len=5,
+        video_spans=(VideoTokenSpan(start=3, latent_grid=(1, 1, 2), role="target"),),
+    )

    attention._run_packed_attention(
        q,
@@ -829,6 +837,9 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
        max_seqlen=5,
        packed_total=8,
+        video_layout=video_layout,
+        vsa_prefix_segments=(1, 2),
+        gate_compress=torch.ones_like(q),
    )

    metadata = attention.attention.metadata
@@ -836,6 +847,11 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
    assert metadata.attn_mask is None
    assert metadata.extra["valid_kv_length"] == 5
    assert metadata.extra["npu_attn_varlen"] is True
+    assert metadata.video_layout is video_layout
+    assert metadata.extra["vsa_h3_prefix_segments"] == (1, 2)
+    assert "vsa_h3_video_shape" not in metadata.extra
+    assert "vsa_h3_target_start" not in metadata.extra
+    assert metadata.extra["gate_compress"].shape == (1, *q.shape)
    packed_padding = metadata.packed_padding
    assert packed_padding is not None
    assert packed_padding.q_length == 5
@@ -437,6 +437,83 @@ def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path):
        _fusion(tmp_path).check_serving_contract(partition="ref2va", **contract)


+def test_a_vsa_variant_accepts_pure_ulysses_and_refuses_ring(tmp_path):
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    od_config = SimpleNamespace(
+        diffusion_attention_backend="FASTVIDEO_VSA",
+        parallel_config=SimpleNamespace(sequence_parallel_size=8, ulysses_degree=8, ring_degree=1),
+    )
+    fusion = _load(sparse.parent)
+    fusion.check_serving_contract(partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0)
+
+    od_config.parallel_config.ring_degree = 2
+    with pytest.raises(ValueError, match="pure Ulysses"):
+        fusion.check_serving_contract(partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0)
+
+
+def test_a_vsa_variant_accepts_normalized_attention_config(tmp_path):
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    od_config = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(default=SimpleNamespace(backend="FASTVIDEO_VSA")),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    _load(sparse.parent).check_serving_contract(
+        partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0
+    )
+
+
+def test_a_vsa_variant_reads_the_backend_the_dit_will_actually_resolve(tmp_path):
+    # The 50 DiT blocks carry attention role "self", so a per_role entry - not
+    # the default - decides whether the compression gates ever reach a sparse
+    # kernel.
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    fusion = _load(sparse.parent)
+    contract = {"partition": "fl2va", "video_shift": 12.0, "audio_shift": 3.0}
+
+    dense_dit = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(
+            default=SimpleNamespace(backend="FASTVIDEO_VSA"),
+            per_role={"self": SimpleNamespace(backend="TRTLLM_ATTN")},
+        ),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    with pytest.raises(ValueError, match="TRTLLM_ATTN"):
+        fusion.check_serving_contract(od_config=dense_dit, **contract)
+
+    sparse_dit = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(
+            default=None,
+            per_role={"self": SimpleNamespace(backend="FASTVIDEO_VSA")},
+        ),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    fusion.check_serving_contract(od_config=sparse_dit, **contract)
+
+
+def test_a_gate_that_never_reached_the_model_is_refused(tmp_path):
+    # load_weights only logs a skip for a parameter the model does not have, so
+    # yielding the gate is not evidence it arrived.
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    fusion = _load(sparse.parent)
+    checkpoint = [
+        ("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN))),
+        ("blocks.0.attn.out_proj.weight", torch.zeros((_HIDDEN, _INNER))),
+        ("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))),
+        ("blocks.0.mlp.fc2.weight", torch.zeros((_HIDDEN, _FFN))),
+    ]
+    gate = "blocks.0.attn.to_gate_compress.weight"
+    streamed = {name for name, _ in fusion.apply(checkpoint)}
+    assert gate in streamed
+
+    with pytest.raises(FastH3AdapterError, match="never reached the model"):
+        fusion.validate_fully_applied(streamed - {gate})
+    fusion.validate_fully_applied(streamed)
+
+
 def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path):
    # Only low-rank factors are placed into H3's fused QKV and gate/up layouts,
    # so a .diff aimed at one would otherwise be added transposed.
@@ -67,6 +67,33 @@ def test_packed_attention_is_a_regional_compile_boundary():
    assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)


+def test_denoise_branch_keeps_fast_h3_prefix_geometry_model_local():
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
+        MiniMaxH3DenoiseBranch,
+    )
+
+    packed = {
+        "seq_len": torch.tensor(9),
+        "img_pos": torch.arange(4),
+        "audio_pos": torch.empty(0, dtype=torch.long),
+        "text_pos": torch.arange(2),
+        "update_mask": torch.ones(4, dtype=torch.bool),
+        "cu_seqlens": torch.tensor([0, 9, 9], dtype=torch.int32),
+        "img_position_ids": torch.zeros(9, 3, dtype=torch.long),
+        "video_spans": ({"start": 5, "latent_grid": (1, 2, 2), "role": "target"},),
+    }
+    branch = MiniMaxH3DenoiseBranch(
+        packed=packed,
+        text_embeddings=torch.zeros(2, 4),
+        token_tags=torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2]),
+        device=torch.device("cpu"),
+    )
+
+    layout = branch.static_kwargs["video_token_layout"]
+    assert not hasattr(layout, "prefix_segments")
+    assert branch.static_kwargs["packed_seq_params"]["vsa_prefix_segments"] == (2, 3)
+
+
 def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
    from vllm_omni.diffusion.layers.rope import RotaryEmbedding
    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
@@ -0,0 +1,346 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from types import SimpleNamespace
+
+import cache_dit
+import pytest
+import torch
+from cache_dit import ForwardPattern
+from torch import nn
+
+import vllm_omni.diffusion.cache.cachedit.backend as cachedit_backend_module
+import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
+from tests.diffusion.models.sana_video.test_transformer_sana_video import (  # noqa: F401
+    _TINY_CONFIG,
+    _init_distributed,
+)
+from tests.diffusion.offloader.test_layerwise_backend import DummyEvent, DummyStream, dummy_stream
+from vllm_omni.diffusion.attention import selector as attention_selector
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend
+from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
+from vllm_omni.diffusion.models.sana_video import pipeline_sana_video as pipeline_module
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import (
+    SanaVideoPipeline,
+    _validate_cache_offload_parallelism,
+)
+from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformer3DModel
+from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
+from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@dataclass
+class _PipelineWithTransformer:
+    transformer: nn.Module
+
+
+def _tiny_transformer(monkeypatch) -> SanaVideoTransformer3DModel:
+    monkeypatch.setattr(
+        attention_selector,
+        "_cached_get_backend_cls",
+        lambda *_args, **_kwargs: SDPABackend,
+    )
+    return SanaVideoTransformer3DModel(**(_TINY_CONFIG | {"num_layers": 2}))
+
+
+def _record_cache_adapters(monkeypatch):
+    enabled_adapters = []
+    original_enable_cache = cachedit_backend_module.cache_dit.enable_cache
+
+    def record_enable_cache(block_adapter, **kwargs):
+        enabled_adapters.append(block_adapter)
+        return original_enable_cache(block_adapter, **kwargs)
+
+    monkeypatch.setattr(cachedit_backend_module.cache_dit, "enable_cache", record_enable_cache)
+    return enabled_adapters
+
+
+def test_sana_video_declares_cache_and_layerwise_metadata():
+    adapter_config = SanaVideoTransformer3DModel._cache_dit_adapter_config
+
+    assert adapter_config.block_forward_patterns == {
+        "transformer_blocks": ForwardPattern.Pattern_3,
+    }
+    assert adapter_config.has_separate_cfg is False
+    assert adapter_config.cached_adapter_cls is None
+    assert adapter_config.check_forward_pattern is True
+    assert SanaVideoTransformer3DModel._layerwise_offload_blocks_attrs == ["transformer_blocks"]
+    assert SanaVideoPipeline.default_num_inference_steps == 50
+
+
+def test_two_layer_tiny_transformer_enables_cache_dit(monkeypatch):
+    transformer = _tiny_transformer(monkeypatch)
+    pipeline = _PipelineWithTransformer(transformer=transformer)
+    enabled_adapters = _record_cache_adapters(monkeypatch)
+    backend = CacheDiTBackend()
+
+    try:
+        backend.enable(pipeline)
+
+        assert backend.is_enabled()
+        assert transformer._is_cached is True
+        assert len(enabled_adapters) == 1
+        selected_blocks = cache_dit.BlockAdapter.flatten(enabled_adapters[0].blocks)
+        assert selected_blocks == [transformer.transformer_blocks]
+        assert len(selected_blocks[0]) == 2
+    finally:
+        if enabled_adapters:
+            cache_dit.disable_cache(enabled_adapters[0])
+
+
+def test_second_engine_enables_cache_dit_after_undisabled_first(monkeypatch):
+    """An engine never disables Cache-DiT on shutdown, and cache_dit marks the
+    shared wrapper-pipe class as cached; a second engine in the same process
+    must still get a full enable with a working cache context."""
+    enabled_adapters = _record_cache_adapters(monkeypatch)
+    first = _PipelineWithTransformer(transformer=_tiny_transformer(monkeypatch))
+    second = _PipelineWithTransformer(transformer=_tiny_transformer(monkeypatch))
+    backend = CacheDiTBackend()
+
+    try:
+        CacheDiTBackend().enable(first)
+        backend.enable(second)
+
+        assert backend.is_enabled()
+        assert second.transformer._is_cached is True
+    finally:
+        for adapter in enabled_adapters:
+            cache_dit.disable_cache(adapter)
+
+
+def test_cache_dit_pattern_mismatch_fails_closed(monkeypatch):
+    transformer = _tiny_transformer(monkeypatch)
+    transformer.transformer_blocks[1].forward = lambda unexpected_input: unexpected_input
+    pipeline = _PipelineWithTransformer(transformer=transformer)
+    backend = CacheDiTBackend()
+
+    with pytest.raises(AssertionError, match="No block forward pattern matched"):
+        backend.enable(pipeline)
+
+    assert backend.is_enabled() is False
+    assert not getattr(transformer, "_is_cached", False)
+
+
+@pytest.mark.parametrize(
+    "feature_flags",
+    [
+        {"cache_backend": "cache_dit"},
+        {"enable_cpu_offload": True},
+        {"enable_layerwise_offload": True},
+        {"enable_distributed_layerwise_offload": True},
+    ],
+)
+@pytest.mark.parametrize(
+    "parallel_config",
+    [
+        DiffusionParallelConfig(tensor_parallel_size=2),
+        DiffusionParallelConfig(cfg_parallel_size=2),
+        DiffusionParallelConfig(ulysses_degree=2),
+    ],
+)
+def test_cache_offload_distributed_combinations_fail_closed(feature_flags, parallel_config):
+    config = OmniDiffusionConfig(parallel_config=parallel_config, **feature_flags)
+
+    with pytest.raises(NotImplementedError, match="supported only with TP1, CFG1, and SP1"):
+        _validate_cache_offload_parallelism(config)
+
+
+def test_sana_video_rejects_unvalidated_cache_backends():
+    with pytest.raises(NotImplementedError, match="Cache backend 'tea_cache' is not supported"):
+        _validate_cache_offload_parallelism(OmniDiffusionConfig(cache_backend="tea_cache"))
+
+
+def test_sana_video_rejects_cache_dit_with_distributed_layerwise_offload():
+    config = OmniDiffusionConfig(cache_backend="cache_dit", enable_distributed_layerwise_offload=True)
+
+    with pytest.raises(NotImplementedError, match="Cache-DiT with distributed layerwise offload"):
+        _validate_cache_offload_parallelism(config)
+
+
+def test_sana_video_accepts_distributed_layerwise_offload():
+    _validate_cache_offload_parallelism(OmniDiffusionConfig(enable_distributed_layerwise_offload=True))
+
+
+@pytest.mark.parametrize(
+    ("od_config", "match"),
+    [
+        (
+            OmniDiffusionConfig(
+                cache_backend="cache_dit",
+                parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
+            ),
+            "tensor_parallel_size",
+        ),
+        (
+            OmniDiffusionConfig(cache_backend="cache_dit", enable_distributed_layerwise_offload=True),
+            "Cache-DiT with distributed layerwise offload",
+        ),
+    ],
+)
+def test_validation_fails_before_component_loading(monkeypatch, od_config, match):
+    load_calls = []
+    monkeypatch.setattr(pipeline_module, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(
+        SanaVideoPipeline,
+        "_load_components",
+        lambda *args, **kwargs: load_calls.append((args, kwargs)),
+    )
+
+    with pytest.raises(NotImplementedError, match=match):
+        SanaVideoPipeline(od_config=od_config)
+
+    assert load_calls == []
+
+
+class _TrackingModule:
+    def __init__(self):
+        self.to_calls = []
+
+    def to(self, *args, **kwargs):
+        self.to_calls.append((args, kwargs))
+        return self
+
+
+def test_component_loading_uses_loader_device_not_runtime_device(monkeypatch):
+    component_load_device = torch.device("cpu")
+    text_encoder = _TrackingModule()
+    vae = _TrackingModule()
+    transformer = _TrackingModule()
+    tokenizer = object()
+    scheduler = object()
+    loaded_components = iter([text_encoder, vae])
+
+    monkeypatch.setattr(torch, "get_default_device", lambda: component_load_device)
+    monkeypatch.setattr(pipeline_module, "prefetch_subfolders", lambda *args, **kwargs: None)
+    monkeypatch.setattr(pipeline_module, "_load_sana_tokenizer", lambda *args, **kwargs: tokenizer)
+    monkeypatch.setattr(
+        pipeline_module,
+        "from_pretrained_with_prefetch",
+        lambda *args, **kwargs: next(loaded_components),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "_load_json",
+        lambda _model, filename, _local: ({"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {}),
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "_resolve_vae_class_and_dtype",
+        lambda *_args: (SimpleNamespace(from_pretrained=None), torch.float32),
+    )
+    monkeypatch.setattr(
+        pipeline_module.SanaVideoTransformer3DModel,
+        "from_config",
+        classmethod(lambda _cls, _config: transformer),
+    )
+    monkeypatch.setattr(
+        pipeline_module.DPMSolverMultistepScheduler,
+        "from_pretrained",
+        lambda *args, **kwargs: scheduler,
+    )
+
+    pipeline = object.__new__(SanaVideoPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cuda:7")
+    pipeline.weights_sources = []
+    config = OmniDiffusionConfig(model="local-sana-video", dtype=torch.bfloat16)
+
+    loaded = pipeline._load_components(config, prefix="")
+
+    assert loaded == (tokenizer, text_encoder, vae, transformer, scheduler)
+    assert pipeline.device == torch.device("cuda:7")
+    assert text_encoder.to_calls == [((component_load_device,), {})]
+    assert vae.to_calls == [((component_load_device,), {})]
+    assert transformer.to_calls == [
+        ((), {"dtype": torch.bfloat16, "device": component_load_device}),
+    ]
+
+
+def _patch_layerwise_platform(monkeypatch) -> None:
+    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 _tiny_pipeline(monkeypatch) -> SanaVideoPipeline:
+    pipeline = object.__new__(SanaVideoPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.transformer = _tiny_transformer(monkeypatch)
+    pipeline.text_encoder = nn.Linear(2, 2)
+    pipeline.vae = nn.Linear(2, 2)
+    return pipeline
+
+
+def _enable_layerwise_offload(pipeline: SanaVideoPipeline) -> LayerWiseOffloadBackend:
+    backend = LayerWiseOffloadBackend(
+        OffloadConfig(strategy=OffloadStrategy.LAYER_WISE, pin_cpu_memory=False),
+        device=torch.device("cpu"),
+    )
+    backend.enable(pipeline)
+    return backend
+
+
+def test_two_layer_offload_recovers_from_cached_block_skip_and_cleans_up(monkeypatch):
+    _patch_layerwise_platform(monkeypatch)
+    pipeline = _tiny_pipeline(monkeypatch)
+    discovered = ModuleDiscovery.discover(pipeline)
+    assert discovered.dits == [pipeline.transformer]
+    assert discovered.encoders == [pipeline.text_encoder]
+    assert discovered.vaes == [pipeline.vae]
+
+    backend = _enable_layerwise_offload(pipeline)
+    blocks = list(pipeline.transformer.transformer_blocks)
+    assert backend.is_enabled()
+    assert len(blocks) == 2
+
+    block_1_hook = blocks[1]._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME)
+    assert block_1_hook is not None
+    assert block_1_hook.is_materialized is False
+
+    # Cache-DiT may skip block 0, leaving block 1 without the normal prefetch.
+    block_1_hook.pre_forward(blocks[1])
+    assert block_1_hook.is_materialized is True
+    block_1_hook.post_forward(blocks[1], None)
+    assert block_1_hook.is_materialized is False
+
+    backend.disable()
+    assert backend.is_enabled() is False
+    for block in blocks:
+        assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is None
+
+
+def test_layerwise_offload_then_cache_dit_install_and_cleanup(monkeypatch):
+    _patch_layerwise_platform(monkeypatch)
+    pipeline = _tiny_pipeline(monkeypatch)
+    offload_backend = _enable_layerwise_offload(pipeline)
+    enabled_adapters = _record_cache_adapters(monkeypatch)
+    cache_backend = CacheDiTBackend()
+    blocks = list(pipeline.transformer.transformer_blocks)
+
+    try:
+        cache_backend.enable(pipeline)
+
+        assert offload_backend.is_enabled()
+        assert cache_backend.is_enabled()
+        assert pipeline.transformer._is_cached is True
+        for block in blocks:
+            assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is not None
+    finally:
+        if enabled_adapters:
+            cache_dit.disable_cache(enabled_adapters[0])
+        offload_backend.disable()
+
+    assert not getattr(pipeline.transformer, "_is_cached", False)
+    assert offload_backend.is_enabled() is False
+    for block in blocks:
+        assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is None
@@ -0,0 +1,582 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""CFG dispatch tests for the native SANA-Video pipelines.
+
+These exercise the CFG world-size-1 batch path, the CFG world-size-2 dispatch,
+and the fp32 combine semantics on CPU / mocked distributed. The real all-gather
+path is covered by the multi-process suite, not here.
+"""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import SanaVideoPipeline
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v import SanaImageToVideoPipeline
+from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+    SanaVideoTransformerOutput,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def test_predict_noise_calls_transformer_with_return_dict_false():
+    """predict_noise must call the transformer with return_dict=False.
+
+    The transformer returns a SanaVideoTransformerOutput dataclass when
+    return_dict is True (its default), which is not subscriptable. The CFG
+    mixin's default predict_noise does ``result[0]`` and would therefore
+    raise TypeError on SANA; the pipeline override must force return_dict=False
+    so the tuple form is returned and indexed.
+    """
+    captured = {}
+    expected = torch.full((1, 2, 3), 7.0)
+
+    class _SpyTransformer:
+        def __call__(self, *args, return_dict=True, **kwargs):
+            captured["return_dict"] = return_dict
+            # Mirror the real transformer: tuple on False, dataclass on True.
+            return (expected,) if not return_dict else SanaVideoTransformerOutput(sample=expected)
+
+    pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)  # bypass heavy __init__
+    pipe.transformer = _SpyTransformer()
+
+    out = pipe.predict_noise(
+        torch.zeros(1, 2, 3),
+        encoder_hidden_states=torch.zeros(1, 1, 4),
+        encoder_attention_mask=None,
+        timestep=torch.zeros(1),
+    )
+
+    assert captured["return_dict"] is False
+    assert torch.equal(out, expected)
+
+
+class _IdentityScheduler:
+    """Scheduler stub whose step is a no-op, so diffuse tests isolate the CFG
+    dispatch from real scheduler math."""
+
+    order = 1
+
+    def step(self, noise_pred, t, latents, **kwargs):
+        return (latents,)
+
+
+def _set_cfg_world_size(monkeypatch, n):
+    """Pipelines always run inside an initialized CFG group; unit tests stub its
+    world size instead of spinning up torch.distributed. Patch both pipeline
+    modules since each resolves the symbol in its own namespace."""
+    import vllm_omni.diffusion.distributed.cfg_parallel as cfg
+    import vllm_omni.diffusion.models.sana_video.pipeline_sana_video as t2v
+    import vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v as i2v
+
+    monkeypatch.setattr(t2v, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
+    monkeypatch.setattr(i2v, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
+    monkeypatch.setattr(cfg, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
+    monkeypatch.setattr(t2v, "model_parallel_is_initialized", lambda: True, raising=False)
+    monkeypatch.setattr(i2v, "model_parallel_is_initialized", lambda: True, raising=False)
+
+
+def _bare_t2v_pipeline(transformer, guidance_scale):
+    pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)  # bypass heavy __init__
+    pipe.transformer = transformer
+    pipe.scheduler = _IdentityScheduler()
+    pipe._interrupt = False
+    pipe._guidance_scale = guidance_scale
+    return pipe
+
+
+def test_cfg_world_size_1_does_single_batch2_forward_per_step(monkeypatch):
+    """CFG world size 1 with guidance must do exactly one batch-2 transformer
+    forward per step (the batch path), not the mixin's sequential two-forward
+    path."""
+    _set_cfg_world_size(monkeypatch, 1)
+    batch_dims = []
+
+    class _SpyTransformer:
+        dtype = torch.float32
+
+        def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
+            batch_dims.append(hidden_states.shape[0])
+            # out_channels == in_channels here (no learned-sigma split).
+            return (torch.zeros_like(hidden_states),)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
+
+    out = pipe.diffuse(
+        latents=torch.randn(b, 4, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=None,
+    )
+
+    assert len(batch_dims) == 1, f"expected 1 forward this step, got {len(batch_dims)}"
+    assert batch_dims[0] == 2 * b, f"expected batch-2 (neg+pos) forward, got batch {batch_dims[0]}"
+    assert out.shape == (b, 4, 2, 4, 4)
+
+
+class _DeterministicTransformer:
+    """Learned-sigma transformer whose output depends on the row's encoder
+    embeds and timestep, so uncond != text and the CFG combine is non-trivial."""
+
+    dtype = torch.float32
+
+    def __init__(self, latent_channels):
+        self.out_channels = 2 * latent_channels  # learned sigma
+
+    def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
+        b, c = hidden_states.shape[0], hidden_states.shape[1]
+        sig = encoder_hidden_states.float().mean(dim=(1, 2)).view(b, 1, 1, 1, 1)
+        base = hidden_states.float() + sig + timestep.float().view(b, 1, 1, 1, 1)
+        return (base.repeat(1, self.out_channels // c, 1, 1, 1),)
+
+
+class _LinearScheduler:
+    """Deterministic stepping scheduler so multi-step divergence compounds."""
+
+    order = 1
+
+    def step(self, noise_pred, t, latents, return_dict=False, **kwargs):
+        return (latents - 0.1 * noise_pred,)
+
+
+def _pr1_reference_loop(
+    transformer,
+    scheduler,
+    latents,
+    timesteps,
+    prompt_embeds,
+    prompt_mask,
+    negative_embeds,
+    negative_mask,
+    guidance_scale,
+    latent_channels,
+    dtype,
+):
+    """The original CFG1 denoise math, verbatim, as the parity oracle."""
+    prompt_embeds = torch.cat([negative_embeds, prompt_embeds], dim=0)
+    prompt_mask = torch.cat([negative_mask, prompt_mask], dim=0)
+    for t in timesteps:
+        latent_model_input = torch.cat([latents] * 2)
+        timestep = t.expand(latent_model_input.shape[0])
+        noise_pred = transformer(
+            latent_model_input.to(dtype=dtype),
+            encoder_hidden_states=prompt_embeds.to(dtype=dtype),
+            encoder_attention_mask=prompt_mask,
+            timestep=timestep,
+            return_dict=False,
+        )[0]
+        noise_pred = noise_pred.float()
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+        if transformer.out_channels // 2 == latent_channels:
+            noise_pred = noise_pred.chunk(2, dim=1)[0]
+        latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+    return latents
+
+
+def test_cfg1_refactored_diffuse_is_bit_identical_to_pr1_loop(monkeypatch):
+    """The refactored diffuse must reproduce the original CFG1 latents bit-for-bit
+    across multiple steps, preserving the TP1/CFG1 golden output."""
+    _set_cfg_world_size(monkeypatch, 1)
+    torch.manual_seed(0)
+    b, c = 2, 4
+    latents0 = torch.randn(b, c, 2, 4, 4)
+    timesteps = torch.tensor([900.0, 600.0, 300.0])
+    prompt_embeds = torch.randn(b, 3, 8)
+    negative_embeds = torch.randn(b, 3, 8)
+    prompt_mask = torch.ones(b, 3)
+    negative_mask = torch.ones(b, 3)
+    guidance_scale = 6.0
+
+    transformer = _DeterministicTransformer(latent_channels=c)
+    scheduler = _LinearScheduler()
+
+    ref = _pr1_reference_loop(
+        transformer,
+        scheduler,
+        latents0.clone(),
+        timesteps,
+        prompt_embeds,
+        prompt_mask,
+        negative_embeds,
+        negative_mask,
+        guidance_scale,
+        latent_channels=c,
+        dtype=torch.float32,
+    )
+
+    pipe = _bare_t2v_pipeline(transformer, guidance_scale=guidance_scale)
+    pipe.scheduler = scheduler
+    got = pipe.diffuse(
+        latents=latents0.clone(),
+        timesteps=timesteps,
+        prompt_embeds=prompt_embeds,
+        prompt_attention_mask=prompt_mask,
+        negative_prompt_embeds=negative_embeds,
+        negative_prompt_attention_mask=negative_mask,
+        guidance_scale=guidance_scale,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=c,
+    )
+
+    assert torch.equal(got, ref)
+
+
+def test_combine_cfg_noise_upcasts_to_fp32_before_combine():
+    """combine_cfg_noise must upcast positive/negative to fp32 before the CFG
+    formula, so the CFG2 (bf16 all-gather) combine is bit-identical to CFG1's
+    fp32 combine."""
+    pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)
+    torch.manual_seed(0)
+    pos = torch.randn(1, 4, 2, 4, 4).bfloat16()
+    neg = torch.randn(1, 4, 2, 4, 4).bfloat16()
+    scale = 6.0
+
+    out = pipe.combine_cfg_noise(pos, neg, scale, cfg_normalize=False)
+
+    expected = neg.float() + scale * (pos.float() - neg.float())
+    assert out.dtype == torch.float32
+    assert torch.equal(out, expected)
+
+
+def test_combine_cfg_noise_accepts_tuple_predictions_from_gather():
+    """The CFG-parallel gather path hands combine_cfg_noise single-element
+    tuples, not bare tensors; the override must upcast and combine them."""
+    pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)
+    torch.manual_seed(0)
+    pos = torch.randn(1, 4, 2, 4, 4).bfloat16()
+    neg = torch.randn(1, 4, 2, 4, 4).bfloat16()
+    scale = 6.0
+
+    out = pipe.combine_cfg_noise((pos,), (neg,), scale, cfg_normalize=False)
+
+    expected = neg.float() + scale * (pos.float() - neg.float())
+    assert out.dtype == torch.float32
+    assert torch.equal(out, expected)
+
+
+def test_cfg2_dispatch_builds_independent_branch_kwargs(monkeypatch):
+    """In CFG world size 2, diffuse must build separate batch-B positive/negative
+    kwargs (no pre-cat), and call the mixin with cfg_normalize=False and the
+    learned-sigma output_slice."""
+    _set_cfg_world_size(monkeypatch, 2)
+
+    b, c = 1, 4
+    captured = {}
+
+    def _spy_cfg(**kwargs):
+        captured.update(kwargs)
+        return torch.zeros(b, c, 2, 4, 4)
+
+    prompt_embeds = torch.randn(b, 3, 8)
+    negative_embeds = torch.randn(b, 3, 8)
+    dtype = torch.float32
+
+    pipe = _bare_t2v_pipeline(transformer=None, guidance_scale=6.0)
+    pipe.predict_noise_maybe_with_cfg = _spy_cfg
+
+    pipe.diffuse(
+        latents=torch.randn(b, c, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=prompt_embeds,
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=negative_embeds,
+        negative_prompt_attention_mask=torch.zeros(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={},
+        dtype=dtype,
+        output_slice=c,
+    )
+
+    assert captured, "predict_noise_maybe_with_cfg was not called in CFG world size 2"
+    assert captured["do_true_cfg"] is True
+    assert captured["true_cfg_scale"] == 6.0
+    assert captured["cfg_normalize"] is False
+    assert captured["output_slice"] == c
+
+    pos = captured["positive_kwargs"]
+    neg = captured["negative_kwargs"]
+    assert pos["hidden_states"].shape[0] == b, "positive branch must be batch-B, not pre-concatenated"
+    assert neg["hidden_states"].shape[0] == b, "negative branch must be batch-B, not pre-concatenated"
+    assert pos["timestep"].shape[0] == b
+    assert torch.equal(pos["encoder_hidden_states"], prompt_embeds.to(dtype))
+    assert torch.equal(neg["encoder_hidden_states"], negative_embeds.to(dtype))
+
+
+# ── I2V CFG dispatch ──
+
+
+class _I2VSpyScheduler:
+    """Records the frame count it steps on and mutates non-first frames so
+    first-frame preservation is observable."""
+
+    order = 1
+
+    def __init__(self):
+        self.stepped_frame_counts = []
+
+    def step(self, noise_pred, t, latents, return_dict=False, **kwargs):
+        self.stepped_frame_counts.append(latents.shape[2])
+        return (latents + 1.0,)
+
+
+def _bare_i2v_pipeline(transformer, scheduler, guidance_scale):
+    pipe = SanaImageToVideoPipeline.__new__(SanaImageToVideoPipeline)
+    pipe.transformer = transformer
+    pipe.scheduler = scheduler
+    pipe._interrupt = False
+    pipe._guidance_scale = guidance_scale
+    return pipe
+
+
+def _conditioning_mask(b, frames):
+    mask = torch.zeros(b, 1, frames, 1, 1)
+    mask[:, :, 0] = 1
+    return mask
+
+
+def test_i2v_cfg1_preserves_first_frame_and_steps_only_tail(monkeypatch):
+    """I2V CFG1 must keep the conditioning first frame bit-exact and only run the
+    scheduler on frames [1:]."""
+    _set_cfg_world_size(monkeypatch, 1)
+    b, c, frames = 1, 4, 3
+
+    class _SpyTransformer:
+        dtype = torch.float32
+        out_channels = 2 * c
+
+        def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
+            return (torch.zeros(hidden_states.shape[0], 2 * c, frames, 4, 4),)
+
+    scheduler = _I2VSpyScheduler()
+    pipe = _bare_i2v_pipeline(_SpyTransformer(), scheduler, guidance_scale=6.0)
+
+    latents0 = torch.randn(b, c, frames, 4, 4)
+    out = pipe.diffuse(
+        latents=latents0.clone(),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        conditioning_mask=_conditioning_mask(b, frames),
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=c,
+    )
+
+    assert torch.equal(out[:, :, :1], latents0[:, :, :1]), "first (conditioning) frame must be unchanged"
+    assert torch.equal(out[:, :, 1:], latents0[:, :, 1:] + 1.0), "tail frames must be the scheduler output"
+    assert scheduler.stepped_frame_counts == [frames - 1], "scheduler must step only on frames [1:]"
+
+
+def test_i2v_cfg2_dispatch_builds_masked_branch_kwargs(monkeypatch):
+    """I2V CFG world size 2 must build batch-B per-branch kwargs with the
+    conditioning-frame timestep zeroed (mask not concatenated) and preserve the
+    first frame."""
+    _set_cfg_world_size(monkeypatch, 2)
+    b, c, frames = 1, 4, 3
+    captured = {}
+
+    def _spy_cfg(**kwargs):
+        captured.update(kwargs)
+        return torch.zeros(b, c, frames, 4, 4)
+
+    scheduler = _I2VSpyScheduler()
+    pipe = _bare_i2v_pipeline(transformer=None, scheduler=scheduler, guidance_scale=6.0)
+    pipe.predict_noise_maybe_with_cfg = _spy_cfg
+
+    prompt_embeds = torch.randn(b, 3, 8)
+    negative_embeds = torch.randn(b, 3, 8)
+    latents0 = torch.randn(b, c, frames, 4, 4)
+    mask = _conditioning_mask(b, frames)
+
+    out = pipe.diffuse(
+        latents=latents0.clone(),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=prompt_embeds,
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=negative_embeds,
+        negative_prompt_attention_mask=torch.zeros(b, 3),
+        guidance_scale=6.0,
+        conditioning_mask=mask,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=c,
+    )
+
+    assert captured["do_true_cfg"] is True
+    assert captured["cfg_normalize"] is False
+    assert captured["output_slice"] == c
+    pos = captured["positive_kwargs"]
+    neg = captured["negative_kwargs"]
+    assert pos["hidden_states"].shape[0] == b, "positive branch must be batch-B"
+    assert neg["hidden_states"].shape[0] == b, "negative branch must be batch-B"
+    assert tuple(pos["timestep"].shape) == tuple(mask.shape), "mask must not be concatenated in CFG2"
+    assert torch.equal(pos["timestep"][:, :, 0], torch.zeros(b, 1, 1, 1)), "conditioning frame timestep must be 0"
+    assert torch.equal(pos["encoder_hidden_states"], prompt_embeds)
+    assert torch.equal(neg["encoder_hidden_states"], negative_embeds)
+    assert torch.equal(out[:, :, :1], latents0[:, :, :1]), "first frame must be preserved"
+
+
+def test_diffuse_runs_cfg_parallel_validity_check(monkeypatch, mocker):
+    """diffuse must run the CFG-parallel validity check so a scale<=1 or missing
+    negative prompt warns the user instead of silently wasting a rank."""
+    _set_cfg_world_size(monkeypatch, 2)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(transformer=None, guidance_scale=6.0)
+    pipe.predict_noise_maybe_with_cfg = lambda **kwargs: torch.zeros(b, 4, 2, 4, 4)
+    check = mocker.patch.object(SanaVideoPipeline, "check_cfg_parallel_validity", return_value=True)
+
+    pipe.diffuse(
+        latents=torch.randn(b, 4, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=None,
+    )
+
+    check.assert_called_once_with(6.0, True)
+
+
+def test_i2v_diffuse_runs_cfg_parallel_validity_check(monkeypatch, mocker):
+    """I2V has its own diffuse; it must also run the CFG-parallel validity check."""
+    _set_cfg_world_size(monkeypatch, 2)
+    b, c, frames = 1, 4, 3
+
+    pipe = _bare_i2v_pipeline(None, _I2VSpyScheduler(), guidance_scale=6.0)
+    pipe.predict_noise_maybe_with_cfg = lambda **kwargs: torch.zeros(b, c, frames, 4, 4)
+    check = mocker.patch.object(SanaVideoPipeline, "check_cfg_parallel_validity", return_value=True)
+
+    pipe.diffuse(
+        latents=torch.randn(b, c, frames, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        conditioning_mask=_conditioning_mask(b, frames),
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=c,
+    )
+
+    check.assert_called_once_with(6.0, True)
+
+
+def test_diffuse_runs_without_initialized_parallel_groups():
+    """Single-process runs (unit tests, offline scripts) never initialize the
+    distributed groups; diffuse must fall back to the serial path instead of
+    tripping the parallel-state assertion. No world-size patching on purpose."""
+
+    class _SpyTransformer:
+        dtype = torch.float32
+
+        def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
+            return (torch.zeros_like(hidden_states),)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
+
+    out = pipe.diffuse(
+        latents=torch.randn(b, 4, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={},
+        dtype=torch.float32,
+        output_slice=None,
+    )
+
+    assert out.shape == (b, 4, 2, 4, 4)
+
+
+def test_diffuse_propagates_group_error_when_parallel_initialized(monkeypatch):
+    """Inside an initialized parallel run a failing CFG-group lookup is a real
+    error and must propagate instead of silently flipping to the serial path."""
+    import vllm_omni.diffusion.models.sana_video.pipeline_sana_video as t2v
+
+    monkeypatch.setattr(t2v, "model_parallel_is_initialized", lambda: True, raising=False)
+
+    def _broken_group_lookup():
+        raise AssertionError("classifier_free_guidance group is not initialized")
+
+    monkeypatch.setattr(t2v, "get_classifier_free_guidance_world_size", _broken_group_lookup)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(None, guidance_scale=6.0)
+
+    with pytest.raises(AssertionError, match="classifier_free_guidance"):
+        pipe.diffuse(
+            latents=torch.randn(b, 4, 2, 4, 4),
+            timesteps=torch.tensor([1000.0]),
+            prompt_embeds=torch.randn(b, 3, 8),
+            prompt_attention_mask=torch.ones(b, 3),
+            negative_prompt_embeds=torch.randn(b, 3, 8),
+            negative_prompt_attention_mask=torch.ones(b, 3),
+            guidance_scale=6.0,
+            extra_step_kwargs={},
+            dtype=torch.float32,
+            output_slice=None,
+        )
+
+
+def test_t2v_forwards_eta_and_generator_into_scheduler(monkeypatch):
+    """T2V must keep passing eta/generator (extra_step_kwargs) into scheduler.step."""
+    _set_cfg_world_size(monkeypatch, 1)
+    captured = {}
+
+    class _KwargScheduler:
+        order = 1
+
+        def step(self, noise_pred, t, latents, **kwargs):
+            captured.update(kwargs)
+            return (latents,)
+
+    class _SpyTransformer:
+        dtype = torch.float32
+
+        def __call__(self, hidden_states, **kwargs):
+            return (torch.zeros_like(hidden_states),)
+
+    b = 1
+    pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
+    pipe.scheduler = _KwargScheduler()
+    generator = torch.Generator().manual_seed(0)
+
+    pipe.diffuse(
+        latents=torch.randn(b, 4, 2, 4, 4),
+        timesteps=torch.tensor([1000.0]),
+        prompt_embeds=torch.randn(b, 3, 8),
+        prompt_attention_mask=torch.ones(b, 3),
+        negative_prompt_embeds=torch.randn(b, 3, 8),
+        negative_prompt_attention_mask=torch.ones(b, 3),
+        guidance_scale=6.0,
+        extra_step_kwargs={"eta": 0.5, "generator": generator},
+        dtype=torch.float32,
+        output_slice=None,
+    )
+
+    assert captured.get("eta") == 0.5
+    assert captured.get("generator") is generator
+    assert captured.get("return_dict") is False
@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Topology validation for SANA-Video parallelism.
+
+The native pipeline only supports tensor/CFG parallel of size 1 or 2 and rejects
+sequence parallel and text-encoder TP before any large weight is allocated.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import SanaVideoPipeline
+from vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v import SanaImageToVideoPipeline
+from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+    validate_sana_video_parallel_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _pc(
+    tp: int = 1,
+    cfg: int = 1,
+    sp: int | None = None,
+    text_tp: int = 1,
+    pp: int = 1,
+    hsdp: bool = False,
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        tensor_parallel_size=tp,
+        cfg_parallel_size=cfg,
+        sequence_parallel_size=sp,
+        text_encoder_tp_size=text_tp,
+        pipeline_parallel_size=pp,
+        use_hsdp=hsdp,
+    )
+
+
+@pytest.mark.parametrize("tp", [1, 2])
+def test_supported_tp_sizes_pass(tp: int) -> None:
+    validate_sana_video_parallel_config(_pc(tp=tp))
+
+
+@pytest.mark.parametrize("tp", [3, 4, 8])
+def test_unsupported_tp_sizes_rejected(tp: int) -> None:
+    with pytest.raises(NotImplementedError, match="tensor_parallel_size"):
+        validate_sana_video_parallel_config(_pc(tp=tp))
+
+
+def test_cfg3_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="cfg_parallel_size"):
+        validate_sana_video_parallel_config(_pc(cfg=3))
+
+
+def test_sequence_parallel_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="[Ss]equence parallel"):
+        validate_sana_video_parallel_config(_pc(sp=2))
+
+
+def test_pipeline_parallel_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="pipeline parallel"):
+        validate_sana_video_parallel_config(_pc(pp=2))
+
+
+def test_hsdp_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="HSDP"):
+        validate_sana_video_parallel_config(_pc(hsdp=True))
+
+
+def test_text_encoder_tp_rejected() -> None:
+    with pytest.raises(NotImplementedError, match="text encoder"):
+        validate_sana_video_parallel_config(_pc(text_tp=2))
+
+
+def test_tp2_cfg2_combo_passes() -> None:
+    validate_sana_video_parallel_config(_pc(tp=2, cfg=2))
+
+
+@pytest.mark.parametrize("pipeline_cls", [SanaVideoPipeline, SanaImageToVideoPipeline])
+def test_pipeline_rejects_unsupported_topology_before_loading(pipeline_cls, mocker) -> None:
+    """The pipeline must reject an unsupported topology before it loads weights."""
+    load = mocker.patch.object(SanaVideoPipeline, "_load_components")
+    od_config = SimpleNamespace(parallel_config=_pc(tp=3))
+
+    with pytest.raises(NotImplementedError, match="tensor_parallel_size"):
+        pipeline_cls(od_config=od_config)
+
+    load.assert_not_called()
@@ -1,11 +1,37 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import os
+
 import pytest
 import torch

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

+
+@pytest.fixture(autouse=True)
+def _init_distributed(monkeypatch):
+    """The native transformer uses vLLM parallel linear layers, which require a
+    tensor-parallel group; initialize a single-process group for CPU tests."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29501")
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
 _TINY_CONFIG = {
    "in_channels": 4,
    "out_channels": 4,
@@ -323,6 +349,7 @@ def test_tiny_transformer_matches_diffusers_and_frozen_output():
    from vllm_omni.diffusion.attention.layer import Attention
    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
    from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+        SanaDistributedRMSNorm,
        SanaLinearAttention,
        SanaRMSNorm,
        SanaVideoTransformerConfig,
@@ -342,10 +369,10 @@ def test_tiny_transformer_matches_diffusers_and_frozen_output():
    assert isinstance(block.attn1, SanaLinearAttention)
    assert isinstance(block.attn2.attn, Attention)
    assert block.attn2.attn.role == "cross"
-    assert isinstance(block.attn1.norm_q, SanaRMSNorm)
-    assert isinstance(block.attn1.norm_k, SanaRMSNorm)
-    assert isinstance(block.attn2.norm_q, SanaRMSNorm)
-    assert isinstance(block.attn2.norm_k, SanaRMSNorm)
+    assert isinstance(block.attn1.norm_q, SanaDistributedRMSNorm)
+    assert isinstance(block.attn1.norm_k, SanaDistributedRMSNorm)
+    assert isinstance(block.attn2.norm_q, SanaDistributedRMSNorm)
+    assert isinstance(block.attn2.norm_k, SanaDistributedRMSNorm)
    assert isinstance(model.caption_norm, SanaRMSNorm)

    torch.manual_seed(11)
@@ -0,0 +1,445 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""TP unit tests for SANA-Video: distributed RMSNorm, sharded attention and
+sharded GLUMB temporal conv.
+
+CPU-only. The TP world size and the collective are mocked so a single process
+simulates the shard/reduce logic a real multi-rank run depends on.
+"""
+
+import os
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
+    GLUMBTempConv,
+    SanaCrossAttention,
+    SanaDistributedRMSNorm,
+    SanaLinearAttention,
+    SanaRMSNorm,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+_MODULE = "vllm_omni.diffusion.models.sana_video.transformer_sana_video"
+
+
+@pytest.fixture
+def tp1_group():
+    """Real single-process TP group so the parallel linear layers construct and
+    run on CPU at tensor_parallel_size=1."""
+    from vllm.distributed.parallel_state import (
+        cleanup_dist_env_and_memory,
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    os.environ.setdefault("MASTER_PORT", "29501")
+    init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://")
+    initialize_model_parallel()
+    yield
+    cleanup_dist_env_and_memory()
+
+
+@pytest.fixture
+def force_default_gemm(monkeypatch):
+    """Force CPU-compatible GEMM dispatch for the parallel linear layers."""
+    from vllm.model_executor.layers.utils import default_unquantized_gemm
+
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
+        lambda: default_unquantized_gemm,
+    )
+
+
+def test_distributed_rms_norm_tp1_matches_sana_rms_norm(mocker) -> None:
+    """TP1 distributed norm must be bit-identical to the checkpoint SanaRMSNorm."""
+    dim = 24
+    torch.manual_seed(0)
+
+    dense = SanaRMSNorm(dim, eps=1e-5, elementwise_affine=True)
+    dense.weight.data.normal_()
+
+    dist = SanaDistributedRMSNorm(dim, eps=1e-5)
+    dist.weight.data.copy_(dense.weight.data)
+
+    x = torch.randn(2, 5, dim)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1)
+    out = dist(x)
+    ref = dense(x)
+
+    torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+def test_distributed_rms_norm_tp1_bf16_weight_matches_sana_rms_norm(mocker) -> None:
+    """bf16 weight must follow SanaRMSNorm's cast-to-weight-dtype branch.
+
+    An unconditional cast back to the input dtype would diverge here; this pins
+    the fp16/bf16-only cast that SanaRMSNorm performs.
+    """
+    dim = 24
+    torch.manual_seed(0)
+
+    dense = SanaRMSNorm(dim, eps=1e-5, elementwise_affine=True)
+    dense.weight.data.normal_()
+    dense.weight.data = dense.weight.data.to(torch.bfloat16)
+
+    dist = SanaDistributedRMSNorm(dim, eps=1e-5)
+    dist.weight.data = dense.weight.data.clone()
+
+    x = torch.randn(2, 5, dim, dtype=torch.bfloat16)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1)
+    out = dist(x)
+    ref = dense(x)
+
+    assert out.dtype == torch.bfloat16
+    torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+def test_distributed_rms_norm_tp2_shards_aggregate_to_dense(mocker) -> None:
+    """Two TP2 shards must reproduce dense SanaRMSNorm.
+
+    The RMS must use the full (unsharded) hidden size: each rank reduces its
+    local sum-of-squares across ranks, not a per-shard mean over its half.
+    """
+    full_dim, tp_size = 24, 2
+    half = full_dim // tp_size
+    torch.manual_seed(0)
+
+    dense = SanaRMSNorm(full_dim, eps=1e-5, elementwise_affine=True)
+    dense.weight.data.normal_()
+
+    dist0 = SanaDistributedRMSNorm(half, eps=1e-5)
+    dist1 = SanaDistributedRMSNorm(half, eps=1e-5)
+    dist0.weight.data.copy_(dense.weight.data[:half])
+    dist1.weight.data.copy_(dense.weight.data[half:])
+
+    x = torch.randn(2, 5, full_dim)
+    x0, x1 = x[..., :half], x[..., half:]
+
+    # A real all-reduce returns the sum of every rank's local sum-of-squares.
+    global_sum_sq = x.to(torch.float32).pow(2).sum(-1, keepdim=True)
+
+    def fake_all_reduce(tensor: torch.Tensor) -> torch.Tensor:
+        return global_sum_sq
+
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size)
+    mocker.patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=fake_all_reduce)
+    out0 = dist0(x0)
+    out1 = dist1(x1)
+
+    out = torch.cat([out0, out1], dim=-1)
+    ref = dense(x)
+
+    torch.testing.assert_close(out, ref, rtol=0, atol=0)
+
+
+def test_distributed_rms_norm_weight_loader_shards_dim0(mocker) -> None:
+    """Each rank loads its own dim-0 slice of the global affine weight."""
+    full_dim, tp_size, rank = 24, 2, 1
+    half = full_dim // tp_size
+
+    norm = SanaDistributedRMSNorm(half, eps=1e-5)
+    full_weight = torch.randn(full_dim)
+
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=tp_size)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank)
+    norm.weight.weight_loader(norm.weight, full_weight)
+
+    torch.testing.assert_close(norm.weight.data, full_weight[rank * half : (rank + 1) * half])
+
+
+def test_distributed_rms_norm_weight_loader_tp1_copies_full() -> None:
+    dim = 24
+    norm = SanaDistributedRMSNorm(dim, eps=1e-5)
+    full_weight = torch.randn(dim)
+
+    norm.weight.weight_loader(norm.weight, full_weight)
+
+    torch.testing.assert_close(norm.weight.data, full_weight)
+
+
+_SELF_ATTN_GOLDEN = torch.tensor(
+    [
+        5.0944157,
+        -16.9981823,
+        55.4457436,
+        42.5174408,
+        -9.3055859,
+        19.3744793,
+        -47.4345360,
+        -5.8110709,
+        -14.0679598,
+        -22.3440742,
+        -8.3834534,
+        8.9109459,
+    ]
+)
+
+
+def _fill_by_name(module: torch.nn.Module) -> None:
+    """Deterministic per-parameter fill that is stable across TP refactors."""
+    torch.manual_seed(0)
+    for _, param in sorted(module.named_parameters()):
+        param.data.normal_()
+
+
+def _self_attn_inputs(seq_len: int, dim: int, head_dim: int):
+    torch.manual_seed(1)
+    x = torch.randn(2, seq_len, dim)
+    freqs_cos = torch.randn(1, seq_len, 1, head_dim)
+    freqs_sin = torch.randn(1, seq_len, 1, head_dim)
+    return x, (freqs_cos, freqs_sin)
+
+
+def test_self_attention_tp1_matches_golden_and_keeps_param_names(tp1_group, force_default_gemm) -> None:
+    """TP1 self-attention stays bit-identical to PR1 and preserves checkpoint keys."""
+    dim, num_heads, head_dim, seq_len = 24, 4, 6, 8
+    attn = SanaLinearAttention(
+        dim=dim, num_heads=num_heads, head_dim=head_dim, dropout=0.0, bias=False, qk_norm="rms_norm_across_heads"
+    )
+    attn.eval()
+    _fill_by_name(attn)
+    x, rotary_emb = _self_attn_inputs(seq_len, dim, head_dim)
+
+    with torch.no_grad():
+        out = attn(x, rotary_emb=rotary_emb)
+
+    assert {name for name, _ in attn.named_parameters()} == {
+        "to_q.weight",
+        "to_k.weight",
+        "to_v.weight",
+        "norm_q.weight",
+        "norm_k.weight",
+        "to_out.0.weight",
+        "to_out.0.bias",
+    }
+    torch.testing.assert_close(out.flatten()[:12], _SELF_ATTN_GOLDEN, rtol=1e-5, atol=1e-5)
+
+
+def _mock_tp2(mocker):
+    """Mock a TP world size of 2 for module construction (no real group)."""
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2)
+    mocker.patch(
+        "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
+        return_value=2,
+    )
+    tp_group = mocker.MagicMock()
+    tp_group.world_size = 2
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=tp_group)
+
+
+def test_self_attention_tp2_shards_projections_and_heads(mocker) -> None:
+    """TP2 splits heads column-wise (q/k/v) and the output projection row-wise."""
+    dim, num_heads, head_dim = 24, 4, 6
+    inner_dim = num_heads * head_dim
+    _mock_tp2(mocker)
+
+    attn = SanaLinearAttention(
+        dim=dim, num_heads=num_heads, head_dim=head_dim, dropout=0.0, bias=False, qk_norm="rms_norm_across_heads"
+    )
+
+    assert attn.heads == num_heads // 2
+    assert attn.to_q.weight.shape == (inner_dim // 2, dim)
+    assert attn.to_k.weight.shape == (inner_dim // 2, dim)
+    assert attn.to_v.weight.shape == (inner_dim // 2, dim)
+    assert attn.to_out[0].weight.shape == (dim, inner_dim // 2)
+    assert attn.norm_q.weight.shape == (inner_dim // 2,)
+    assert attn.norm_k.weight.shape == (inner_dim // 2,)
+
+
+def test_cross_attention_tp2_shards_projections_and_heads(mocker) -> None:
+    """TP2 cross-attention shards q/k/v column-wise, to_out row-wise, heads locally."""
+    dim, cross_dim, num_heads, head_dim = 24, 16, 4, 6
+    inner_dim = num_heads * head_dim
+    _mock_tp2(mocker)
+
+    attn = SanaCrossAttention(
+        dim=dim,
+        cross_attention_dim=cross_dim,
+        num_heads=num_heads,
+        head_dim=head_dim,
+        dropout=0.0,
+        bias=False,
+        out_bias=True,
+        qk_norm="rms_norm_across_heads",
+        prefix="attn2",
+    )
+
+    assert attn.heads == num_heads // 2
+    assert attn.to_q.weight.shape == (inner_dim // 2, dim)
+    assert attn.to_k.weight.shape == (inner_dim // 2, cross_dim)
+    assert attn.to_v.weight.shape == (inner_dim // 2, cross_dim)
+    assert attn.to_out[0].weight.shape == (dim, inner_dim // 2)
+    assert isinstance(attn.norm_q, SanaDistributedRMSNorm)
+    assert attn.norm_q.weight.shape == (inner_dim // 2,)
+
+
+_TINY_TRANSFORMER_CONFIG = {
+    "in_channels": 4,
+    "out_channels": 4,
+    "num_attention_heads": 2,
+    "attention_head_dim": 12,
+    "num_layers": 1,
+    "num_cross_attention_heads": 2,
+    "cross_attention_head_dim": 12,
+    "cross_attention_dim": 24,
+    "caption_channels": 8,
+    "mlp_ratio": 2.0,
+    "patch_size": (1, 2, 2),
+    "sample_size": 4,
+    "rope_max_seq_len": 32,
+}
+
+
+def test_tp2_load_weights_shards_full_checkpoint_without_missing(mocker) -> None:
+    """The full checkpoint loads under TP2 via weight_loaders with no missing/unexpected keys."""
+    from diffusers import SanaVideoTransformer3DModel as DiffusersTransformer
+
+    from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel
+
+    _mock_tp2(mocker)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=0)
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+
+    reference = DiffusersTransformer(**_TINY_TRANSFORMER_CONFIG)
+    model = SanaVideoTransformer3DModel(**_TINY_TRANSFORMER_CONFIG)
+
+    loaded = model.load_weights(list(reference.state_dict().items()))
+
+    assert loaded == set(model.state_dict())
+
+
+# ── GLUMB temporal conv TP ──
+
+# hidden_channels = expand_ratio * in_channels = 8, so conv_inverted emits 2h = 16.
+_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO = 4, 4, 2.0
+
+_GLUMB_GOLDEN = torch.tensor(
+    [
+        51.4757385,
+        -10.9912930,
+        -4.1482553,
+        -26.1580830,
+        -19.7727261,
+        17.9901562,
+        -1.6007471,
+        13.3993073,
+        -8.6459389,
+        -10.4008112,
+        -7.4962354,
+        1.5446399,
+    ]
+)
+
+
+def _glumb_inputs() -> torch.Tensor:
+    torch.manual_seed(1)
+    return torch.randn(1, 2, 3, 3, _GLUMB_IN)  # [B, F, H, W, C]
+
+
+def test_glumb_tp1_matches_golden_and_keeps_param_names(tp1_group) -> None:
+    """TP1 GLUMB stays bit-identical to the PR1 conv output and preserves keys."""
+    glumb = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False)
+    glumb.eval()
+    _fill_by_name(glumb)
+    x = _glumb_inputs()
+
+    with torch.no_grad():
+        out = glumb(x)
+
+    assert {name for name, _ in glumb.named_parameters()} == {
+        "conv_inverted.weight",
+        "conv_inverted.bias",
+        "conv_depth.weight",
+        "conv_depth.bias",
+        "conv_point.weight",
+        "conv_temp.weight",
+    }
+    torch.testing.assert_close(out.flatten()[:12], _GLUMB_GOLDEN, rtol=1e-5, atol=1e-5)
+
+
+def _random_glumb_checkpoint(hidden_channels: int) -> dict[str, torch.Tensor]:
+    """Checkpoint-format (nn.Conv2d) GLUMB weights with per-segment asymmetry.
+
+    conv_inverted / conv_depth carry two h-channel segments (value then gate).
+    Filling the gate half from a different distribution than the value half is
+    what exposes a naive contiguous dim-0 split: that bug would place the whole
+    value segment on rank 0 and the whole gate segment on rank 1, so each rank's
+    local chunk(2) would slice value-vs-value instead of value-vs-gate.
+    """
+    torch.manual_seed(0)
+    h = hidden_channels
+    value_w = torch.randn(h, _GLUMB_IN, 1, 1)
+    gate_w = torch.randn(h, _GLUMB_IN, 1, 1) * 3.0 + 5.0
+    value_b = torch.randn(h)
+    gate_b = torch.randn(h) * 3.0 + 5.0
+    value_dw = torch.randn(h, 1, 3, 3)
+    gate_dw = torch.randn(h, 1, 3, 3) * 3.0 + 5.0
+    value_db = torch.randn(h)
+    gate_db = torch.randn(h) * 3.0 + 5.0
+    return {
+        "conv_inverted.weight": torch.cat([value_w, gate_w], dim=0),
+        "conv_inverted.bias": torch.cat([value_b, gate_b], dim=0),
+        "conv_depth.weight": torch.cat([value_dw, gate_dw], dim=0),
+        "conv_depth.bias": torch.cat([value_db, gate_db], dim=0),
+        "conv_point.weight": torch.randn(_GLUMB_OUT, h, 1, 1),
+        "conv_temp.weight": torch.randn(_GLUMB_OUT, _GLUMB_OUT, 3, 1),
+    }
+
+
+def _mock_glumb_tp(mocker, world_size: int, rank: int) -> None:
+    """Mock a TP group of ``world_size`` at ``rank`` with an identity all-reduce,
+    so a single process can build and run one rank's shard in isolation."""
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=world_size)
+    mocker.patch(f"{_MODULE}.get_tensor_model_parallel_rank", return_value=rank)
+    # Identity conv_point all-reduce returns each rank's local partial for the
+    # test to sum.
+    mocker.patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=lambda t: t)
+
+
+def _load_glumb_checkpoint(glumb: GLUMBTempConv, checkpoint: dict[str, torch.Tensor]) -> None:
+    """Load like AutoWeightsLoader: custom weight_loader when present (the sharded
+    convs), plain copy for replicated params (conv_temp)."""
+    params = dict(glumb.named_parameters())
+    for name, tensor in checkpoint.items():
+        param = params[name]
+        loader = getattr(param, "weight_loader", None)
+        if loader is None:
+            param.data.copy_(tensor)
+        else:
+            loader(param, tensor)
+
+
+def test_glumb_tp2_shards_aggregate_to_dense(mocker) -> None:
+    """The merged value/gate two-segment shard must aggregate to the dense conv.
+
+    Everything after conv_point (temporal conv + residual) is linear and
+    bias-free, so summing each rank's full forward (with the conv_point
+    all-reduce mocked to identity) reproduces the replicated dense output.
+    """
+    hidden_channels = int(_GLUMB_RATIO * _GLUMB_IN)
+    checkpoint = _random_glumb_checkpoint(hidden_channels)
+    x = _glumb_inputs()
+
+    _mock_glumb_tp(mocker, world_size=1, rank=0)
+    dense = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False)
+    dense.eval()
+    _load_glumb_checkpoint(dense, checkpoint)
+    with torch.no_grad():
+        ref = dense(x)
+
+    partials = []
+    for rank in (0, 1):
+        _mock_glumb_tp(mocker, world_size=2, rank=rank)
+        shard = GLUMBTempConv(_GLUMB_IN, _GLUMB_OUT, _GLUMB_RATIO, norm_type=None, residual_connection=False)
+        shard.eval()
+        _load_glumb_checkpoint(shard, checkpoint)
+        with torch.no_grad():
+            partials.append(shard(x))
+
+    aggregated = partials[0] + partials[1]
+    torch.testing.assert_close(aggregated, ref, rtol=1e-5, atol=1e-5)
@@ -14,6 +14,7 @@ from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 pytestmark = [pytest.mark.diffusion]

@@ -133,12 +134,7 @@ class _CompileTrackingModel:


 def _make_request():
-    sampling_params = SimpleNamespace(
-        generator=None,
-        seed=None,
-        generator_device=None,
-        num_inference_steps=4,
-    )
+    sampling_params = OmniDiffusionSamplingParams(num_inference_steps=4)
    return SimpleNamespace(
        request_id="req-test",
        prompt="a prompt",
@@ -249,6 +245,58 @@ def _make_compile_runner(
    return runner


+class _EnabledCacheBackend:
+    def __init__(self):
+        self.refresh_calls = []
+
+    def is_enabled(self):
+        return True
+
+    def refresh(self, pipeline, num_inference_steps, verbose=True):
+        self.refresh_calls.append((pipeline, num_inference_steps, verbose))
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_refresh_cache_prefers_request_steps_then_schedule_then_pipeline_default():
+    cache_backend = _EnabledCacheBackend()
+    runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit")
+    runner.pipeline.default_num_inference_steps = 50
+    req = _make_request()
+
+    custom_timesteps = torch.linspace(999.0, 0.0, 8)
+    custom_sigmas = [0.9, 0.5, 0.1]
+    cases = [
+        (20, None, None),
+        (30, custom_timesteps, custom_sigmas),
+        (None, custom_timesteps, None),
+        (None, custom_timesteps, custom_sigmas),
+        (None, None, custom_sigmas),
+        (None, None, None),
+    ]
+    for num_inference_steps, timesteps, sigmas in cases:
+        req.sampling_params.num_inference_steps = num_inference_steps
+        req.sampling_params.timesteps = timesteps
+        req.sampling_params.sigmas = sigmas
+        DiffusionModelRunner._refresh_cache_for_requests(runner, [req], od_config=runner.od_config)
+
+    assert [steps for _, steps, _ in cache_backend.refresh_calls] == [20, 30, 8, 8, 3, 50]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_refresh_cache_without_request_or_pipeline_default_warns(caplog):
+    cache_backend = _EnabledCacheBackend()
+    runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit")
+    req = _make_request()
+    req.sampling_params.num_inference_steps = None
+
+    DiffusionModelRunner._refresh_cache_for_requests(runner, [req], od_config=runner.od_config)
+
+    assert cache_backend.refresh_calls == []
+    assert "requires num_inference_steps to be passed explicitly" in caplog.text
+
+
@pytest.mark.core_model
@pytest.mark.cpu
 def test_update_states_carries_prepared_layout() -> None:
@@ -451,16 +499,6 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa
@pytest.mark.core_model
@hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeypatch):
-    class _EnabledCacheBackend:
-        def __init__(self):
-            self.refresh_calls = []
-
-        def is_enabled(self):
-            return True
-
-        def refresh(self, pipeline, num_inference_steps, verbose=True):
-            self.refresh_calls.append((pipeline, num_inference_steps, verbose))
-
    cache_backend = _EnabledCacheBackend()
    runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit")
    req = _make_request()
@@ -10,8 +10,9 @@ import sys
 import time
 from abc import ABC, abstractmethod
 from base64 import b64decode, b64encode
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterator, Mapping
 from contextlib import contextmanager
+from dataclasses import dataclass
 from fractions import Fraction
 from hashlib import sha1
 from io import BytesIO
@@ -81,6 +82,59 @@ def model_output_dir(parent_dir: Path, model: str) -> Path:
    return path


+DEFAULT_DEVICE_PROFILE = "default"
+
+
+@dataclass(frozen=True, slots=True)
+class SimilarityThresholds:
+    """SSIM/PSNR floors for accuracy comparisons."""
+
+    ssim: float
+    psnr: float
+
+
+def resolve_device_profile(
+    device_name: str | None = None,
+    *,
+    profiles: Mapping[str, Any],
+) -> str:
+    """Map a device marketing name to a key in *profiles*.
+
+    When *device_name* is omitted, uses ``current_omni_platform.get_device_name()``.
+    Non-``default`` keys are matched as substrings; unmatched devices use ``default``.
+    """
+    if device_name is None:
+        try:
+            from vllm_omni.platforms import current_omni_platform
+
+            device_name = current_omni_platform.get_device_name()
+        except Exception:
+            return DEFAULT_DEVICE_PROFILE
+
+    if not device_name:
+        return DEFAULT_DEVICE_PROFILE
+
+    for profile in profiles:
+        if profile == DEFAULT_DEVICE_PROFILE:
+            continue
+        if profile in device_name:
+            return profile
+    return DEFAULT_DEVICE_PROFILE
+
+
+def resolve_similarity_thresholds(
+    table: Mapping[str, SimilarityThresholds],
+    device_profile: str | None = None,
+    *,
+    device_name: str | None = None,
+) -> SimilarityThresholds:
+    """Pick SSIM/PSNR thresholds from *table* for the current (or given) device profile."""
+    if DEFAULT_DEVICE_PROFILE not in table:
+        raise ValueError("threshold table must include a 'default' entry")
+    profile = device_profile or resolve_device_profile(device_name, profiles=table)
+    return table.get(profile, table[DEFAULT_DEVICE_PROFILE])
+
+
 _SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
 _PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")

@@ -94,15 +94,28 @@ def _default_ar_dit_devices() -> tuple[str, str]:


 AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices()
-MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
-HUNYUAN_IMAGE_REF_PATH = get_asset_path("hunyuan/hunyuan_image_ref.png")
-NUM_INFERENCE_STEPS = 50
+
+# Model configuration based on HUNYUAN_MODEL_PATH
+DISTIL_MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct-Distil"
+INSTRUCT_MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
+MODEL_PATH = os.environ.get("HUNYUAN_MODEL_PATH", DISTIL_MODEL_NAME)
+
+# Determine model type and set parameters accordingly
+IS_DISTIL = "distil" in MODEL_PATH.lower()
+MODEL_NAME = DISTIL_MODEL_NAME if IS_DISTIL else INSTRUCT_MODEL_NAME
+NUM_INFERENCE_STEPS = 8 if IS_DISTIL else 50
 GUIDANCE_SCALE = 2.5

+# Asset paths based on model type
+ASSET_SUBDIR = "hunyuan_image3"
+REF_IMAGE_NAME = "hunyuan_image_distill_ref.png" if IS_DISTIL else "hunyuan_image_instruct_ref.png"
+COT_REF_NAME = "hunyuan_image_distill_cot_ref.txt" if IS_DISTIL else "hunyuan_image3_instruct_cot_ref.txt"
+HUNYUAN_IMAGE_REF_PATH = get_asset_path(f"{ASSET_SUBDIR}/{REF_IMAGE_NAME}")
+HUNYUAN_COT_REF_PATH = get_asset_path(f"{ASSET_SUBDIR}/{COT_REF_NAME}")
+
 # ============================================================================
 # Constants
 # ============================================================================
-MODEL_PATH = os.environ.get("HUNYUAN_MODEL_PATH", MODEL_NAME)
 # Test input
 PROMPT = "基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴"
 TEST_IMAGE_URLS = [
@@ -113,23 +126,52 @@ SEED = 42
 AR_TP_SIZE = len(AR_DEVICES.split(","))
 DIT_TP_SIZE = len(DIT_DEVICES.split(","))

-# Precision thresholds
+# Precision thresholds organized by model variant and test type
 THRESHOLDS = {
-    # AR text comparison
-    "text_prefix_match": 10,  # First 10 characters must match exactly
-    "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
-    # Image comparison
-    "clip_score": 90,  # CLIP image semantic similarity
-    "ssim": 0.26,  # Structural similarity
-    "psnr": 12.5,  # Peak signal-to-noise ratio (dB)
+    "instruct": {
+        # AR text comparison
+        "text_prefix_match": 10,  # First 10 characters must match exactly
+        "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
+        # Image comparison
+        "clip_score": 90,  # CLIP image semantic similarity
+        "ssim": 0.26,  # Structural similarity
+        "psnr": 12.5,  # Peak signal-to-noise ratio (dB)
+    },
+    # The 8-step distilled model has more structural variance across GPU
+    # runtimes; CLIP and PSNR retain the same semantic and pixel safeguards.
+    "distil": {
+        "text_prefix_match": 10,
+        "cot_semantic_sim": 0.9,
+        "clip_score": 90,
+        "ssim": 0.23,
+        "psnr": 12.5,
+    },
+    # Thresholds for quantized DiT accuracy tests (bf16 vs quant comparison)
+    "quant": {
+        "psnr": 10.0,
+        "ssim": 0.20,
+        "clip_score": 20.0,  # Minimum CLIP score for quantized output
+        "clip_score_drop": float(
+            os.environ.get("HUNYUAN_IMAGE3_QUANT_CLIP_SCORE_DROP_THRESHOLD", "5.0")
+        ),  # Max allowed drop from bf16
+    },
 }

+
+# Helper to get thresholds for current model variant
+def _get_thresholds() -> dict:
+    """Return appropriate thresholds based on test context."""
+    return THRESHOLDS["distil" if IS_DISTIL else "instruct"]
+
+
+def _get_quant_thresholds() -> dict:
+    """Return thresholds for quantization accuracy tests."""
+    return THRESHOLDS["quant"]
+
+
+# Quant test configuration
 QUANT_PROMPT = "A brown and white dog is running on the grass."
 QUANT_HEIGHT, QUANT_WIDTH = 1024, 1024
-QUANT_PSNR_THRESHOLD = 10.0
-QUANT_SSIM_THRESHOLD = 0.20
-QUANT_CLIP_SCORE_THRESHOLD = 20.0
-QUANT_CLIP_SCORE_DROP_THRESHOLD = float(os.environ.get("HUNYUAN_IMAGE3_QUANT_CLIP_SCORE_DROP_THRESHOLD", "5.0"))
 QUANT_RUN_ENV = "HUNYUAN_IMAGE3_RUN_QUANT_ACCURACY"
 QUANT_BF16_ENV = "HUNYUAN_IMAGE3_BF16_MODEL"
 QUANT_FP8_ENV = "HUNYUAN_IMAGE3_FP8_MODEL"
@@ -167,7 +209,6 @@ _DEPLOY_CONFIG = {
                "top_p": 1,
                "top_k": -1,
                "max_tokens": 8192,
-                "stop_token_ids": [128025],
                "detokenize": True,
                "skip_special_tokens": False,
                "include_stop_str_in_output": True,
@@ -246,19 +287,6 @@ def _quant_accuracy_cases() -> list[pytest.ParameterSet]:
    return params


-# fmt: off
-COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，"
-           "它们被放置在灰色的织物背景上。这些冰箱贴的关键特征是其材质：它们具有光滑、高光的珐琅或烤漆质感，边缘有明显的金属包边，整体呈现出一种立体的、有厚度的实体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）"
-           "转化为一个具有特定物理属性（材质和立体感）的三维物体。核心任务是保留logo的视觉识别性，同时赋予其冰箱贴的实体质感。为了构建答案图像，我会将图一的圆形logo作为基础形状。然后，我会将图二中冰箱贴的材质特性应用到这个logo上。具体来说，logo中的每一个色块（深蓝、中蓝、浅蓝、白色）"
-           "都会被渲染成具有高光泽度的珐琅质感，表面会反射出柔和的环境光，形成自然的高光。logo中不同颜色区域之间的分界线，将被处理成纤细的、带有金属光泽的凸起边缘，这既能清晰地勾勒出图案，也符合珐琅工艺品的典型特征。整个冰箱贴会呈现出轻微的厚度和圆润的边缘，使其看起来像一个真实的、可触摸的物体。"
-           "最后，将这个制作完成的冰箱贴放置在图二所示的灰色织物背景上，并为其添加一个微妙的、柔和的阴影，以增强其立体感和与背景的融合度，最终呈现出一个精致、逼真的产品展示图。</think><recaption>这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。"
-           "冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。"
-           "柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。"
-           "</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi>"
-           "<img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>")
-# fmt: on
-
-
 def _make_config(enable_kv_reuse: bool, path: Path) -> None:
    config = copy.deepcopy(_DEPLOY_CONFIG)
    config["stages"][0]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse}
@@ -432,34 +460,36 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path) -> None:
    online_cot = online_cot.lstrip("\n")
    scorer = SemanticSimilarityScorer()
    clip_scorer = CLIPScorer()
-    cot_results = scorer.text_similarity(online_cot, COT_REF)
+    cot_ref = HUNYUAN_COT_REF_PATH.read_text(encoding="utf-8")
+    cot_results = scorer.text_similarity(online_cot, cot_ref)
    image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB")
    image_clip_score = clip_scorer.image_image_score(online_image, image_ref)
    ssim_value, psnr_value = compute_image_ssim_psnr(prediction=online_image, reference=image_ref, compare_mode="RGB")

+    thresholds = _get_thresholds()
    table = [
-        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644],
-        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29],
-        ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538],
-        ["SSIM", f"{ssim_value:.4f}", 0.242],
-        ["PSNR (dB)", f"{psnr_value:.2f}", 14.1],
+        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", thresholds["cot_semantic_sim"]],
+        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", thresholds["text_prefix_match"]],
+        ["Image-Image similarity", f"{image_clip_score:.4f}", thresholds["clip_score"]],
+        ["SSIM", f"{ssim_value:.4f}", thresholds["ssim"]],
+        ["PSNR (dB)", f"{psnr_value:.2f}", thresholds["psnr"]],
    ]
-    print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
+    print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "Threshold"], tablefmt="grid"))

    _assert_cot_structural_markers(online_cot, "ONLINE")

-    assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
-        f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {THRESHOLDS['cot_semantic_sim']}"
+    assert cot_results["cot_semantic_sim"] >= thresholds["cot_semantic_sim"], (
+        f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {thresholds['cot_semantic_sim']}"
    )
-    assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], (
-        f"[ONLINE] COT prefix match {cot_results['text_prefix_match_count']} below threshold {THRESHOLDS['text_prefix_match']}"
+    assert cot_results["text_prefix_match_count"] >= thresholds["text_prefix_match"], (
+        f"[ONLINE] COT prefix match {cot_results['text_prefix_match_count']} below threshold {thresholds['text_prefix_match']}"
    )
-    assert image_clip_score >= THRESHOLDS["clip_score"], (
-        f"[ONLINE] Image-Image similarity {image_clip_score:.4f} below threshold {THRESHOLDS['clip_score']}"
+    assert image_clip_score >= thresholds["clip_score"], (
+        f"[ONLINE] Image-Image similarity {image_clip_score:.4f} below threshold {thresholds['clip_score']}"
    )
-    assert ssim_value >= THRESHOLDS["ssim"], f"[ONLINE] SSIM {ssim_value:.4f} below threshold {THRESHOLDS['ssim']}"
-    assert psnr_value >= THRESHOLDS["psnr"], (
-        f"[ONLINE] PSNR {psnr_value:.2f} dB below threshold {THRESHOLDS['psnr']} dB"
+    assert ssim_value >= thresholds["ssim"], f"[ONLINE] SSIM {ssim_value:.4f} below threshold {thresholds['ssim']}"
+    assert psnr_value >= thresholds["psnr"], (
+        f"[ONLINE] PSNR {psnr_value:.2f} dB below threshold {thresholds['psnr']} dB"
    )


@@ -546,34 +576,36 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path) -> None:

    scorer = SemanticSimilarityScorer()
    clip_scorer = CLIPScorer()
-    cot_results = scorer.text_similarity(omni_cot, COT_REF)
+    cot_ref = HUNYUAN_COT_REF_PATH.read_text(encoding="utf-8")
+    cot_results = scorer.text_similarity(omni_cot, cot_ref)
    image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB")
    image_clip_score = clip_scorer.image_image_score(omni_image, image_ref)
    ssim_value, psnr_value = compute_image_ssim_psnr(prediction=omni_image, reference=image_ref, compare_mode="RGB")

+    thresholds = _get_thresholds()
    table = [
-        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644],
-        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29],
-        ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538],
-        ["SSIM", f"{ssim_value:.4f}", 0.242],
-        ["PSNR (dB)", f"{psnr_value:.2f}", 14.1],
+        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", thresholds["cot_semantic_sim"]],
+        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", thresholds["text_prefix_match"]],
+        ["Image-Image similarity", f"{image_clip_score:.4f}", thresholds["clip_score"]],
+        ["SSIM", f"{ssim_value:.4f}", thresholds["ssim"]],
+        ["PSNR (dB)", f"{psnr_value:.2f}", thresholds["psnr"]],
    ]

-    print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
+    print(tabulate(table, headers=["Metric", "Value", "Threshold"], tablefmt="grid"))

    _assert_cot_structural_markers(omni_cot, "OFFLINE")

-    assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
-        f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}"
+    assert cot_results["cot_semantic_sim"] >= thresholds["cot_semantic_sim"], (
+        f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {thresholds['cot_semantic_sim']}"
    )
-    assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], (
-        f"COT prefix match count {cot_results['text_prefix_match_count']} is below threshold {THRESHOLDS['text_prefix_match']}"
+    assert cot_results["text_prefix_match_count"] >= thresholds["text_prefix_match"], (
+        f"COT prefix match count {cot_results['text_prefix_match_count']} is below threshold {thresholds['text_prefix_match']}"
    )
-    assert image_clip_score >= THRESHOLDS["clip_score"], (
-        f"Image-Image similarity{image_clip_score:.4f} is below threshold {THRESHOLDS['clip_score']}"
+    assert image_clip_score >= thresholds["clip_score"], (
+        f"Image-Image similarity{image_clip_score:.4f} is below threshold {thresholds['clip_score']}"
    )
-    assert ssim_value >= THRESHOLDS["ssim"], f"SSIM {ssim_value:.4f} is below threshold {THRESHOLDS['ssim']}"
-    assert psnr_value >= THRESHOLDS["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {THRESHOLDS['psnr']} dB"
+    assert ssim_value >= thresholds["ssim"], f"SSIM {ssim_value:.4f} is below threshold {thresholds['ssim']}"
+    assert psnr_value >= thresholds["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {thresholds['psnr']} dB"


@pytest.mark.parametrize("case", _quant_accuracy_cases())
@@ -607,12 +639,13 @@ def test_quantized_dit_matches_bf16_accuracy(
        prediction=quant_image,
        reference=bf16_image,
    )
+    quant_thresholds = _get_quant_thresholds()
    assert_similarity(
        model_name=f"{MODEL_NAME} {case.name} vs bf16",
        vllm_image=quant_image,
        diffusers_image=bf16_image,
-        ssim_threshold=QUANT_SSIM_THRESHOLD,
-        psnr_threshold=QUANT_PSNR_THRESHOLD,
+        ssim_threshold=quant_thresholds["ssim"],
+        psnr_threshold=quant_thresholds["psnr"],
        width=QUANT_WIDTH,
        height=QUANT_HEIGHT,
    )
@@ -647,15 +680,15 @@ def test_quantized_dit_matches_bf16_accuracy(
    print(f"  bf16 model:       {bf16_model}")
    print(f"  quant model:      {quant_model}")
    print(f"  BF16 CLIP score:  {bf16_clip_score:.4f}")
-    print(f"  quant CLIP score: {quant_clip_score:.4f} threshold>={QUANT_CLIP_SCORE_THRESHOLD:.4f}")
-    print(f"  CLIP score drop:  {clip_score_drop:.4f} threshold<={QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}")
+    print(f"  quant CLIP score: {quant_clip_score:.4f} threshold>={quant_thresholds['clip_score']:.4f}")
+    print(f"  CLIP score drop:  {clip_score_drop:.4f} threshold<={quant_thresholds['clip_score_drop']:.4f}")
    print(f"  metrics:          {metrics_path}")

-    assert quant_clip_score >= QUANT_CLIP_SCORE_THRESHOLD, (
+    assert quant_clip_score >= quant_thresholds["clip_score"], (
        f"{case.name} CLIP score below threshold: got {quant_clip_score:.4f}, "
-        f"expected >= {QUANT_CLIP_SCORE_THRESHOLD:.4f}"
+        f"expected >= {quant_thresholds['clip_score']:.4f}"
    )
-    assert clip_score_drop <= QUANT_CLIP_SCORE_DROP_THRESHOLD, (
+    assert clip_score_drop <= quant_thresholds["clip_score_drop"], (
        f"{case.name} CLIP score drop too large: got {clip_score_drop:.4f}, "
-        f"expected <= {QUANT_CLIP_SCORE_DROP_THRESHOLD:.4f}"
+        f"expected <= {quant_thresholds['clip_score_drop']:.4f}"
    )
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations

 import base64
@@ -7,18 +10,26 @@ import logging
 import os
 import tempfile
 from pathlib import Path
+from typing import Any

 import pytest
 import requests
 import yaml
 from PIL import Image

-from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
+from tests.e2e.accuracy.helpers import (
+    SimilarityThresholds,
+    assert_images_pixel_close,
+    assert_similarity,
+    model_output_dir,
+    resolve_device_profile,
+    resolve_similarity_thresholds,
+)
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import get_asset_path
 from tests.helpers.runtime import OmniServer

-pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+pytestmark = [pytest.mark.diffusion]

 logger = logging.getLogger(__name__)

@@ -31,15 +42,23 @@ WIDTH = 1024
 PROMPT = "A brown and white dog is running on the grass."
 MEAN_THRESHOLD = 3e-2
 P99_THRESHOLD = 3e-1
-SSIM_THRESHOLD = 0.97
-PSNR_THRESHOLD_CUDA = 30.0
 PSNR_THRESHOLD_NPU = 26.0

+# Per-device SSIM/PSNR for online/offline vs baseline. Unlisted devices use ``default``.
+SIMILARITY_THRESHOLDS_BY_DEVICE: dict[str, SimilarityThresholds] = {
+    "B200": SimilarityThresholds(ssim=0.96, psnr=28.0),
+    "default": SimilarityThresholds(ssim=0.97, psnr=30.0),
+}

-def _psnr_threshold() -> float:
+
+def _psnr_threshold(thresholds: SimilarityThresholds | None = None) -> float:
    from vllm_omni.platforms import current_omni_platform

-    return PSNR_THRESHOLD_NPU if current_omni_platform.is_npu() else PSNR_THRESHOLD_CUDA
+    if current_omni_platform.is_npu():
+        return PSNR_THRESHOLD_NPU
+    if thresholds is None:
+        thresholds = resolve_similarity_thresholds(SIMILARITY_THRESHOLDS_BY_DEVICE)
+    return thresholds.psnr


 _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
@@ -57,7 +76,7 @@ def _baseline_path() -> Path:
 _OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"

 # DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml).
-_DEPLOY_CONFIG = {
+_DEPLOY_CONFIG: dict[str, Any] = {
    "pipeline": "hunyuan_image3_dit",
    "async_chunk": False,
    "trust_remote_code": True,
@@ -121,7 +140,7 @@ def _devices() -> str:


 def _write_deploy_config(path: Path) -> None:
-    config = copy.deepcopy(_DEPLOY_CONFIG)
+    config: dict[str, Any] = copy.deepcopy(_DEPLOY_CONFIG)
    devices = _devices()
    num_devices = len(devices.split(","))
    config["stages"][0]["devices"] = devices
@@ -253,6 +272,15 @@ def _run_vllm_omni_hunyuan_image3_offline(*, model: str, deploy_config: str, out
 def _assert_against_baseline(image: Image.Image, label: str) -> None:
    baseline_path = _baseline_path()
    baseline_image = Image.open(baseline_path).convert("RGB")
+    device_profile = resolve_device_profile(profiles=SIMILARITY_THRESHOLDS_BY_DEVICE)
+    thresholds = resolve_similarity_thresholds(SIMILARITY_THRESHOLDS_BY_DEVICE, device_profile)
+    psnr_threshold = _psnr_threshold(thresholds)
+    logger.info(
+        "HunyuanImage3 pixel accuracy thresholds: profile=%s ssim>=%s psnr>=%s",
+        device_profile,
+        thresholds.ssim,
+        psnr_threshold,
+    )

    assert_images_pixel_close(
        model_name=f"{MODEL_NAME} ({label} vs baseline)",
@@ -265,11 +293,12 @@ def _assert_against_baseline(image: Image.Image, label: str) -> None:
        model_name=f"{MODEL_NAME} ({label} vs baseline)",
        vllm_image=image,
        diffusers_image=baseline_image,
-        ssim_threshold=SSIM_THRESHOLD,
-        psnr_threshold=_psnr_threshold(),
+        ssim_threshold=thresholds.ssim,
+        psnr_threshold=psnr_threshold,
    )


+@pytest.mark.full_model
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4)
 def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> None:
    model = _model_name()
@@ -284,6 +313,7 @@ def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> N
    _assert_against_baseline(image, "online")


+@pytest.mark.full_model
@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4)
 def test_hunyuan_image3_pixel_accuracy_offline(accuracy_artifact_root: Path) -> None:
    model = _model_name()
@@ -6,54 +6,24 @@ from __future__ import annotations
 import os
 import subprocess
 import sys
-from base64 import b64decode
 from pathlib import Path

 import pytest
-import requests
 import torch
-from diffusers import UniPCMultistepScheduler
-from PIL import Image

 from tests.e2e.accuracy.helpers import (
    assert_video_metadata,
    assert_video_similarity_metrics,
+    build_online_image_reference,
+    online_timeout_seconds,
+    probe_binary,
+    probe_video,
+    resolve_device_profile,
+    resolve_similarity_thresholds,
+    send_video_request_with_timeout,
+    validate_image_source,
    video_artifact_dir,
 )
-from tests.e2e.accuracy.helpers import (
-    build_online_image_reference as _build_online_image_reference,
-)
-from tests.e2e.accuracy.helpers import (
-    online_timeout_seconds as _online_timeout_seconds,
-)
-from tests.e2e.accuracy.helpers import (
-    parse_psnr_score as _parse_psnr_score,
-)
-from tests.e2e.accuracy.helpers import (
-    parse_ssim_score as _parse_ssim_score,
-)
-from tests.e2e.accuracy.helpers import (
-    parse_video_metadata as _parse_video_metadata,
-)
-from tests.e2e.accuracy.helpers import (
-    probe_binary as _probe_binary,
-)
-from tests.e2e.accuracy.helpers import (
-    probe_video as _probe_video,
-)
-from tests.e2e.accuracy.helpers import (
-    send_video_request_with_timeout as _send_video_request_with_timeout,
-)
-from tests.e2e.accuracy.helpers import (
-    validate_image_source as _validate_image_source,
-)
-from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
-    _configure_scheduler,
-    _ensure_wan_ftfy_fallback,
-    _IdentityFtfy,
-    _offline_cuda_device,
-    _resize_to_target,
-)
 from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import (
    FLOW_SHIFT,
    FPS,
@@ -65,178 +35,16 @@ from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import (
    NUM_FRAMES,
    NUM_INFERENCE_STEPS,
    PROMPT,
-    PSNR_THRESHOLD,
    RABBIT_IMAGE_URL,
    SEED,
+    SIMILARITY_THRESHOLDS_BY_DEVICE,
    SIZE,
-    SSIM_THRESHOLD,
    WIDTH,
 )
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams

-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-
-def test_parse_video_metadata_extracts_dimensions_and_fps() -> None:
-    payload = {
-        "streams": [
-            {
-                "width": 832,
-                "height": 480,
-                "avg_frame_rate": "16/1",
-                "nb_read_frames": "93",
-            }
-        ]
-    }
-
-    metadata = _parse_video_metadata(payload)
-
-    assert metadata["width"] == 832
-    assert metadata["height"] == 480
-    assert metadata["fps"] == 16.0
-    assert metadata["frame_count"] == 93
-
-
-def test_parse_ssim_summary_extracts_all_score() -> None:
-    output = """
-    [Parsed_ssim_0 @ 000001] SSIM Y:0.971903 (15.512007) U:0.965077 (14.569044) V:0.962414 (14.252637) All:0.968311 (15.035654)
-    """
-
-    assert _parse_ssim_score(output) == 0.968311
-
-
-def test_parse_psnr_summary_extracts_average_score() -> None:
-    output = """
-    [Parsed_psnr_0 @ 000001] PSNR y:32.670157 u:31.844621 v:31.513839 average:32.148004 min:31.563744 max:33.201457
-    """
-
-    assert _parse_psnr_score(output) == 32.148004
-
-
-def test_build_diffusers_command_uses_python_runner_path(tmp_path: Path) -> None:
-    runner_path = tmp_path / "run_wan22_i2v_diffusers_cp.py"
-    command = _build_diffusers_command(
-        runner_path=runner_path,
-        image_source=RABBIT_IMAGE_URL,
-        output_path=tmp_path / "offline.mp4",
-        metadata_path=tmp_path / "offline.json",
-    )
-
-    assert command[:2] == [
-        sys.executable,
-        str(runner_path),
-    ]
-    assert "--output" in command
-    assert "--metadata-output" in command
-
-
-def test_resolve_image_source_prefers_existing_local_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
-    image_path = tmp_path / "rabbit.png"
-    Image.new("RGB", (8, 8), color=(255, 128, 64)).save(image_path)
-
-    assert _resolve_image_source(str(image_path)) == str(image_path.resolve())
-
-
-def test_build_online_image_reference_uses_data_url_for_local_path(tmp_path: Path) -> None:
-    image_path = tmp_path / "rabbit.png"
-    Image.new("RGB", (4, 2), color=(10, 20, 30)).save(image_path)
-
-    reference = _build_online_image_reference(str(image_path))
-
-    assert reference.startswith("data:image/png;base64,")
-    encoded = reference.split(",", 1)[1]
-    assert len(b64decode(encoded)) > 0
-
-
-def test_send_video_request_with_timeout_uses_requested_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
-    class DummyResponse:
-        def raise_for_status(self) -> None:
-            return None
-
-        def json(self) -> dict[str, str]:
-            return {"id": "video_123"}
-
-    class DummyClient:
-        def __init__(self) -> None:
-            self.run_level = "core_model"
-            self.wait_args: tuple[str, int, int] | None = None
-
-        def _build_url(self, path: str) -> str:
-            return f"http://localhost:8000{path}"
-
-        def _wait_until_video_completed(
-            self,
-            video_id: str,
-            poll_interval_seconds: int = 2,
-            timeout_seconds: int = 300,
-        ) -> None:
-            self.wait_args = (video_id, poll_interval_seconds, timeout_seconds)
-
-        def _download_video_content(self, video_id: str) -> bytes:
-            return b"video-bytes"
-
-    monkeypatch.setattr(requests, "post", lambda *args, **kwargs: DummyResponse())
-
-    client = DummyClient()
-    result = _send_video_request_with_timeout(
-        client,
-        {
-            "form_data": {"prompt": "test"},
-            "image_reference": RABBIT_IMAGE_URL,
-        },
-        timeout_seconds=1200,
-    )
-
-    assert result == b"video-bytes"
-    assert client.wait_args == ("video_123", 2, 1200)
-
-
-def test_online_timeout_defaults_to_1200() -> None:
-    assert _online_timeout_seconds(None) == 1200
-
-
-def test_artifact_dir_is_under_repo_result_folder(tmp_path: Path) -> None:
-    artifact_dir = _artifact_dir(str(tmp_path / "rabbit.png"))
-
-    assert artifact_dir.parent == Path(__file__).parent / "result"
-    assert artifact_dir.name.startswith("rabbit-")
-
-
-def test_offline_cuda_device_uses_indexed_cuda_device() -> None:
-    assert _offline_cuda_device() == torch.device("cuda:0")
-
-
-def test_ensure_wan_ftfy_fallback_sets_identity(monkeypatch: pytest.MonkeyPatch) -> None:
-    from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module
-
-    monkeypatch.delattr(wan_i2v_module, "ftfy", raising=False)
-    _ensure_wan_ftfy_fallback()
-
-    assert hasattr(wan_i2v_module, "ftfy")
-    assert isinstance(wan_i2v_module.ftfy, _IdentityFtfy)
-    assert wan_i2v_module.ftfy.fix_text("abc") == "abc"
-
-
-def test_resize_to_target_matches_requested_dimensions() -> None:
-    image = Image.new("RGB", (100, 50), color=(1, 2, 3))
-
-    resized = _resize_to_target(image, width=832, height=480)
-
-    assert resized.size == (832, 480)
-
-
-def test_configure_scheduler_switches_to_unipc() -> None:
-    class DummyPipe:
-        def __init__(self) -> None:
-            self.scheduler = UniPCMultistepScheduler(num_train_timesteps=1000, flow_shift=1.0)
-
-    pipe = DummyPipe()
-    _configure_scheduler(pipe, flow_shift=5.0)
-
-    assert isinstance(pipe.scheduler, UniPCMultistepScheduler)
-    assert pipe.scheduler.config.flow_shift == 5.0
-
+pytestmark = [pytest.mark.diffusion]

 REPO_ROOT = Path(__file__).resolve().parents[4]
 WORKSPACE_ROOT = REPO_ROOT.parent
@@ -312,7 +120,7 @@ def _runner_env() -> dict[str, str]:
    if existing_pythonpath:
        pythonpath_parts.append(existing_pythonpath)
    env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
-    env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
    return env


@@ -324,12 +132,8 @@ def _resolve_image_source(configured: str | None) -> str:
    return configured


-def _artifact_dir(image_source: str) -> Path:
-    return video_artifact_dir(RESULT_ROOT, image_source)
-
-
 def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
-    artifact_dir = _artifact_dir(image_source)
+    artifact_dir = video_artifact_dir(RESULT_ROOT, image_source)
    artifact_dir.mkdir(parents=True, exist_ok=True)
    return (
        artifact_dir / "online.mp4",
@@ -343,7 +147,7 @@ def _generate_online_video(
    omni_server,
    online_client,
    image_source: str,
-    online_timeout_seconds: int,
+    online_timeout_seconds_value: int,
 ) -> Path:
    online_path, _, _ = _artifact_paths(image_source)
    request_config = {
@@ -360,12 +164,12 @@ def _generate_online_video(
            "num_inference_steps": NUM_INFERENCE_STEPS,
            "seed": SEED,
        },
-        "image_reference": _build_online_image_reference(image_source),
+        "image_reference": build_online_image_reference(image_source),
    }
-    online_video_bytes = _send_video_request_with_timeout(
+    online_video_bytes = send_video_request_with_timeout(
        online_client,
        request_config,
-        timeout_seconds=_online_timeout_seconds(online_timeout_seconds),
+        timeout_seconds=online_timeout_seconds(online_timeout_seconds_value),
    )
    online_path.write_bytes(online_video_bytes)
    return online_path
@@ -390,6 +194,7 @@ def _generate_offline_video(*, image_source: str) -> tuple[Path, Path]:


@pytest.mark.benchmark
+@pytest.mark.full_model
@hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_wan22_i2v_diffusers_offline_generates_video(
    wan22_i2v_image_source: str | None,
@@ -397,20 +202,21 @@ def test_wan22_i2v_diffusers_offline_generates_video(
    if not torch.cuda.is_available():
        pytest.skip("Wan2.2 I2V diffusers offline test requires CUDA.")

-    _probe_binary("ffprobe")
+    probe_binary("ffprobe")
    if not RUNNER_PATH.exists():
        raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}")

    image_source = _resolve_image_source(wan22_i2v_image_source)
-    _validate_image_source(image_source)
+    validate_image_source(image_source)
    offline_path, offline_metadata_path = _generate_offline_video(image_source=image_source)
    assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
    assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"
-    offline_metadata = _probe_video(offline_path)
+    offline_metadata = probe_video(offline_path)
    assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)


@pytest.mark.benchmark
+@pytest.mark.full_model
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
 def test_wan22_i2v_online_serving_generates_video(
@@ -422,21 +228,22 @@ def test_wan22_i2v_online_serving_generates_video(
    if not torch.cuda.is_available() or torch.accelerator.device_count() < 2:
        pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")

-    _probe_binary("ffprobe")
+    probe_binary("ffprobe")
    image_source = _resolve_image_source(wan22_i2v_image_source)
-    _validate_image_source(image_source)
+    validate_image_source(image_source)
    online_path = _generate_online_video(
        omni_server=omni_server,
        online_client=online_client,
        image_source=image_source,
-        online_timeout_seconds=wan22_i2v_online_timeout_seconds,
+        online_timeout_seconds_value=wan22_i2v_online_timeout_seconds,
    )
    assert online_path.exists(), f"Expected online video artifact at {online_path}"
-    online_metadata = _probe_video(online_path)
+    online_metadata = probe_video(online_path)
    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)


@pytest.mark.benchmark
+@pytest.mark.full_model
@hardware_test(res={"cuda": "H100"}, num_cards=2)
 def test_wan22_i2v_serving_matches_diffusers_video_similarity(
    wan22_i2v_image_source: str | None,
@@ -444,13 +251,13 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity(
    if not torch.cuda.is_available() or torch.accelerator.device_count() < 2:
        pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")

-    _probe_binary("ffmpeg")
-    _probe_binary("ffprobe")
+    probe_binary("ffmpeg")
+    probe_binary("ffprobe")
    if not RUNNER_PATH.exists():
        raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}")

    image_source = _resolve_image_source(wan22_i2v_image_source)
-    _validate_image_source(image_source)
+    validate_image_source(image_source)
    online_path, offline_path, offline_metadata_path = _artifact_paths(image_source)

    if not online_path.exists():
@@ -462,8 +269,8 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity(
    assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
    assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"

-    online_metadata = _probe_video(online_path)
-    offline_metadata = _probe_video(offline_path)
+    online_metadata = probe_video(online_path)
+    offline_metadata = probe_video(offline_path)
    assert online_metadata == offline_metadata, (
        f"Video metadata mismatch:\n"
        f"online={online_metadata}\n"
@@ -472,11 +279,17 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity(
        f"offline_path={offline_path}"
    )
    assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
+    device_profile = resolve_device_profile(profiles=SIMILARITY_THRESHOLDS_BY_DEVICE)
+    thresholds = resolve_similarity_thresholds(SIMILARITY_THRESHOLDS_BY_DEVICE, device_profile)
+    print(
+        f"wan22_i2v similarity thresholds: device_profile={device_profile}, "
+        f"ssim>={thresholds.ssim:.6f}, psnr>={thresholds.psnr:.6f}"
+    )
    assert_video_similarity_metrics(
        label="wan22_i2v",
        online_path=online_path,
        offline_path=offline_path,
-        ssim_threshold=SSIM_THRESHOLD,
-        psnr_threshold=PSNR_THRESHOLD,
+        ssim_threshold=thresholds.ssim,
+        psnr_threshold=thresholds.psnr,
    )
    print(f"offline_metadata={offline_metadata_path}")
@@ -1,5 +1,7 @@
 from __future__ import annotations

+from tests.e2e.accuracy.helpers import SimilarityThresholds
+
 RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png"
 MODEL_NAME = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
 SIZE = "832x480"
@@ -13,8 +15,16 @@ FLOW_SHIFT = 5.0
 NUM_INFERENCE_STEPS = 50
 SEED = 42
 BOUNDARY_RATIO = 0.875
-SSIM_THRESHOLD = 0.94
-PSNR_THRESHOLD = 28.0
+
+
+# Per-device overrides for online-vs-diffusers video similarity.
+# Only listed profiles are matched against ``get_device_name()``; everything
+# else (including H100/L4) uses ``default``.
+SIMILARITY_THRESHOLDS_BY_DEVICE: dict[str, SimilarityThresholds] = {
+    "B200": SimilarityThresholds(ssim=0.93, psnr=28.0),
+    "default": SimilarityThresholds(ssim=0.94, psnr=28.0),
+}
+

 PROMPT = """一只棕色野兔的正面特写镜头，采用低角度仰拍视角，营造亲密而庄严的视觉冲击。兔子一双圆润漆黑的大眼睛直视镜头深处，眼神中交织着野生动物的警觉与一丝难以言喻的温柔好奇，仿佛在与观者建立跨越物种的静默对话。它毛色呈现层次丰富的棕褐渐变，从浅奶油色腹部过渡到深棕背部，每根毛发纹理清晰可辨，在侧光下泛着丝绸般的光泽。细长洁白的胡须共有三对，随呼吸节奏微微颤动，偶尔因捕捉气流信息而轻轻摇摆。
 两只标志性的长耳完全竖立，耳廓外侧覆盖短密棕毛，内侧则露出粉嫩的血管网络，薄如蝉翼的皮肤下血液流动隐约可见，耳朵以细微幅度不时转动，精准定位声源方向。背景是一片澄澈的蔚蓝天空，形态蓬松的白色积云以缓慢速度横向漂移，云影在兔子头顶交替变化，光线随之明暗流转。晴朗天气的明媚阳光从画面左上方45度角倾泻而下，在兔脸右侧形成柔和的伦勃朗式阴影，强化了面部立体感和皮毛质感。
@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.experimental.fullduplex.core import protocol as ev
+from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig
+from vllm_omni.experimental.fullduplex.mage_vl import (
+    MageVLCodecWindow,
+    MageVLDuplexAdapter,
+    MageVLDuplexRuntime,
+    MageVLGateDecision,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+async def _feed(events):
+    for event in events:
+        yield event
+
+
+def _collector():
+    out: list[dict] = []
+
+    async def emit(event: dict) -> None:
+        out.append(event)
+
+    return out, emit
+
+
+def test_mage_vl_capabilities_match_backend_contract():
+    capabilities = MageVLDuplexAdapter().capabilities()
+
+    assert capabilities.input_modalities == frozenset({"text", "video"})
+    assert capabilities.output_modalities == frozenset({"text"})
+    assert capabilities.proactive
+
+
+@pytest.mark.asyncio
+async def test_mage_vl_proactive_gate_speaks_once_for_event():
+    async def gate(session, windows):
+        del session
+        if windows[-1].segment_id == "goal":
+            return MageVLGateDecision(True, text="Goal scored.", event_id="goal-1", score=0.91)
+        return MageVLGateDecision(False)
+
+    adapter = MageVLDuplexAdapter(gate=gate, window_size=1)
+    session = DuplexSession(
+        "mage",
+        DuplexSessionConfig(
+            input_modalities=("video", "text"),
+            output_modalities=("text",),
+            proactive=True,
+        ),
+    )
+    rt = MageVLDuplexRuntime(session, adapter)
+    out, emit = _collector()
+
+    await rt.run(
+        _feed(
+            [
+                {"type": ev.INPUT_APPEND, "modality": "video", "data": {"segment_id": "warmup", "frames": ["f0"]}},
+                {"type": ev.INPUT_APPEND, "modality": "video", "data": {"segment_id": "goal", "frames": ["f1"]}},
+                {"type": ev.INPUT_APPEND, "modality": "video", "data": {"segment_id": "goal", "frames": ["f2"]}},
+                {"type": ev.CLOSE},
+            ]
+        ),
+        emit,
+    )
+
+    deltas = [event["data"] for event in out if event["type"] == ev.RESPONSE_DELTA]
+    assert deltas == ["Goal scored."]
+
+
+@pytest.mark.asyncio
+async def test_mage_vl_query_bypasses_silent_gate_and_preserves_codec_metadata():
+    seen: dict[str, object] = {}
+
+    async def gate(session, windows):
+        del session, windows
+        return {"should_respond": False}
+
+    async def generate(session, windows, query, gate_decision):
+        del session, gate_decision
+        seen["query"] = query
+        seen["window"] = windows[-1]
+        yield "The player is moving downfield."
+
+    adapter = MageVLDuplexAdapter(gate=gate, generate=generate, window_size=1)
+    session = DuplexSession(
+        "mage",
+        DuplexSessionConfig(
+            input_modalities=("video", "codec_window", "text"),
+            output_modalities=("text",),
+            proactive=True,
+        ),
+    )
+    rt = MageVLDuplexRuntime(session, adapter)
+    out, emit = _collector()
+
+    await rt.run(
+        _feed(
+            [
+                {"type": ev.INPUT_APPEND, "modality": "text", "data": "What changed?"},
+                {
+                    "type": ev.INPUT_APPEND,
+                    "modality": "video",
+                    "data": {
+                        "kind": "h264",
+                        "segment_id": "seg-7",
+                        "pts_ms": 7000,
+                        "duration_ms": 1000,
+                        "codec": {"motion_vectors": "mv", "residual_energy": "res"},
+                        "metadata": {"gop": "P"},
+                    },
+                },
+                {"type": ev.CLOSE},
+            ]
+        ),
+        emit,
+    )
+
+    window = seen["window"]
+    assert seen["query"] == "What changed?"
+    assert isinstance(window, MageVLCodecWindow)
+    assert window.kind == "h264"
+    assert window.segment_id == "seg-7"
+    assert window.pts_ms == 7000
+    assert window.metadata == {"gop": "P"}
+    assert [event["data"] for event in out if event["type"] == ev.RESPONSE_DELTA] == ["The player is moving downfield."]
+
+
+@pytest.mark.asyncio
+async def test_mage_vl_session_state_is_isolated():
+    seen: dict[str, tuple[MageVLCodecWindow, ...]] = {}
+
+    async def generate(session, windows, query, gate_decision):
+        del query, gate_decision
+        seen[session.session_id] = tuple(windows)
+        yield session.session_id
+
+    adapter = MageVLDuplexAdapter(generate=generate, window_size=1)
+    session_a = DuplexSession(
+        "mage-a",
+        DuplexSessionConfig(input_modalities=("video", "text"), output_modalities=("text",), proactive=True),
+    )
+    session_b = DuplexSession(
+        "mage-b",
+        DuplexSessionConfig(input_modalities=("video", "text"), output_modalities=("text",), proactive=True),
+    )
+
+    await adapter.on_input(session_a, "video", {"segment_id": "a", "frames": ["a0"]})
+    await adapter.on_input(session_b, "video", {"segment_id": "b", "frames": ["b0"]})
+    await adapter.on_input(session_a, "text", "describe a")
+    await adapter.on_input(session_b, "text", "describe b")
+    assert adapter.should_respond(session_a)
+    assert adapter.should_respond(session_b)
+
+    async for _ in adapter.respond(session_a):
+        pass
+    async for _ in adapter.respond(session_b):
+        pass
+    assert [window.segment_id for window in seen["mage-a"]] == ["a"]
+    assert [window.segment_id for window in seen["mage-b"]] == ["b"]
@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import asyncio
+import time
+
+import pytest
+from fastapi.testclient import TestClient
+from starlette.websockets import WebSocketDisconnect
+
+from vllm_omni.experimental.fullduplex.mage_vl import MageVLDuplexAdapter
+from vllm_omni.experimental.fullduplex.mage_vl.serving import backend as backend_module
+from vllm_omni.experimental.fullduplex.mage_vl.serving.server import (
+    MageVLServingConfig,
+    create_app,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_mage_vl_codec_backend_requires_preprocessor(monkeypatch):
+    monkeypatch.setattr(backend_module.shutil, "which", lambda _name: None)
+
+    with pytest.raises(RuntimeError, match="cv-preinfer"):
+        backend_module.MageVLTransformersBackend(video_backend="codec")
+
+
+def test_mage_vl_frames_backend_does_not_require_preprocessor(monkeypatch):
+    def unexpected_lookup(_name):
+        raise AssertionError("frames backend must not inspect codec dependencies")
+
+    monkeypatch.setattr(backend_module.shutil, "which", unexpected_lookup)
+    backend = backend_module.MageVLTransformersBackend(video_backend="frames")
+
+    assert backend.video_backend == "frames"
+
+
+def _factory():
+    async def gate(_session, windows):
+        return {
+            "should_respond": windows[-1].segment_id == "goal",
+            "event_id": "goal-1",
+            "score": 0.9,
+        }
+
+    async def generate(_session, windows, query, _decision):
+        yield f"{query or 'proactive'}:{windows[-1].segment_id}"
+
+    return MageVLDuplexAdapter(gate=gate, generate=generate, window_size=1)
+
+
+def test_mage_vl_health_ready_and_models():
+    with TestClient(create_app(_factory)) as client:
+        assert client.get("/health").json() == {"status": "ok"}
+        assert client.get("/ready").json() == {"status": "ready", "active_sessions": 0}
+        assert client.get("/v1/models").json()["data"][0]["id"] == "microsoft/Mage-VL"
+
+
+def test_mage_vl_websocket_runs_gate_and_generation():
+    with TestClient(create_app(_factory)) as client:
+        with client.websocket_connect("/v1/mage-vl/duplex?session_id=test-session") as websocket:
+            created = websocket.receive_json()
+            assert created["type"] == "session.created"
+            assert created["session_id"] == "test-session"
+            websocket.send_json(
+                {
+                    "type": "input.append",
+                    "modality": "video",
+                    "data": {"segment_id": "goal", "frames": ["frame-1"]},
+                }
+            )
+            assert websocket.receive_json()["type"] == "response.created"
+            delta = websocket.receive_json()
+            assert delta == {
+                "type": "response.delta",
+                "response_index": 1,
+                "modality": "text",
+                "data": "proactive:goal",
+            }
+            assert websocket.receive_json()["type"] == "response.done"
+            websocket.send_json({"type": "close"})
+
+
+def test_mage_vl_websocket_requires_bearer_token():
+    app = create_app(_factory, MageVLServingConfig(auth_token="secret"))
+    with TestClient(app) as client:
+        with pytest.raises(WebSocketDisconnect) as denied:
+            with client.websocket_connect("/v1/mage-vl/duplex"):
+                pass
+        assert denied.value.code == 4401
+
+        with client.websocket_connect(
+            "/v1/mage-vl/duplex",
+            headers={"Authorization": "Bearer secret"},
+        ) as websocket:
+            assert websocket.receive_json()["type"] == "session.created"
+            websocket.send_json({"type": "close"})
+
+
+def test_mage_vl_rejects_duplicate_session_and_releases_disconnect():
+    app = create_app(_factory, MageVLServingConfig(max_sessions=1))
+    with TestClient(app) as client:
+        with client.websocket_connect("/v1/mage-vl/duplex?session_id=one") as first:
+            assert first.receive_json()["session_id"] == "one"
+            with pytest.raises(WebSocketDisconnect) as denied:
+                with client.websocket_connect("/v1/mage-vl/duplex?session_id=two") as second:
+                    second.receive_json()
+            assert denied.value.code == 4429
+            first.send_json({"type": "close"})
+
+        with client.websocket_connect("/v1/mage-vl/duplex?session_id=two") as second:
+            assert second.receive_json()["session_id"] == "two"
+            second.send_json({"type": "close"})
+
+
+def test_mage_vl_rejects_malformed_event():
+    with TestClient(create_app(_factory)) as client:
+        with client.websocket_connect("/v1/mage-vl/duplex") as websocket:
+            websocket.receive_json()
+            websocket.send_text("[]")
+            event = websocket.receive_json()
+            assert event["type"] == "error"
+            assert "JSON object" in event["message"]
+
+
+def test_mage_vl_cancel_is_not_blocked_by_slow_gate():
+    async def slow_gate(_session, _windows):
+        await asyncio.sleep(2)
+        return {"should_respond": True, "event_id": "late"}
+
+    def factory():
+        return MageVLDuplexAdapter(gate=slow_gate, window_size=1)
+
+    with TestClient(create_app(factory)) as client:
+        with client.websocket_connect("/v1/mage-vl/duplex") as websocket:
+            websocket.receive_json()
+            websocket.send_json(
+                {
+                    "type": "input.append",
+                    "modality": "video",
+                    "data": {"segment_id": "slow", "frames": ["frame"]},
+                }
+            )
+            started = time.monotonic()
+            websocket.send_json({"type": "response.cancel"})
+            event = websocket.receive_json()
+            assert time.monotonic() - started < 0.5
+            assert event["type"] == "response.cancelled"
+            websocket.send_json({"type": "close"})
@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """E2E offline inference tests for the Gepard-1.0 single-stage native-AR TTS.

 Zero-shot only (the model's default learned voice). These tests need a GPU and
@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 E2E offline tests for MiniCPM-o 4.5 model with multimodal input and audio / text output.
 """
@@ -44,7 +47,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None:

@pytest.mark.full_model
@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
    """Test processing audio, generating text output."""
@@ -57,7 +60,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None:

@pytest.mark.full_model
@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_image_to_text(omni_runner, omni_runner_handler) -> None:
    """Test processing image, generating text output."""
@@ -68,7 +71,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None:

@pytest.mark.full_model
@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_video_to_text(omni_runner, omni_runner_handler) -> None:
    """Test processing video, generating text output."""
@@ -79,7 +82,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None:

@pytest.mark.full_model
@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
    """Test processing text and generating audio through Talker and Code2Wav."""
@@ -109,7 +112,7 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None:

@pytest.mark.full_model
@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1)
+@hardware_test(res={"cuda": ["H100", "B200"], "npu": "A3"}, num_cards=1)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
 def test_video_to_audio(omni_runner, omni_runner_handler) -> None:
    """Test processing video, generating audio output."""
@@ -78,7 +78,7 @@ quant_params = [(QUANTIZED_MODEL, stage_config)]


@_SKIP_ISSUE_5652
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
 def test_text_to_text(omni_runner, offline_client):
    """Text input → text output with W4A16 quantized Qwen3-Omni."""
@@ -97,7 +97,7 @@ def test_text_to_text(omni_runner, offline_client):


@_SKIP_ISSUE_5652
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
 def test_audio_to_text(omni_runner, offline_client):
    """Audio input → text output with W4A16 quantized Qwen3-Omni."""
@@ -121,7 +121,7 @@ def test_audio_to_text(omni_runner, offline_client):


@_SKIP_ISSUE_5652
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
 def test_image_to_text(omni_runner, offline_client):
    """Image input → text output with W4A16 quantized Qwen3-Omni."""
@@ -143,7 +143,7 @@ def test_image_to_text(omni_runner, offline_client):


@_SKIP_ISSUE_5652
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
 def test_video_to_text(omni_runner, offline_client):
    """Video input → text output with W4A16 quantized Qwen3-Omni."""
@@ -165,7 +165,7 @@ def test_video_to_text(omni_runner, offline_client):


@_SKIP_ISSUE_5652
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
 def test_video_to_audio(omni_runner, offline_client):
    """Video input → audio output with W4A16 quantized Qwen3-Omni."""
@@ -186,7 +186,7 @@ def test_video_to_audio(omni_runner, offline_client):


@_SKIP_ISSUE_5652
-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
 def test_mix_to_audio(omni_runner, offline_client):
    """Mixed-modality input → audio output with W4A16 quantized Qwen3-Omni."""
@@ -72,7 +72,7 @@ def _stage_config():
 quant_params = [(QUANTIZED_MODEL, _stage_config())]


-@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=2)
@pytest.mark.parametrize("omni_runner", quant_params, indirect=True)
 def test_marlin_fallback_no_nan_collapse(omni_runner, offline_client):
    """H100 Marlin FP4 fallback: text in -> coherent text out.
@@ -6,7 +6,8 @@ L4 expansion coverage for Boogu-Image Edit and Edit-Turbo (image editing / TI2I)

 The nightly smoke module (``test_boogu_image_edit.py``) already covers single-image
 text-guided editing and the text+image double-guidance path. This expansion
-module adds the two edit paths smoke does not exercise:
+module adds the edit paths and CFG-parallel configurations smoke does not
+exercise:

 - ``test_image_only_guidance`` — ``guidance_scale=1.0`` (text CFG off) with
  ``guidance_scale_2=2.0`` (image CFG on): the image-only branch of upstream's
@@ -15,13 +16,19 @@ module adds the two edit paths smoke does not exercise:
 - ``test_align_res_output_size`` — a non-square reference with no requested
  ``height`` / ``width``: asserts the output resolution follows the reference
  (upstream ``align_res``, on by default for a single reference).
+- ``test_double_guidance_cfg_parallel`` — three-branch double guidance with
+  CFG sizes 2 and 3.
+- ``test_cfg_off_on_cfg_parallel_server`` — both scales at 1.0 on a size-2 CFG
+  server.

 Both Edit checkpoints share ``BooguImagePipeline`` with the Base checkpoint;
 the TI2I path activates when the request carries a reference image. Only a
 single reference image is supported. The regular Edit cases stay smoke-depth
 (``num_inference_steps=2``), while Edit-Turbo uses its tuned four-step
-configuration. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet
-supported for these models and are intentionally omitted.
+configuration. Regular Edit CFG-parallel coverage includes two-card
+text-only/image-only/CFG-off paths, two-card round-robin double guidance, and
+full three-card double guidance. Edit-Turbo remains single-card-only here;
+CPU offload and Cache-DiT are intentionally omitted for both checkpoints.

 From ``tests/``::

@@ -47,16 +54,49 @@ EDIT_TURBO_REVISION = "hotfix-1k-20260708"
 EDIT_PROMPT = "Change the style to a colored pencil drawing."

 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+PARALLEL_2_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+PARALLEL_3_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=3)


 def _get_diffusion_feature_cases(model: str):
-    """Return the single-card server configuration supported for Boogu-Image-Edit today."""
+    """Return single-card and two-card CFG configurations for two-branch paths."""
    return [
        pytest.param(
            OmniServerParams(model=model),
            id="default",
            marks=SINGLE_CARD_FEATURE_MARKS,
        ),
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]),
+            id="cfg_parallel_2",
+            marks=PARALLEL_2_FEATURE_MARKS,
+        ),
+    ]
+
+
+def _get_double_guidance_cases(model: str):
+    """Return two-rank round-robin and full three-rank double-guidance rows."""
+    return [
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]),
+            id="cfg_parallel_2_round_robin",
+            marks=PARALLEL_2_FEATURE_MARKS,
+        ),
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "3"]),
+            id="cfg_parallel_3_full",
+            marks=PARALLEL_3_FEATURE_MARKS,
+        ),
+    ]
+
+
+def _get_cfg_off_cases(model: str):
+    return [
+        pytest.param(
+            OmniServerParams(model=model, server_args=["--cfg-parallel-size", "2"]),
+            id="cfg_parallel_2_cfg_off",
+            marks=PARALLEL_2_FEATURE_MARKS,
+        )
    ]


@@ -106,6 +146,42 @@ def test_image_only_guidance(omni_server: OmniServer, openai_client: OpenAIClien
    openai_client.send_diffusion_request(request_config)


+@pytest.mark.parametrize("omni_server", _get_double_guidance_cases(MODEL), indirect=True)
+def test_double_guidance_cfg_parallel(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """Double guidance exercises three branches on two and three CFG ranks."""
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "num_inference_steps": 2,
+            "guidance_scale": 5.0,
+            "guidance_scale_2": 2.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.parametrize("omni_server", _get_cfg_off_cases(MODEL), indirect=True)
+def test_cfg_off_on_cfg_parallel_server(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
+    """A CFG process group must not force negative embeddings or guidance."""
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "num_inference_steps": 2,
+            "guidance_scale": 1.0,
+            "guidance_scale_2": 1.0,
+            "seed": 42,
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
@pytest.mark.parametrize(
    ("omni_server", "sampling_params"),
    _get_align_res_cases(),
Author	SHA1	Message	Date
hsliuustc0106	52c8f8dbc8	Support one-sided MoT biases Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-09-03 17:05:57 +00:00
WeiQing Chen	dff28e6867	[Doc] Add production diffusion model skill (#6097 ) Signed-off-by: David Chen <530634352@qq.com>	2026-09-03 20:09:28 +08:00
Mu GuanLin	1dae537313	[Bugfix][LTX] Restore distilled two-stage video serving (#6847 ) Signed-off-by: mglyn <1203789601@qq.com>	2026-09-03 18:48:43 +08:00
Eugene Zhao	ec22c7ac06	[Bugfix] Restore Omni health route ownership (#6723 ) Signed-off-by: qujing226 <qujing226@gmail.com> Signed-off-by: Eugene Zhao <113281545+qujing226@users.noreply.github.com>	2026-09-03 00:54:36 -07:00
Jim Ban	711daacabf	[Bugfix][TTS] Make CosyVoice3 TRT plan publication concurrency-safe (#6955 ) Signed-off-by: BANANASJIM <bananasjim1@gmail.com>	2026-09-03 06:44:42 +00:00
Chiquan MoandMciG-ggg	2b47faf601	[Bugfix] Validate --stage-overrides shape in parse_stage_overrides (drop field allowlist) (#6230 ) Signed-off-by: Chiquan Mo <mcigggg399026@gmail.com> Signed-off-by: MciG-ggg <MciG-ggg@users.noreply.github.com> Co-authored-by: MciG-ggg <MciG-ggg@users.noreply.github.com>	2026-09-02 23:19:08 -07:00
Tai An	d1a5a7c178	[Bugfix][Qwen2.5-Omni] Fix two stale call sites in the speech helpers (#6886 ) Signed-off-by: Anai-Guo <antai12232931@outlook.com>	2026-09-02 21:41:05 -07:00
Joshna-MedisettyandCursor	24c49cf392	[Bugfix][XPU] Add MoTRMSNorm.forward_xpu delegating to forward_native (#5568 ) Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com> Signed-off-by: Joshna-Medisetty <joshna.medisetty@intel.com> Co-authored-by: Cursor <cursoragent@cursor.com>	2026-09-02 23:06:17 -05:00
ZhanshenzhuiandSy03	48ae412a2f	[model]Add experimental Mage-VL full-duplex adapter and offline integration (#6537 ) Signed-off-by: Zhanshenzhui <99891002+zyforsure@users.noreply.github.com> Co-authored-by: Sy03 <1370724210@qq.com>	2026-09-02 20:25:32 -07:00
Zeng Chuang	b5e898f295	[Feat] Add support and testcase for HunyuanImage-3.0-Instruct-Distil (#4048 ) Signed-off-by: zengchuang <zengchuang3@huawei.com> Signed-off-by: zengchuang-hw <zengchuang1005@gmail.com>	2026-09-03 02:43:13 +00:00
0z5aand0z5a	4927fd5412	[Core] Optimize diffusion LoRA expand accumulation (#6268 ) Signed-off-by: 0z5a <0z5a@users.noreply.github.com> Signed-off-by: 0z5a <dinadegttyarenko6@gmail.com> Co-authored-by: 0z5a <0z5a@users.noreply.github.com>	2026-09-03 10:37:35 +08:00
hurukawa	8f284b343a	[Feature][Diffusion]: CFG parallelism support for Boogu-Image (#6786 ) Signed-off-by: nagisa <1434936049@qq.com> Signed-off-by: nagisa-kun <1434936049@qq.com>	2026-09-03 10:37:21 +08:00
Rahul SteigerandRahul Steiger	376b28aced	[Bugfix][Cosmos3] Support multi-chunk transfer with distributed VAE (#6920 ) Signed-off-by: Rahul Steiger <rsteiger@nvl72d147-T16.cm.cluster> Co-authored-by: Rahul Steiger <rsteiger@nvl72d147-T16.cm.cluster>	2026-09-03 09:01:53 +08:00
maithilijoshi20andmaithilijoshi20	63ed0fdef4	[Bugfix] Fix vLLM serve error response imports (#6707 ) Signed-off-by: maithilijoshi20 <maithilijoshi2003@gmail.com> Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com> Co-authored-by: maithilijoshi20 <maithilijoshi2003@gmail.com>	2026-09-02 15:48:01 -07:00
DanaerLee	ca6b4e2112	[Perf][Bugfix][OmniVoice] Restore float16 serving, and fuse the generator hot loop (#6317 ) Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>	2026-09-02 19:54:55 +00:00
Jim Ban	b9a88bf51e	[Bugfix][TTS] Isolate concurrent StepAudio2 TRT builds (#6957 ) Signed-off-by: BANANASJIM <bananasjim1@gmail.com>	2026-09-02 19:53:02 +00:00
GXINandHongsheng Liu	3ab582c993	[Bugfix][Qwen3-TTS] Reject task mismatches before EngineCore dispatch (#6113 ) Signed-off-by: XIN GAO <1037396230@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-09-02 11:16:55 -07:00
vOv Hongsheng Liu Alicia	3b6ce98421	[Model] Add Cache-DiT and CPU offload support for SANA-Video 2B (#5882 ) Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>	2026-09-02 23:40:48 +08:00
wangyu lishunyang12 Cursor Bo Li	95a9311696	[CI] Select mirror_hardwares presets via MIRROR_HW (H100/B200) (#5543 ) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: lishunyang12 <lishunyang12@163.com> Signed-off-by: [Your Name] <your.email@example.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com> Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> Co-authored-by: lishunyang12 <lishunyang12@163.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Bo Li <22713281+bobboli@users.noreply.github.com>	2026-09-02 15:39:49 +00:00
wangyu	d5fd9eb6bb	[Bugfix][Diffusion] Register Sana I2V pipeline as video output (#6953 ) Signed-off-by: wangyu <410167048@qq.com>	2026-09-02 23:16:58 +08:00
Bo Li	5d20f6b900	[Doc] Refine TRTLLM attention guidance (#6724 ) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>	2026-09-02 19:42:37 +08:00
汪志鹏	c59b48b47e	[Feature][MiniMax-H3] Complete VSA and Ulysses support (#6909 ) Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>	2026-09-02 07:19:21 -04:00
vOv	57c08d34c3	[Model] Add TP and CFG parallelism to SANA-Video 2B (#5861 ) Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>	2026-09-02 17:44:41 +08:00
wangyuandamy-why-3459	b6e72f3879	[CI] Migrate L4 jobs to l4-k8s and split by GPU count (#6890 ) Signed-off-by: wangyu <410167048@qq.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>	2026-09-02 17:09:59 +08:00
jingchengtian	67ef0b64b9	[Hardware][Ascend] Use npu_rotary_mul fused kernel for MOSS-TTS codec RoPE (#6908 ) Signed-off-by: jingchengtian <tjc1995@126.com>	2026-09-02 08:58:23 +00:00
XenorynandClaude	e51fe6ec1b	[Bugfix][Diffusion] Keep the GQA ratio when padding Ulysses heads (#5716 ) Signed-off-by: Xenoryn <timberh2000@outlook.com> Co-authored-by: Claude <noreply@anthropic.com>	2026-09-02 08:24:51 +08:00
leaandHongsheng Liu	5f49848536	[Bugfix][Ascend NPU][Boogu] Replace complex RoPE with real-valued operations (#6571 ) Signed-off-by: lea_lei <JiaWen_Lei_7@163.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-09-02 07:11:01 +08:00
Arifuzzaman Joy	3cddda9781	[Bugfix] Prevent build_engine_args_dict from mutating stage_config.engine_args (#6783 ) Signed-off-by: Arifuzzamanjoy <s1710374103@ru.ac.bd>	2026-09-01 16:45:28 +00:00
				`@@ -0,0 +1 @@`
				首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，<recaption>这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>
				`@@ -0,0 +1 @@`
				首先，我分析所有输入图像：图像1是一个圆形logo，背景为黑色，主体由深蓝、中蓝、浅蓝和白色组成的抽象曲线构成，整体呈现出一种动态的、类似旋涡的视觉效果。图像2展示了四个并排摆放在灰色织物表面上的冰箱贴。这些冰箱贴的造型分别是长颈鹿、小鸟、蜜蜂和海豚，其关键特征是材质：它们具有光滑、高反光的珐琅质感，边缘有明显的银色金属包边，呈现出立体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）转化为一个具有特定三维物理属性（冰箱贴材质）的实体物品。核心在于保留logo的设计精髓，同时精确复刻参考图像中冰箱贴的材质和工艺感。为了构建最终的答案图像，我会将图一的圆形logo作为新冰箱贴的主体图案。这个圆形冰箱贴将拥有与图二冰箱贴一致的厚度和立体感。其表面会覆盖一层高光泽的透明釉质，完美再现珐琅的质感。logo内部的蓝色和白色区域将保持其原有的色彩，但会因为表面的光泽而出现明亮的高光反射。整个圆形的边缘会包裹一圈纤细而有光泽的银色金属边框，这不仅增加了产品的精致感，也与图二中的金属包边工艺相呼应。最后，将这个新制作的圆形冰箱贴放置在图二所示的灰色织物背景上，确保其光照、阴影和反射效果与周围环境和其他冰箱贴完全一致，使其看起来就像是和它们一起被拍摄的，从而创造出一个逼真且风格统一的视觉效果。</think><recaption>这幅图像以产品摄影的精准风格，展示了一枚根据`image_1`的logo设计定制的圆形冰箱贴。最终图像使用`image_2`的分辨率。该冰箱贴居中放置于`image_2`的灰色织物背景之上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝及白色构成的动态旋涡状抽象图案。整个冰箱贴被赋予了与`image_2`中冰箱贴完全一致的高级珐琅质感，表面覆盖着一层晶莹剔透的高光釉面，在柔和的顶光照射下，呈现出清晰而温润的高光反射。一圈纤薄而精致的银色金属边框紧密包裹着圆形主体，其金属光泽与釉面的质感形成精妙对比。冰箱贴在织物背景上投下了一层极其自然的柔和阴影，证明了其与环境光影的无缝融合，整体视觉效果真实可信，宛如一件精心制作的实体工艺品。<relation_1>最终图像完整保留了`image_1`中logo的全部核心设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色四种颜色构成的抽象旋涡图案的精确形态、色彩和布局。</relation_1><relation_2>最终图像的分辨率、背景和材质风格完全来自`image_2`。这包括其作为背景的灰色织物表面的独特纹理，以及应用于新冰箱贴上的、与`image_2`中冰箱贴完全相同的珐琅材质（高光釉面质感）和银色金属边框的物理属性。画面的整体光照环境，包括柔和的光源方向和质感，也与`image_2`保持一致。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep><guidance><timestep_r>[<img>]{3600}<eoi></answer>