[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>

.buildkite/common/ci_mirror_hardwares.yml

@@ -3,12 +3,24 @@
 # ``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_2, l4_3, 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)
+#
+# 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 /
@@ -380,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"

.buildkite/common/scripts/skip_ci.py

@@ -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
 
 

.buildkite/common/scripts/upload_pipeline.py

@@ -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,28 +254,249 @@ def _load_mirror_hardwares() -> dict[str, dict[str, Any]]:
     return presets
 
 
-def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any]:
+@lru_cache(maxsize=1)
-    """Replace uploader-only ``mirror_hardwares`` with preset fields from ci_mirror_hardwares.yml."""
+def _cuda_mirror_chips() -> tuple[str, ...]:
-    hardware = step.get("mirror_hardwares")
+    """Lowercase CUDA SKUs that have at least one ``{chip}_{n}`` preset.
-    if hardware is None:
-        return step
 
-    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}",
+            f"unknown mirror_hardwares {preset_name!r} in step {step_label!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",
         )
 
     expanded = copy.deepcopy(preset)
@@ -301,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"}
@@ -309,14 +551,11 @@ def _process_test_steps(
             processed.append(new_step)
             continue
 
-        if deps is not None:
+        leaf = {key: value for key, value in step.items() if key != "source_file_dependencies"}
-            processed.append(
+        expanded = _expand_mirror_hardwares(leaf)
-                _expand_mirror_hardwares(
+        if expanded is None:
-                    {key: value for key, value in step.items() if key != "source_file_dependencies"},
+            continue
-                ),
+        processed.append(expanded)
-            )
-        else:
-            processed.append(_expand_mirror_hardwares(step))
 
     return processed
 
@@ -342,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
@@ -371,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 if force_all or e2e_only else ctx.changed_files
-        changed_files = None
-    else:
-        changed_files = ctx.changed_files
 
     doc = yaml.safe_load(text)
     if not isinstance(doc, dict):

.buildkite/cuda/test-merge.yml

@@ -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

.buildkite/cuda/test-nightly.yml

@@ -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
@@ -15,14 +19,12 @@ steps:
       - label: ":full_moon: Omni · Function Test with H100 · Single-GPU"
         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
+          - 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
-        mirror_hardwares: h100_1
 
       - 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
+          - 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
-        mirror_hardwares: h100_2
 
       - label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
         timeout_in_minutes: 50
@@ -266,7 +268,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
@@ -369,7 +370,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"
@@ -419,7 +419,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

.buildkite/cuda/test-ready.yml

@@ -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

.buildkite/cuda/test-weekly.yml

@@ -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

docs/contributing/ci/ci_settings.md

@@ -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:
 

docs/contributing/ci/test_system_overview.md

@@ -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)).

docs/user_guide/diffusion/attention_backends/dense_backends.md

@@ -21,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
+is FA4 only (`flash_attn.cute` / `vllm-omni[fa4]`). If FA4 is unavailable,
-FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
+explicit `FLASH_ATTN` is rejected and automatic selection continues to another
-Ampere it is the preferred automatic route when a compatible package is
+compatible Blackwell backend. Hopper-only FA2/FA3 wheels are not used, even if
-installed.
+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
@@ -39,8 +40,10 @@ 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
+Blackwell. On Blackwell, `FLASH_ATTN` is FA4 only: Hopper-only FA2/FA3
-compatible FlashAttention 3 or 2 path.
+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`
 
@@ -89,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
 ```

docs/user_guide/diffusion/attention_backends/huggingface_hub.md

@@ -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
+If `kernels` is unavailable, an explicit `FLASH_ATTN_HUB` or
-corresponding local FlashAttention implementation. On CUDA devices older than
+`FLASH_ATTN_3_HUB` selection raises. Automatic / `FLASH_ATTN` selection in the
-Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
+Diffusers adapter may still walk hub then local FlashAttention backends.
 
 ## Usage
 

docs/user_guide/fault_injection_reliability_matrix.md

@@ -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 `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 `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 `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 `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 `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 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) |
 

tests/buildkite/test_skip_ci.py

@@ -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

tests/buildkite/test_upload_pipeline.py

@@ -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": [
             {
@@ -182,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()

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -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,

tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json

@@ -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,

tests/dfx/perf/tests/test_runner_metadata.py

@@ -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)
 
 

tests/diffusion/attention/test_attention_config.py

@@ -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
+        Regression: ``str(str, Enum)`` later produced ``RainFusionPrecision.MIX``
-        produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it).
+        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,
+                )

tests/diffusion/attention/test_cuda_backend_selection.py

@@ -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

tests/diffusion/attention/test_cudnn_attn.py

@@ -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(
+def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
-    ("query_sequence_length", "key_sequence_length", "expected_backend"),
+    """Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V."""
-    [
-        (1, 1, SDPBackend.MATH),
-        (8, 1, SDPBackend.MATH),
-        (1, 8, SDPBackend.MATH),
-        (8, 8, SDPBackend.CUDNN_ATTENTION),
-    ],
-)
-def test_cudnn_routes_single_token_attention_to_math(
-    monkeypatch, query_sequence_length, key_sequence_length, expected_backend
-):
     selected_backends = []
 
+    @contextmanager
     def fake_sdpa_kernel(backends):
-        selected_backends.append(backends)
+        selected_backends.append(tuple(backends))
-        return nullcontext()
+        yield
 
-    monkeypatch.setattr(
+    def fake_sdpa(query, key, value, **kwargs):
-        "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
+        return query
-        fake_sdpa_kernel,
-    )
-    monkeypatch.setattr(
-        torch.nn.functional,
-        "scaled_dot_product_attention",
-        lambda query, _key, _value, **_kwargs: query,
-    )
-    impl = CuDNNAttentionImpl(
-        num_heads=2,
-        head_size=4,
-        softmax_scale=0.5,
-    )
-    query = torch.randn(1, query_sequence_length, 2, 4)
-    key = torch.randn(1, key_sequence_length, 2, 4)
-    value = torch.randn_like(key)
 
-    output = impl.forward_cuda(query, key, value)
+    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()

tests/diffusion/attention/test_flashinfer_attn.py

@@ -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

tests/diffusion/attention/test_kernels_hub.py

@@ -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):
+def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
-    """Test that when kernels package is not available, the platform fallback logic
-
-    routes to local FLASH_ATTN.
-    """
     from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
 
     # Temporarily hide/remove kernels module if present
@@ -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
+    with pytest.raises(ImportError, match="explicitly selected"):
-    backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
+        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_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)
+    with pytest.raises(ValueError, match="require.*Hopper GPU"):
-    assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
+        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

tests/diffusion/attention/test_ring_backend_selection.py

@@ -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)

tests/diffusion/attention/test_ring_fa4.py

@@ -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)

tests/diffusion/attention/test_ring_utils.py

@@ -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))))

tests/diffusion/attention/test_sage_attn.py

@@ -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)

tests/diffusion/attention/test_sage_attn3.py

@@ -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)
+def test_sage_attn3_rejects_mask_instead_of_ignoring_it(monkeypatch: pytest.MonkeyPatch):
-    assert sdpa_calls["enable_gqa"] is True
+    def fake_kernel(*args, **kwargs):
-    expected = (query.permute(0, 2, 1, 3) + 1).permute(0, 2, 1, 3)
+        raise AssertionError("sageattn3_blackwell should not run with an unsupported mask")
-    assert torch.allclose(output, expected)
+
+    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)

tests/diffusion/attention/test_selector.py

@@ -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")

tests/diffusion/cache/test_cache_dit.py

@@ -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(
+    with _force_torch_sdpa():
-        in_channels=16,
+        model = LTX2VideoTransformer3DModel(
-        out_channels=16,
+            in_channels=16,
-        patch_size=1,
+            out_channels=16,
-        patch_size_t=1,
+            patch_size=1,
-        num_attention_heads=2,
+            patch_size_t=1,
-        attention_head_dim=8,
+            num_attention_heads=2,
-        cross_attention_dim=16,
+            attention_head_dim=8,
-        audio_in_channels=16,
+            cross_attention_dim=16,
-        audio_out_channels=16,
+            audio_in_channels=16,
-        audio_num_attention_heads=2,
+            audio_out_channels=16,
-        audio_attention_head_dim=8,
+            audio_num_attention_heads=2,
-        audio_cross_attention_dim=16,
+            audio_attention_head_dim=8,
-        num_layers=2,
+            audio_cross_attention_dim=16,
-        caption_channels=16,
+            num_layers=2,
-    )
+            caption_channels=16,
+        )
 
     # NOTE: This is currently using the LTX2 custom enabler, but the custom
     # enablers will be consolidated after

tests/diffusion/diffusion_backend/test_diffusers_backend.py

@@ -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,6 +143,84 @@ 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",
         [

tests/diffusion/distributed/test_sp_plan_hooks.py

@@ -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:

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -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
@@ -63,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.
 

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -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):

tests/e2e/accuracy/helpers.py

@@ -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.]+)")
 

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -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,
+        ssim_threshold=thresholds.ssim,
-        psnr_threshold=_psnr_threshold(),
+        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()

tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py

@@ -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]
+pytestmark = [pytest.mark.diffusion]
-
-
-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
-
 
 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("ffmpeg")
-    _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)
     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)
+    online_metadata = probe_video(online_path)
-    offline_metadata = _probe_video(offline_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,
+        ssim_threshold=thresholds.ssim,
-        psnr_threshold=PSNR_THRESHOLD,
+        psnr_threshold=thresholds.psnr,
     )
     print(f"offline_metadata={offline_metadata_path}")

tests/e2e/accuracy/wan22_i2v/wan22_i2v_video_similarity_common.py

@@ -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度角倾泻而下，在兔脸右侧形成柔和的伦勃朗式阴影，强化了面部立体感和皮毛质感。

tests/e2e/offline_inference/test_gepard_tts.py

@@ -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

tests/e2e/offline_inference/test_minicpmo_4_5.py

@@ -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."""

tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py

@@ -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."""

tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py

@@ -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.

tests/e2e/online_serving/test_minicpmo_4_5.py

@@ -152,7 +152,7 @@ def test_text_to_text_001(omni_server, openai_client) -> 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_server", test_params, indirect=True)
 def test_text_to_audio_001(omni_server, openai_client) -> None:
     """
@@ -223,7 +223,7 @@ def test_text_to_audio_with_default_reference(omni_server, openai_client) -> Non
 
 @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_server", test_params, indirect=True)
 def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
     """
@@ -254,7 +254,7 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> 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_server", test_params, indirect=True)
 def test_image_to_text_audio_001(omni_server, openai_client) -> None:
     """
@@ -283,7 +283,7 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> 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_server", test_params, indirect=True)
 def test_video_to_text_audio_001(omni_server, openai_client) -> None:
     """

tests/e2e/online_serving/test_minicpmo_4_5_expansion.py

@@ -87,7 +87,7 @@ _TTS_EXTRA_BODY = {
 }
 
 
-@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_server", test_params, indirect=True)
 def test_text_video_to_text_001(omni_server, openai_client) -> None:
     """
@@ -114,7 +114,7 @@ def test_text_video_to_text_001(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@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_server", test_params, indirect=True)
 def test_sequential_requests_independent(omni_server, openai_client) -> None:
     """
@@ -157,7 +157,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None:
     )
 
 
-@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_server", test_params, indirect=True)
 def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
     """
@@ -185,7 +185,7 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@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_server", test_params, indirect=True)
 def test_text_to_audio_long_form_001(omni_server, openai_client) -> None:
     """
@@ -208,7 +208,7 @@ def test_text_to_audio_long_form_001(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@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_server", test_params, indirect=True)
 def test_chinese_text_to_audio(omni_server, openai_client) -> None:
     """

tests/e2e/online_serving/test_nemotron_voicechat_duplex.py

@@ -32,7 +32,7 @@ if TOKENIZER:
 pytestmark = [pytest.mark.full_model, pytest.mark.omni]
 
 
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=1)
 @pytest.mark.parametrize(
     "omni_server",
     [

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -136,7 +136,7 @@ def get_max_batch_size(size_type="few"):
     return batch_sizes.get(size_type, 5)
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
 def test_text_to_text_audio_001(omni_server, online_client) -> None:
     """
@@ -156,7 +156,7 @@ def test_text_to_text_audio_001(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_video_to_text_001(omni_server, online_client) -> None:
     """
@@ -180,7 +180,7 @@ def test_text_video_to_text_001(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params + reduced_token_params, indirect=True)
 def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
     """
@@ -205,7 +205,7 @@ def test_text_audio_to_text_audio_002(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
     """
@@ -233,7 +233,7 @@ def test_large_image_to_text_audio_001(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
     """
@@ -258,7 +258,7 @@ def test_text_video_to_text_audio_001(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config)
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_audio_in_video_001(omni_server, online_client) -> None:
     """
@@ -284,7 +284,7 @@ def test_audio_in_video_001(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config)
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_audio_in_video_002(omni_server, online_client) -> None:
     """
@@ -311,7 +311,7 @@ def test_audio_in_video_002(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_audio_in_video_default_loader_sampling_regression(omni_server, online_client) -> None:
     """
@@ -349,7 +349,7 @@ def test_audio_in_video_default_loader_sampling_regression(omni_server, online_c
     online_client.send_omni_request(dict(base_config))
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_one_word_prompt_001(omni_server, online_client) -> None:
     """Catastrophic-regression gate on one-word pronunciation.
@@ -403,7 +403,7 @@ def test_one_word_prompt_001(omni_server, online_client) -> None:
     )
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_speaker_001(omni_server, online_client) -> None:
     """
@@ -428,7 +428,7 @@ def test_speaker_001(omni_server, online_client) -> None:
     online_client.send_omni_request(request_config)
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_speaker_002(omni_server, online_client) -> None:
     """
@@ -463,7 +463,7 @@ def test_speaker_002(omni_server, online_client) -> None:
             print(f"Gender assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
 
 
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@hardware_test(res={"cuda": ["H100", "B200"], "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_audio_long_output_001(omni_server, online_client) -> None:
     """

tests/helpers/__init__.py

@@ -11,7 +11,7 @@ affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based
 ``OmniServer`` tests.
 """
 
-import pytest
+from __future__ import annotations
 
 
 def skip_if_gated_repo_inaccessible(
@@ -36,11 +36,15 @@ def skip_if_gated_repo_inaccessible(
     try:
         hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)
     except GatedRepoError as exc:
+        import pytest
+
         pytest.skip(
             f"Skipping: gated HF repo {repo_id!r} inaccessible to the current "
             f"HF_TOKEN ({exc}). See docs/contributing/ci/hf_credentials.md."
         )
     except RepositoryNotFoundError as exc:
+        import pytest
+
         pytest.skip(f"Skipping: HF repo {repo_id!r} not found ({exc}).")
     except Exception:
         return

tests/helpers/mark.py

@@ -16,6 +16,11 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     import pytest
 
+    # CUDA (and only CUDA) may list several SKUs so one item matches H100 nightly
+    # and B200 mirror collection: ``res={"cuda": ["H100", "B200"]}``.
+    # Kept under TYPE_CHECKING so Python 3.9 agents do not evaluate ``str | list``.
+    SkuSpec = str | list[str] | tuple[str, ...]
+
 # Marker description tags in ``pyproject.toml`` ``tool.pytest.ini_options.markers``.
 # Example: ``"H100: [hardware-resource] [cuda] Tests that require H100 GPU"``.
 _HARDWARE_RESOURCE_MARKER_TAG = "[hardware-resource]"
@@ -108,18 +113,25 @@ def get_level_markers() -> frozenset[str]:
 
 
 @lru_cache(maxsize=8)
-def _skus_for_platform(platform: str) -> frozenset[str]:
+def get_skus_for_platform(platform: str) -> frozenset[str]:
     """SKU names tagged ``[hardware-resource]`` and ``[platform]`` (e.g. ``[cuda]``)."""
     return get_hardware_mark_list() & _marker_names_with_tag(f"[{platform}]")
 
 
 def _require_sku(platform: str, res: str) -> None:
-    allowed = _skus_for_platform(platform)
+    allowed = get_skus_for_platform(platform)
     if res not in allowed:
         supported = ", ".join(sorted(allowed)) or f"(none tagged [hardware-resource] [{platform}])"
         raise ValueError(f"Invalid {platform} resource type: {res}. Supported: {supported}")
 
 
+def _as_sku_list(res: SkuSpec) -> list[str]:
+    skus = [res] if isinstance(res, str) else list(res)
+    if not skus:
+        raise ValueError("resource list must not be empty")
+    return skus
+
+
 def _cards_mark(num_cards: int) -> pytest.MarkDecorator:
     """Return ``cards_{n}`` so collection can filter with ``-m cards_2`` (etc.).
 
@@ -137,14 +149,15 @@ def _cards_mark(num_cards: int) -> pytest.MarkDecorator:
     return getattr(pytest.mark, f"cards_{num_cards}")
 
 
-def _cuda_marks(*, res: str, num_cards: int):
+def _cuda_marks(*, res: SkuSpec, num_cards: int):
     import pytest
     from vllm.platforms import current_platform
 
-    _require_sku("cuda", res)
+    skus = _as_sku_list(res)
-    test_resource = getattr(pytest.mark, res)
+    for sku in skus:
-    test_platform_detail = pytest.mark.cuda
+        _require_sku("cuda", sku)
-    marks = [test_resource, test_platform_detail, _cards_mark(num_cards)]
+    marks = [getattr(pytest.mark, sku) for sku in skus]
+    marks.extend([pytest.mark.cuda, _cards_mark(num_cards)])
     if num_cards == 1 or not current_platform.is_cuda():
         return marks
     return marks + [
@@ -204,10 +217,10 @@ def _res_platforms() -> frozenset[str]:
     A ``[hardware-platform]`` name is a ``res`` key only if some SKU is tagged
     with that platform (``[cuda]``, ``[npu]``, …). ``cpu`` / ``gpu`` have no SKUs.
     """
-    return frozenset(p for p in get_supported_platforms() if _skus_for_platform(p))
+    return frozenset(p for p in get_supported_platforms() if get_skus_for_platform(p))
 
 
-def _normalize_num_cards(res: dict[str, str], num_cards: int | dict[str, int]) -> dict[str, int]:
+def _normalize_num_cards(res: dict[str, SkuSpec], num_cards: int | dict[str, int]) -> dict[str, int]:
     allowed = _res_platforms()
     device_class = get_supported_platforms() - allowed
     for platform in res:
@@ -232,20 +245,24 @@ def _normalize_num_cards(res: dict[str, str], num_cards: int | dict[str, int]) -
     return num_cards_dict
 
 
-def _marks_for_platform(platform: str, resource: str, num_cards: int) -> list[pytest.MarkDecorator]:
+def _marks_for_platform(platform: str, resource: SkuSpec, num_cards: int) -> list[pytest.MarkDecorator]:
     import pytest
 
-    builders = {
+    if platform == "cuda":
-        "cuda": _cuda_marks,
+        marks = _cuda_marks(res=resource, num_cards=num_cards)
-        "rocm": _rocm_marks,
+    else:
-        "xpu": _xpu_marks,
+        if not isinstance(resource, str):
-        "musa": _musa_marks,
+            raise ValueError(f"{platform} resource must be a string, got {resource!r}")
-        "npu": _npu_marks,
+        builders = {
-    }
+            "rocm": _rocm_marks,
-    builder = builders.get(platform)
+            "xpu": _xpu_marks,
-    if builder is None:
+            "musa": _musa_marks,
-        raise ValueError(f"Unsupported platform: {platform}")
+            "npu": _npu_marks,
-    marks = builder(res=resource, num_cards=num_cards)
+        }
+        builder = builders.get(platform)
+        if builder is None:
+            raise ValueError(f"Unsupported platform: {platform}")
+        marks = builder(res=resource, num_cards=num_cards)
     if platform in _gpu_res_platforms():
         return [pytest.mark.gpu] + marks
     return marks
@@ -272,7 +289,7 @@ def _apply_marks(func, marks: list[pytest.MarkDecorator]):
     return func
 
 
-def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
+def hardware_marks(*, res: dict[str, SkuSpec], num_cards: int | dict[str, int] = 1):
     """Return marks for a **single** pytest item.
 
     Different per-platform ``num_cards`` cannot share one item: ``cards_2`` and
@@ -295,7 +312,7 @@ def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
     return all_marks
 
 
-def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
+def hardware_test(*, res: dict[str, SkuSpec], num_cards: int | dict[str, int] = 1):
     """Apply hardware marks; split collected items when card counts differ by platform.
 
     A single item cannot carry both ``cards_4`` (CUDA) and ``cards_2`` (ROCm):

tests/helpers/tests/test_mark.py

@@ -5,6 +5,8 @@ import pytest
 
 from tests.helpers.mark import (
     _gpu_res_platforms,
+    get_skus_for_platform,
+    get_supported_card_counts,
     hardware_marks,
     hardware_test,
 )
@@ -21,6 +23,18 @@ def test_gpu_platforms_come_from_pyproject_gpu_tag():
     assert "gpu" not in _gpu_res_platforms()
 
 
+def test_get_skus_for_platform_cuda_excludes_other_accelerators():
+    cuda_skus = get_skus_for_platform("cuda")
+    assert cuda_skus >= {"H100", "H200", "H800", "L4", "B200"}
+    assert "MI325" not in cuda_skus
+    assert "A2" not in cuda_skus
+    assert get_skus_for_platform("rocm") == {"MI325"}
+
+
+def test_get_supported_card_counts_comes_from_pyproject():
+    assert get_supported_card_counts() == frozenset(range(1, 9))
+
+
 def test_default_num_cards_adds_cards_1():
     names = _mark_names(hardware_marks(res={"cuda": "H100"}))
     assert "cards_1" in names
@@ -137,3 +151,14 @@ def test_cuda_registered_skus_attach_sku_and_cards(sku):
 def test_cuda_rejects_sku_from_another_platform():
     with pytest.raises(ValueError, match="Invalid cuda resource type"):
         hardware_marks(res={"cuda": "MI325"})
+
+
+def test_cuda_multi_sku_attaches_all_skus_and_one_cards_mark():
+    names = _mark_names(hardware_marks(res={"cuda": ["H100", "B200"]}, num_cards=2))
+    assert names >= {"H100", "B200", "cuda", "gpu", "cards_2"}
+    assert "cards_1" not in names
+
+
+def test_non_cuda_rejects_sku_list():
+    with pytest.raises(ValueError, match="rocm resource must be a string"):
+        hardware_marks(res={"rocm": ["MI325"]})

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -2,13 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for CosyVoice3 components."""
 
+from contextlib import contextmanager
 from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.nn as nn
 
-from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
     CausalHiFTGenerator,
 )
@@ -16,6 +18,17 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@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
+
+
 @pytest.fixture
 def causal_hift():
     return CausalHiFTGenerator(
@@ -53,8 +66,6 @@ class TestPreLookaheadLayer:
 
         return PreLookaheadLayer(in_channels=512, channels=512, pre_lookahead_len=3)
 
-    @pytest.mark.core_model
-    @pytest.mark.cpu
     def test_forward_shape(self, layer):
         """Test that output shape matches input shape."""
         batch, seq_len, channels = 2, 10, 512
@@ -64,8 +75,6 @@ class TestPreLookaheadLayer:
 
         assert out.shape == x.shape
 
-    @pytest.mark.core_model
-    @pytest.mark.cpu
     def test_forward_with_context(self, layer):
         """Test forward with context for streaming."""
         batch, seq_len, channels = 1, 10, 512
@@ -77,8 +86,6 @@ class TestPreLookaheadLayer:
 
         assert out.shape == x.shape
 
-    @pytest.mark.core_model
-    @pytest.mark.cpu
     def test_residual_connection(self, layer):
         """Test that residual connection is applied."""
         batch, seq_len, channels = 1, 5, 512
@@ -98,10 +105,9 @@ class TestDiTAttention:
     def attention(self):
         from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention
 
-        return DiTAttention(dim=512, heads=8, dim_head=64, dropout=0.0)
+        with _force_torch_sdpa():
+            return DiTAttention(dim=512, heads=8, dim_head=64, dropout=0.0)
 
-    @pytest.mark.core_model
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, attention):
         """Test attention output shape."""
         batch, seq_len, dim = 2, 16, 512
@@ -111,8 +117,6 @@ class TestDiTAttention:
 
         assert out.shape == x.shape
 
-    @pytest.mark.core_model
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_with_mask(self, attention):
         """Test attention with mask."""
         batch, seq_len, dim = 2, 16, 512
@@ -126,8 +130,6 @@ class TestDiTAttention:
         # Masked positions should be zero
         assert torch.allclose(out[:, -3:], torch.zeros_like(out[:, -3:]))
 
-    @pytest.mark.core_model
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_qkv_projections(self, attention):
         """Test that Q/K/V projections exist and have correct dimensions."""
         assert hasattr(attention, "to_q")
@@ -145,10 +147,9 @@ class TestDiTBlock:
     def block(self):
         from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTBlock
 
-        return DiTBlock(dim=512, heads=8, dim_head=64, ff_mult=4, dropout=0.0)
+        with _force_torch_sdpa():
+            return DiTBlock(dim=512, heads=8, dim_head=64, ff_mult=4, dropout=0.0)
 
-    @pytest.mark.core_model
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, block):
         """Test block output shape."""
         batch, seq_len, dim = 2, 16, 512
@@ -159,8 +160,6 @@ class TestDiTBlock:
 
         assert out.shape == x.shape
 
-    @pytest.mark.core_model
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_adalayernorm_modulation(self, block):
         """Test that AdaLayerNorm modulates based on timestep."""
         batch, seq_len, dim = 1, 8, 512
@@ -182,21 +181,20 @@ class TestDiT:
     def dit(self):
         from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT
 
-        return DiT(
+        with _force_torch_sdpa():
-            dim=256,
+            return DiT(
-            depth=2,
+                dim=256,
-            heads=4,
+                depth=2,
-            dim_head=64,
+                heads=4,
-            dropout=0.0,
+                dim_head=64,
-            ff_mult=2,
+                dropout=0.0,
-            mel_dim=80,
+                ff_mult=2,
-            mu_dim=80,
+                mel_dim=80,
-            spk_dim=80,
+                mu_dim=80,
-            long_skip_connection=True,
+                spk_dim=80,
-        )
+                long_skip_connection=True,
+            )
 
-    @pytest.mark.core_model
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_forward_shape(self, dit):
         """Test DiT forward output shape."""
         batch, mel_dim, seq_len = 1, 80, 32
@@ -211,8 +209,6 @@ class TestDiT:
 
         assert out.shape == (batch, mel_dim, seq_len)
 
-    @pytest.mark.core_model
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_timestep_embedding(self, dit):
         """Test that different timesteps produce different outputs."""
         batch, mel_dim, seq_len = 1, 80, 16
@@ -245,8 +241,6 @@ class TestCFM:
 
         return DummyEstimator()
 
-    @pytest.mark.core_model
-    @pytest.mark.cpu
     def test_causal_conditional_cfm_forward(self, dummy_estimator):
         """Test CausalConditionalCFM forward pass."""
         from omegaconf import DictConfig
@@ -432,18 +426,17 @@ class TestCFM:
 class TestSDPAFallback:
     """Test SDPA fallback for float32 inputs."""
 
-    @pytest.mark.core_model
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_float32_uses_sdpa(self):
         """Test that float32 inputs use SDPA fallback."""
         from vllm_omni.diffusion.attention.layer import Attention
 
-        attn = Attention(
+        with _force_torch_sdpa():
-            num_heads=8,
+            attn = Attention(
-            head_size=64,
+                num_heads=8,
-            causal=False,
+                head_size=64,
-            softmax_scale=1.0 / 8.0,
+                causal=False,
-        )
+                softmax_scale=1.0 / 8.0,
+            )
 
         batch, seq_len, heads, dim = 1, 16, 8, 64
         q = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32)

tools/nightly/buildkite_testcase_statistics.py

@@ -1,4 +1,7 @@
 #!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Parse pytest commands from Buildkite .buildkite/cuda/test-{ready,merge,nightly,weekly}.yml;
 collect test cases (including parametrized) via pytest --collect-only -q and produce an HTML report.

tools/nightly/run_nightly_jobs.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # Extract steps from .buildkite/cuda/test-nightly.yml that contain pytest, synthesize
 # small bash wrappers (exports + pytest), run them, and tee output to logs named
@@ -259,7 +259,7 @@ _nightly_log_dir_basename() {
 
 BUILDKITE_REL=".buildkite/cuda/test-nightly.yml"
 
-_find_repo_containing_nightly() {
+_find_repo_containing_yml() {
   local dir="${1:-}"
   [[ -n "$dir" ]] || return 1
   dir="$(cd "$dir" && pwd)" || return 1
@@ -307,10 +307,10 @@ else
     REPO_ROOT="$(git -C "${PWD}" rev-parse --show-toplevel 2>/dev/null || true)"
   fi
   if [[ -z "${REPO_ROOT}" ]]; then
-    REPO_ROOT="$(_find_repo_containing_nightly "${PWD}" || true)"
+    REPO_ROOT="$(_find_repo_containing_yml "${PWD}" || true)"
   fi
   if [[ -z "${REPO_ROOT}" ]]; then
-    REPO_ROOT="$(_find_repo_containing_nightly "${SCRIPT_DIR}" || true)"
+    REPO_ROOT="$(_find_repo_containing_yml "${SCRIPT_DIR}" || true)"
   fi
   if [[ -z "${REPO_ROOT}" ]]; then
     echo "Could not locate ${BUILDKITE_REL}. Set REPO_ROOT or YML, run from inside the vllm-omni clone," >&2
@@ -326,7 +326,7 @@ fi
 
 echo "Log directory: ${LOG_DIR}" >&2
 
-# Require nightly YAML only when at least one non-stability test kind needs it (validated in Python too).
+# Require test-nightly.yml only when at least one non-stability test kind needs it (validated in Python too).
 _needs_nightly_yml=0
 IFS=',' read -ra _TTOK <<< "${TEST_TYPE}"
 for _t in "${_TTOK[@]}"; do
@@ -971,7 +971,6 @@ PERF_RESULTS_SRC="${REPO_ROOT}/tests/dfx/perf/results"
 # Copy perf benchmark JSON written under tests/dfx/perf/results/ during this run into LOG_DIR.
 _collect_perf_result_jsons() {
   local since_epoch="${1:?}"
-  local dest="${LOG_DIR}/perf_results"
 
   if [[ ! -d "${PERF_RESULTS_SRC}" ]]; then
     return 0

tools/pre_commit/shellcheck.sh

@@ -62,7 +62,9 @@ run_shellcheck() {
     local f
     for f in "$@"; do
         if should_lint "$f"; then
-            "$SHELLCHECK_BIN" -s bash "$f"
+            # -x honors `# shellcheck source=` so sourced helpers are followed
+            # (avoids SC1091 / false SC2034 on BUILDKITE_REL in entry scripts).
+            "$SHELLCHECK_BIN" -s bash -x "$f"
         fi
     done
 }
@@ -75,5 +77,5 @@ fi
 # Direct invocation with no args: lint every tracked *.sh.
 while IFS= read -r -d '' f || [ -n "$f" ]; do
     git check-ignore -q "$f" && continue
-    "$SHELLCHECK_BIN" -s bash "$f"
+    "$SHELLCHECK_BIN" -s bash -x "$f"
 done < <(find . -path ./.git -prune -o -name "*.sh" -print0)

tools/run_jobs_common.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # Shared helpers for tools/run_*_jobs.sh (source this file; do not execute).
 #
@@ -15,7 +15,7 @@
 # metadata.
 
 if [[ -n "${_RUN_JOBS_COMMON_LOADED:-}" ]]; then
-  return 0 2>/dev/null || exit 0
+  return 0
 fi
 _RUN_JOBS_COMMON_LOADED=1
 

tools/run_merge_jobs.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # Extract steps from .buildkite/cuda/test-merge.yml that contain pytest, synthesize
 # small bash wrappers (exports + pytest), run them, and tee output to logs named

tools/run_ready_jobs.sh

@@ -1,6 +1,6 @@
 #!/usr/bin/env bash
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # Extract steps from .buildkite/cuda/test-ready.yml that contain pytest, synthesize
 # small bash wrappers (exports + pytest), run them, and tee output to logs named

vllm_omni/diffusion/attention/backends/abstract.py

@@ -70,7 +70,15 @@ class AttentionBackend(ABC):
         return None
 
     @classmethod
-    def supports_attention_mask(cls) -> bool:
+    def supports_attention_mask(cls, attention_spec: object | None = None) -> bool:
+        """Return whether this backend can consume a nontrivial ``attn_mask``.
+
+        ``attention_spec`` is the resolved per-role config when the user picked
+        a backend explicitly. Implementations that depend on kernel variant
+        (for example FlashInfer cute-dsl vs fa2) must consult it so capability
+        probes match a runnable configuration.
+        """
+        del attention_spec
         return False
 
     @staticmethod

vllm_omni/diffusion/attention/backends/cudnn_attn.py

@@ -1,9 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import torch
 from torch.nn.attention import SDPBackend, sdpa_kernel
-from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionBackend,
@@ -12,23 +11,35 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 )
 from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
 
-logger = init_logger(__name__)
-
 
 class CuDNNAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supports_prefix_kv_slicing: bool = True
 
+    # cuDNN 9.5+ FMHA on Blackwell: head_dim divisible by 8 and at most 256
+    # for BF16/FP16. Used by automatic platform selection; explicit CUDNN_ATTN
+    # raises when the configured head size is outside this set.
+    _MAX_HEAD_SIZE = 256
+    _HEAD_SIZE_MULTIPLE = 8
+
     @classmethod
-    def supports_attention_mask(cls) -> bool:
+    def supports_attention_mask(cls, attention_spec: object | None = None) -> bool:
+        del attention_spec
         return True
 
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
-        # cuDNN 9.5+ FMHA on Blackwell supports any head_dim divisible by 8
+        return list(
-        # up to 256 for BF16/FP16. Empty list = "accept any"; the sdpa_kernel
+            range(
-        # fallback chain handles shapes cuDNN can't take.
+                CuDNNAttentionBackend._HEAD_SIZE_MULTIPLE,
-        return []
+                CuDNNAttentionBackend._MAX_HEAD_SIZE + 1,
+                CuDNNAttentionBackend._HEAD_SIZE_MULTIPLE,
+            )
+        )
+
+    @classmethod
+    def supports_head_size(cls, head_size: int) -> bool:
+        return 0 < head_size <= cls._MAX_HEAD_SIZE and head_size % cls._HEAD_SIZE_MULTIPLE == 0
 
     @staticmethod
     def get_name() -> str:
@@ -52,6 +63,9 @@ class CuDNNAttentionImpl(AttentionImpl):
     ) -> None:
         self.causal = causal
         self.softmax_scale = softmax_scale
+        # Set when AttentionConfig (or --diffusion-attention-backend) selected
+        # CUDNN_ATTN. The kv_seq_len=1 MATH path is automatic-only.
+        self.backend_explicit = bool(extra_impl_args.get("backend_explicit", False))
 
     def forward_cuda(
         self,
@@ -84,54 +98,32 @@ class CuDNNAttentionImpl(AttentionImpl):
                 )
 
         enable_gqa = query.shape[2] != key.shape[2]
+        kv_seq_len = key.shape[1]
         query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
-        # cuDNN has no kernel for degenerate attention shapes when either
-        # sequence has one token. Route those degenerate shapes before entering the
-        # cuDNN-only context: under torch.compile, FakeTensor evaluation raises
-        # before the eager RuntimeError fallback below can run.
-        if query.shape[-2] == 1 or key.shape[-2] == 1:
-            with sdpa_kernel([SDPBackend.MATH]):
-                output = torch.nn.functional.scaled_dot_product_attention(
-                    query,
-                    key,
-                    value,
-                    attn_mask=attention_mask,
-                    dropout_p=0.0,
-                    is_causal=self.causal,
-                    scale=self.softmax_scale,
-                    enable_gqa=enable_gqa,
-                )
-            return output.permute(0, 2, 1, 3)
 
-        # Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a
+        # Pin one backend only. A priority list like [CUDNN, FLASH, MATH] hits a
         # PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask,
         # cuDNN gets runtime-disabled in the same call and the dispatcher falls
         # through to MATH even though cuDNN alone handles the mask fine
         # (~11 ms vs ~215 ms for MATH on sm_120 HV-1.5 shapes).
+        # Explicit CUDNN_ATTN must not silently substitute MATH/EFFICIENT.
         #
-        # Fall back to the default SDPA dispatcher if cuDNN rejects the shape,
+        # Automatic-only: cuDNN FMHA rejects KV sequence length 1
-        # e.g. under torch.compile where Dynamo sees a symbolic head_dim and
+        # ("cudnn SDPA does not support key/value sequence length 1"). Dummy
-        # cuDNN's kernel selection fails for some symbolic attention shapes.
+        # warmup and some I2V layers (e.g. LTX-2) hit this when CUDNN_ATTN is
-        # The unpinned dispatcher then picks EFFICIENT/MATH instead of raising.
+        # the platform default. MATH is the only kernel that can run that shape.
-        try:
+        # An explicit CUDNN_ATTN request raises instead of silently switching.
-            with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
+        if kv_seq_len <= 1:
-                output = torch.nn.functional.scaled_dot_product_attention(
+            if self.backend_explicit:
-                    query,
+                raise ValueError(
-                    key,
+                    "CUDNN_ATTN was explicitly selected but cuDNN FMHA does not "
-                    value,
+                    f"support key/value sequence length {kv_seq_len}. "
-                    attn_mask=attention_mask,
+                    "Select TORCH_SDPA or another backend for this shape."
-                    dropout_p=0.0,
-                    is_causal=self.causal,
-                    scale=self.softmax_scale,
-                    enable_gqa=enable_gqa,
                 )
-        except RuntimeError as e:
+            backends = [SDPBackend.MATH]
-            if "No available kernel" not in str(e):
+        else:
-                raise
+            backends = [SDPBackend.CUDNN_ATTENTION]
-            logger.warning_once(
+        with sdpa_kernel(backends):
-                "cuDNN SDPA rejected this shape; falling back to default SDPA dispatcher. "
-                "Set DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA for the full dispatcher path on every call."
-            )
             output = torch.nn.functional.scaled_dot_product_attention(
                 query,
                 key,

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -48,7 +48,8 @@ class FlashAttentionBackend(AttentionBackend):
         return current_omni_platform.is_cuda() or current_omni_platform.is_rocm() or current_omni_platform.is_musa()
 
     @classmethod
-    def supports_attention_mask(cls) -> bool:
+    def supports_attention_mask(cls, attention_spec: object | None = None) -> bool:
+        del attention_spec
         return True
 
     @staticmethod
@@ -64,7 +65,7 @@ class FlashAttentionBackend(AttentionBackend):
         return FlashAttentionImpl
 
 
-class FlashAttentionImpl(AttentionImpl):
+class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
     # Per-platform FP8 KV quantization support.
     # To enable FP8 on a new platform, add its OmniPlatformEnum value here
     # and handle kv_cache_dtype in the corresponding forward_{platform}().
@@ -157,6 +158,8 @@ class FlashAttentionImpl(AttentionImpl):
             flash_attn_varlen_func,
         )
 
+        if flash_attn_varlen_func is None:
+            raise ImportError("Masked variable-length attention requires flash_attn_varlen_func")
         assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
         batch_size, query_length = query.shape[:2]
         if not self.is_cross_attn and query_length == key.size(1):
@@ -250,6 +253,8 @@ class FlashAttentionImpl(AttentionImpl):
             flash_attn_varlen_func,
         )
 
+        if flash_attn_varlen_func is None:
+            raise ImportError("Dense variable-length attention requires flash_attn_varlen_func")
         batch_size, q_len = query.size()[:2]
         k_len = key.size(1)
         cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device)
@@ -378,7 +383,7 @@ class FlashAttentionImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         """CUDA/ROCm/MUSA flash attention implementation."""
         from vllm_omni.diffusion.attention.backends.utils.fa import (
@@ -422,7 +427,7 @@ class FlashAttentionImpl(AttentionImpl):
                 full_attn_spans,
                 self.softmax_scale,
                 attn_func,
-                query_ranges=attn_metadata.query_ranges,
+                query_ranges=None if attn_metadata is None else attn_metadata.query_ranges,
             )
 
         packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
@@ -473,7 +478,7 @@ class FlashAttentionImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         """XPU flash attention implementation."""
         from vllm_omni.diffusion.attention.backends.utils.fa import (
@@ -508,7 +513,7 @@ class FlashAttentionImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         """NPU attention implementation using mindiesd."""
 
@@ -522,7 +527,7 @@ class FlashAttentionImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         from vllm_omni.platforms.npu.quant.kv_quant_npu import fp8_rotate_quant_fa
 
@@ -542,7 +547,7 @@ class FlashAttentionImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         from mindiesd import attention_forward
 
@@ -712,8 +717,12 @@ class FlashAttentionImpl(AttentionImpl):
         key = key[:, :used_k]
         value = value[:, :used_k]
 
-        input_scale = extra.get("laser_input_scale")
+        input_scale_raw = extra.get("laser_input_scale")
-        preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1
+        if isinstance(input_scale_raw, (int, float)) and input_scale_raw > 1:
+            input_scale = float(input_scale_raw)
+        else:
+            input_scale = 1.0
+        preserve_input_range = input_scale > 1
         if preserve_input_range:
             query = query / input_scale
             key = key / input_scale

vllm_omni/diffusion/attention/backends/flash_attn_hub.py

@@ -1,8 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import threading
+from collections.abc import Callable
 from functools import partial
+from typing import Any
 
 import torch
 from vllm.logger import init_logger
@@ -44,6 +46,12 @@ def _get_hub_module(repo_id: str):
         return _hub_modules[repo_id]
 
 
+def _require_varlen_func(func: Callable[..., Any] | None) -> Callable[..., Any]:
+    if func is None:
+        raise RuntimeError("flash_attn_varlen_func is not available from the Hub kernel")
+    return func
+
+
 def _run_varlen_dense(
     varlen_func,
     query: torch.Tensor,
@@ -79,7 +87,8 @@ class FlashAttentionHubBackend(AttentionBackend):
     supports_piecewise_spans: bool = True
 
     @classmethod
-    def supports_attention_mask(cls) -> bool:
+    def supports_attention_mask(cls, attention_spec: object | None = None) -> bool:
+        del attention_spec
         return True
 
     @staticmethod
@@ -95,7 +104,7 @@ class FlashAttentionHubBackend(AttentionBackend):
         return FlashAttentionHubImpl
 
 
-class FlashAttentionHubImpl(AttentionImpl):
+class FlashAttentionHubImpl(AttentionImpl[AttentionMetadata]):
     def __init__(
         self,
         num_heads: int,
@@ -116,8 +125,8 @@ class FlashAttentionHubImpl(AttentionImpl):
             logger.warning("FlashAttentionHubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
         hub_module = _get_hub_module("kernels-community/flash-attn2")
-        self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
+        self.flash_attn_func: Callable[..., Any] | None = getattr(hub_module, "flash_attn_func", None)
-        self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
+        self.flash_attn_varlen_func: Callable[..., Any] | None = getattr(hub_module, "flash_attn_varlen_func", None)
         if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
             raise RuntimeError("Failed to load flash-attn2 kernel from HuggingFace Hub: no functions found")
 
@@ -130,7 +139,7 @@ class FlashAttentionHubImpl(AttentionImpl):
         if attn_func is not None:
             return self._unwrap_flash_output(attn_func(q, k, v, **kwargs))
         return _run_varlen_dense(
-            self.flash_attn_varlen_func,
+            _require_varlen_func(self.flash_attn_varlen_func),
             q,
             k,
             v,
@@ -157,7 +166,7 @@ class FlashAttentionHubImpl(AttentionImpl):
             query, key, value, attention_mask, query_length, _unpad_input
         )
 
-        out_unpad = self.flash_attn_varlen_func(
+        out_unpad = _require_varlen_func(self.flash_attn_varlen_func)(
             q,
             k,
             v,
@@ -180,7 +189,7 @@ class FlashAttentionHubImpl(AttentionImpl):
         value: torch.Tensor,
     ) -> torch.Tensor:
         return _run_varlen_dense(
-            self.flash_attn_varlen_func,
+            _require_varlen_func(self.flash_attn_varlen_func),
             query,
             key,
             value,
@@ -193,7 +202,7 @@ class FlashAttentionHubImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
@@ -213,7 +222,7 @@ class FlashAttentionHubImpl(AttentionImpl):
                 full_attn_spans,
                 self.softmax_scale,
                 attn_func,
-                query_ranges=attn_metadata.query_ranges,
+                query_ranges=None if attn_metadata is None else attn_metadata.query_ranges,
             )
 
         if attention_mask is not None and torch.any(~attention_mask):
@@ -246,7 +255,8 @@ class FlashAttention3HubBackend(AttentionBackend):
     supports_piecewise_spans: bool = True
 
     @classmethod
-    def supports_attention_mask(cls) -> bool:
+    def supports_attention_mask(cls, attention_spec: object | None = None) -> bool:
+        del attention_spec
         return True
 
     @staticmethod
@@ -262,7 +272,7 @@ class FlashAttention3HubBackend(AttentionBackend):
         return FlashAttention3HubImpl
 
 
-class FlashAttention3HubImpl(AttentionImpl):
+class FlashAttention3HubImpl(AttentionImpl[AttentionMetadata]):
     def __init__(
         self,
         num_heads: int,
@@ -283,8 +293,8 @@ class FlashAttention3HubImpl(AttentionImpl):
             logger.warning("FlashAttention3HubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
         hub_module = _get_hub_module("kernels-community/flash-attn3")
-        self.flash_attn_func = getattr(hub_module, "flash_attn_func", None)
+        self.flash_attn_func: Callable[..., Any] | None = getattr(hub_module, "flash_attn_func", None)
-        self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None)
+        self.flash_attn_varlen_func: Callable[..., Any] | None = getattr(hub_module, "flash_attn_varlen_func", None)
         if self.flash_attn_func is None and self.flash_attn_varlen_func is None:
             raise RuntimeError("Failed to load flash-attn3 kernel from HuggingFace Hub: no functions found")
 
@@ -297,7 +307,7 @@ class FlashAttention3HubImpl(AttentionImpl):
         if attn_func is not None:
             return self._unwrap_flash_output(attn_func(q, k, v, **kwargs))
         return _run_varlen_dense(
-            self.flash_attn_varlen_func,
+            _require_varlen_func(self.flash_attn_varlen_func),
             q,
             k,
             v,
@@ -324,7 +334,7 @@ class FlashAttention3HubImpl(AttentionImpl):
             query, key, value, attention_mask, query_length, _unpad_input
         )
 
-        out_unpad = self.flash_attn_varlen_func(
+        out_unpad = _require_varlen_func(self.flash_attn_varlen_func)(
             q,
             k,
             v,
@@ -347,7 +357,7 @@ class FlashAttention3HubImpl(AttentionImpl):
         value: torch.Tensor,
     ) -> torch.Tensor:
         return _run_varlen_dense(
-            self.flash_attn_varlen_func,
+            _require_varlen_func(self.flash_attn_varlen_func),
             query,
             key,
             value,
@@ -360,7 +370,7 @@ class FlashAttention3HubImpl(AttentionImpl):
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
+        attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
@@ -380,7 +390,7 @@ class FlashAttention3HubImpl(AttentionImpl):
                 full_attn_spans,
                 self.softmax_scale,
                 attn_func,
-                query_ranges=attn_metadata.query_ranges,
+                query_ranges=None if attn_metadata is None else attn_metadata.query_ranges,
             )
 
         if attention_mask is not None and torch.any(~attention_mask):

vllm_omni/diffusion/attention/backends/flashinfer_attn.py

@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import torch
 from packaging.version import InvalidVersion, Version
@@ -13,6 +16,9 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionMetadata,
 )
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+
 logger = init_logger(__name__)
 
 try:
@@ -32,9 +38,18 @@ class FlashInferAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @classmethod
-    def supports_attention_mask(cls) -> bool:
+    def supports_attention_mask(cls, attention_spec: object | None = None) -> bool:
-        # Non-CuTe batch-prefill backends accept boolean custom masks. CuTe-DSL
+        # fa2/fa3 batch-prefill accept boolean custom masks. cute-dsl does not:
-        # falls back to SDPA for nontrivial masks.
+        # automatic platform selection may fall back to SDPA, but an explicit
+        # FLASHINFER_ATTN config must not be advertised as mask-capable.
+        requested = "auto"
+        quant = getattr(attention_spec, "quant", None) if attention_spec is not None else None
+        requested_backend = getattr(quant, "flashinfer_backend", None)
+        if requested_backend:
+            requested = requested_backend
+        backend = FlashInferAttentionImpl._select_backend(requested)
+        if backend == "cute-dsl":
+            return attention_spec is None
         return True
 
     @staticmethod
@@ -49,7 +64,7 @@ class FlashInferAttentionBackend(AttentionBackend):
         return "FLASHINFER_ATTN"
 
     @staticmethod
-    def get_impl_cls() -> type["FlashInferAttentionImpl"]:
+    def get_impl_cls() -> type[FlashInferAttentionImpl]:
         return FlashInferAttentionImpl
 
 
@@ -94,13 +109,17 @@ class FlashInferAttentionImpl(AttentionImpl):
         self.dtype_qk = self._check_dtype(quant.get("dtype_qk"), "dtype_qk", self._QK_DTYPES)
         self.dtype_vo = self._check_dtype(quant.get("dtype_vo"), "dtype_vo", self._VO_DTYPES)
         requested_backend = quant.get("flashinfer_backend", "auto")
+        # Set when AttentionConfig (or --diffusion-attention-backend) selected
+        # FLASHINFER_ATTN. Cute-dsl custom-mask SDPA is automatic-only.
+        self.backend_explicit = bool(extra_impl_args.get("backend_explicit", False))
+        self._sdpa_fallback: SDPAImpl | None = None
 
         if not HAS_FLASHINFER:
             raise ImportError("FLASHINFER_ATTN backend requires flashinfer")
 
         self._check_flashinfer_version()
 
-        self.flashinfer_backend = self._select_backend(requested_backend, self.device)
+        self.flashinfer_backend = self._select_backend(requested_backend, device=self.device)
         workspace_size = 0 if self.flashinfer_backend == "cute-dsl" else 128 * 1024 * 1024
         self._workspace = torch.empty(
             workspace_size,
@@ -144,10 +163,15 @@ class FlashInferAttentionImpl(AttentionImpl):
             )
 
     @staticmethod
-    def _select_backend(requested_backend: str, device: torch.device) -> str:
+    def _select_backend(requested_backend: str, device: torch.device | None = None) -> str:
         if requested_backend != "auto":
             return requested_backend
-        major, _minor = torch.cuda.get_device_capability(device)
+        try:
+            major, _minor = (
+                torch.cuda.get_device_capability(device) if device is not None else torch.cuda.get_device_capability()
+            )
+        except Exception:
+            return "fa2"
         if major >= 10:
             return "cute-dsl"
         if major >= 9:
@@ -180,18 +204,14 @@ class FlashInferAttentionImpl(AttentionImpl):
         Returns either ``(qo_len, kv_len)`` (shared across the batch) or
         ``(batch_size, qo_len, kv_len)`` (per-sample), or ``None`` when the
         mask is all-keep (elide). Only boolean masks are handled here;
-        additive/float masks raise ``ValueError`` so the caller falls back to
+        additive/float masks and shape mismatches raise ``ValueError`` because
-        SDPA, which applies them with the correct softmax semantics. Shape
+        an explicitly selected backend must not silently switch to SDPA.
-        mismatches also raise ``ValueError``.
         """
         mask = attn_mask
         if mask.dtype != torch.bool:
             # Additive masks (0 / -inf / -1e4 / finfo.min) cannot be faithfully
             # reduced to a boolean keep-mask here; SDPA handles them correctly.
-            raise ValueError(
+            raise ValueError(f"non-boolean attn_mask (dtype={mask.dtype}); FlashInfer custom_mask is boolean-only")
-                f"non-boolean attn_mask (dtype={mask.dtype}); FlashInfer custom_mask "
-                "is boolean-only — deferring to SDPA"
-            )
         # Diffusion masks arrive as (qo,kv), (1,1,kv), (B,1,1,kv), (B,1,qo,kv)
         # or (B,H,qo,kv). The mask is identical across heads, so collapse the
         # head dim, but keep a real batch dim — indexing mask[0] would reuse
@@ -216,23 +236,6 @@ class FlashInferAttentionImpl(AttentionImpl):
         # kernel, which reads from GPU memory directly.
         return mask.contiguous()
 
-    def _sdpa_fallback(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata | None,
-    ) -> torch.Tensor:
-        from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
-
-        impl = SDPAImpl(
-            num_heads=query.shape[2],
-            head_size=query.shape[3],
-            softmax_scale=self.softmax_scale,
-            causal=self.causal,
-        )
-        return impl.forward_cuda(query, key, value, attn_metadata)
-
     @torch.compiler.disable
     def _plan_wrapper(
         self,
@@ -353,6 +356,28 @@ class FlashInferAttentionImpl(AttentionImpl):
     ) -> torch.Tensor:
         return self._wrapper.run(query, key, value)
 
+    def _sdpa_for_unsupported_mask(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+        reason: str,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+
+        logger.warning_once("FLASHINFER_ATTN falling back to SDPA: %s", reason)
+        fallback = self._sdpa_fallback
+        if fallback is None:
+            fallback = SDPAImpl(
+                num_heads=query.shape[2],
+                head_size=query.shape[3],
+                softmax_scale=self.softmax_scale,
+                causal=self.causal,
+            )
+            self._sdpa_fallback = fallback
+        return fallback.forward_cuda(query, key, value, attn_metadata)
+
     def forward_cuda(
         self,
         query: torch.Tensor,
@@ -360,6 +385,15 @@ class FlashInferAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
+        if not HAS_FLASHINFER:
+            raise ImportError(
+                "FLASHINFER_ATTN backend requires flashinfer. "
+                "Install it or set DIFFUSION_ATTENTION_BACKEND to another backend."
+            )
+
+        # Explicit FLASHINFER_ATTN must not silently switch away from the
+        # requested kernel. Automatic platform selection (Blackwell cute-dsl)
+        # may use SDPA for masks that this FlashInfer variant cannot run.
         batch_size = query.shape[0]
         custom_mask = None
         if attn_metadata is not None and attn_metadata.attn_mask is not None:
@@ -370,14 +404,44 @@ class FlashInferAttentionImpl(AttentionImpl):
                     qo_len=query.shape[1],
                     kv_len=key.shape[1],
                 )
-            except ValueError as error:
+            except ValueError:
-                logger.debug("Falling back to SDPA for mask path: %s", error)
+                if self.backend_explicit:
-                return self._sdpa_fallback(query, key, value, attn_metadata)
+                    raise
+                return self._sdpa_for_unsupported_mask(
+                    query,
+                    key,
+                    value,
+                    attn_metadata,
+                    "attn_mask is not a FlashInfer boolean custom_mask",
+                )
+            # FlashInfer cannot combine causal masking with a custom_mask; rather
+            # than silently dropping the explicit mask (diverging from SDPA),
+            # require an explicit caller to select a compatible backend.
             if custom_mask is not None and self.causal:
-                return self._sdpa_fallback(query, key, value, attn_metadata)
+                if self.backend_explicit:
+                    raise ValueError(
+                        "FLASHINFER_ATTN does not support causal=True together with an explicit custom mask."
+                    )
+                return self._sdpa_for_unsupported_mask(
+                    query,
+                    key,
+                    value,
+                    attn_metadata,
+                    "causal=True with a custom mask",
+                )
 
         if custom_mask is not None and self.flashinfer_backend == "cute-dsl":
-            logger.debug("CuTe DSL does not support custom masks; deferring to SDPA")
+            if self.backend_explicit:
-            return self._sdpa_fallback(query, key, value, attn_metadata)
+                raise ValueError(
+                    "FLASHINFER_ATTN cute-dsl backend does not support custom masks. "
+                    "Select TORCH_SDPA or a FlashInfer backend that accepts custom_mask."
+                )
+            return self._sdpa_for_unsupported_mask(
+                query,
+                key,
+                value,
+                attn_metadata,
+                "cute-dsl does not support custom masks",
+            )
 
         return self._run_batch_prefill(query, key, value, custom_mask)

vllm_omni/diffusion/attention/backends/ring/ring_globals.py

@@ -22,6 +22,11 @@ except (ImportError, ModuleNotFoundError):
 HAS_FA3 = False
 fa3_fwd_func = None  # Low-level forward function (_flash_attn_forward)
 fa3_attn_func = None  # High-level attention function (flash_attn_func)
+# ``None`` means that the extension does not publish an architecture contract,
+# so retain the historical import-only behavior.  The pinned fa3-fwd wheel is
+# handled below because importing its Python module succeeds on Blackwell even
+# though its CUDA binary only contains SM8x/SM90 kernels.
+FA3_SUPPORTED_CUDA_MAJORS: frozenset[int] | None = None
 
 # FA4 detection. The CuTe API returns LSE directly, so it can participate in
 # Ring Attention's numerically stable block-wise output accumulation.
@@ -42,16 +47,36 @@ try:
     from flash_attn_interface import flash_attn_func as fa3_attn_func  # noqa: F401
 
     HAS_FA3 = True
+    # The source-build FA3 interface is the Hopper implementation. Importing
+    # its Python module on Blackwell does not imply that an SM10x kernel exists.
+    FA3_SUPPORTED_CUDA_MAJORS = frozenset({9})
 except (ImportError, ModuleNotFoundError):
     pass
 
-# Fallback: try fa3_fwd_interface (PyPI package, supports Ampere/Ada/Hopper)
+# Secondary FA3 import path: fa3_fwd_interface PyPI package.
 if not HAS_FA3:
     try:
+        from importlib.metadata import PackageNotFoundError, version
+
+        import fa3_fwd_interface
         from fa3_fwd_interface import _flash_attn_forward as fa3_fwd_func  # noqa: F401
         from fa3_fwd_interface import flash_attn_func as fa3_attn_func  # noqa: F401
 
         HAS_FA3 = True
+        published_majors = getattr(fa3_fwd_interface, "SUPPORTED_CUDA_MAJORS", None)
+        if published_majors is not None:
+            FA3_SUPPORTED_CUDA_MAJORS = frozenset(int(major) for major in published_majors)
+        else:
+            try:
+                fa3_fwd_version = version("fa3_fwd")
+            except PackageNotFoundError:
+                fa3_fwd_version = None
+            # requirements/cuda.txt pins this wheel.  Its fat binary contains
+            # SM80 and SM90 cubins, but no Blackwell kernel.  Do not apply this
+            # restriction to an unknown/future build unless it publishes its
+            # own supported-major metadata above.
+            if fa3_fwd_version == "0.0.3":
+                FA3_SUPPORTED_CUDA_MAJORS = frozenset({8, 9})
     except (ImportError, ModuleNotFoundError):
         pass
 
@@ -68,9 +93,8 @@ try:
     HAS_FLASHINFER = True
 except Exception as e:
     # flashinfer may raise RuntimeError at import-time for version/binary mismatches.
-    # Treat it as unavailable so the runtime can gracefully fall back to other backends.
     HAS_FLASHINFER = False
-    logger.warning("FlashInfer is unavailable; falling back to other attention backends. Reason: %s", e)
+    logger.warning("FlashInfer ring kernels are unavailable. Reason: %s", e)
 
 try:
     import aiter  # noqa: F401

vllm_omni/diffusion/attention/backends/ring/ring_kernels.py

@@ -59,10 +59,8 @@ def pytorch_attn_forward(
     k shape (bs, seqlen, nhead, hs)
     v shape (bs, seqlen, nhead, hs)
     """
-    # Fallback logic: Flash Attention does not support float32.
-    # If op_type is 'flash' but dtype is float32, force 'efficient'.
     if op_type == "flash" and q.dtype == torch.float32:
-        op_type = "efficient"
+        raise ValueError("The explicitly requested Flash ring kernel does not support float32.")
 
     # Volta (V100, sm70) does not have native bf16 support for the efficient
     # SDPA kernel. When running Ring attention in bf16 on such GPUs, the
@@ -123,6 +121,7 @@ def pytorch_attn_forward(
     # Keep LSE in fp32 for numerical stability when accumulating across ring
     # steps (update_out_and_lse uses sigmoid/logsigmoid on LSE diffs).
     # Casting LSE down to fp16/bf16 can introduce NaNs on some GPUs/shapes.
+    # PyTorch SDPA returns LSE as (B, H, S).
     lse = lse.to(torch.float32)
 
     if out.dtype != orig_dtype:
@@ -173,6 +172,7 @@ def flash_attn_forward(
             alibi_slopes=alibi_slopes,
             return_softmax=return_softmax,
         )
+    # FA2 softmax_lse is (B, H, S), optionally padded on S.
     return block_out, block_lse
 
 
@@ -198,6 +198,7 @@ def fa3_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, softcap,
         softcap=softcap if softcap else 0.0,
     )
 
+    # FA3 softmax_lse is (B, H, S).
     return out, softmax_lse
 
 
@@ -236,6 +237,7 @@ def flash_attn4_func_forward(
         softcap=softcap or 0.0,
         return_lse=True,
     )
+    # FA4 CuTe softmax_lse is (B, H, S).
     return out, softmax_lse
 
 
@@ -264,6 +266,7 @@ def flash_attn_forward_aiter(
         return_lse=True,
     )
 
+    # Aiter softmax_lse is (B, H, S), matching FA2.
     return block_out, block_lse
 
 
@@ -293,7 +296,7 @@ def flashinfer_attn_forward(
             window_left=window_size[0],
             return_lse=True,
         )
-        lse = lse.transpose(0, 1)
+        # FlashInfer unbatched LSE is (H, S). Canonical ring layout is (B, H, S).
         out, lse = out.unsqueeze(0), lse.unsqueeze(0)
     elif q.ndim == 3:
         out, lse = single_prefill_with_kv_cache(
@@ -306,7 +309,7 @@ def flashinfer_attn_forward(
             window_left=window_size[0],
             return_lse=True,
         )
-        lse = lse.transpose(0, 1)
+        # Unbatched (H, S) matches canonical heads-major layout without a batch dim.
     else:
         raise ValueError(f"Invalid input shape: {q.shape}")
     lse = lse / _LOG2_E

vllm_omni/diffusion/attention/backends/ring/ring_utils.py

@@ -2,91 +2,77 @@
 # Copyright (c) 2024, Jiarui Fang.
 # Adapted from https://github.com/feifeibear/long-context-attention
 
+from typing import Literal
 
 import torch
 import torch.nn.functional as F
-from vllm.logger import init_logger
 
-logger = init_logger(__name__)
+LseLayout = Literal["bhs", "bsh"]
 
 __all__ = ["update_out_and_lse", "flatten_varlen_lse", "unflatten_varlen_lse"]
 
 
-# Remove torch.jit.script for debugging and flexible shape handling
+def _normalize_lse(block_lse: torch.Tensor, block_out: torch.Tensor, lse_layout: LseLayout) -> torch.Tensor:
+    """Normalize a ring kernel LSE to ``(B, S, H, 1)``.
+
+    ``lse_layout`` is required. When ``seq_len == num_heads``, ``(B, H, S)`` and
+    ``(B, S, H)`` are the same shape, so guessing from dimension sizes would
+    silently transpose a non-symmetric LSE.
+    """
+    if block_out.ndim != 4:
+        raise ValueError(f"Ring attention output must be 4D (B, S, H, D), got {tuple(block_out.shape)}")
+
+    batch, seq_len, num_heads, _ = block_out.shape
+    if block_lse.shape[0] != batch:
+        raise ValueError(f"Ring LSE batch dimension {block_lse.shape[0]} does not match output batch {batch}.")
+
+    if block_lse.ndim == 4:
+        if block_lse.shape[-1] != 1:
+            raise ValueError(f"4D Ring LSE must end in a singleton dimension, got {tuple(block_lse.shape)}")
+        block_lse = block_lse.squeeze(-1)
+    elif block_lse.ndim != 3:
+        raise ValueError(f"Ring LSE must be 3D or 4D, got {tuple(block_lse.shape)}")
+
+    if lse_layout == "bhs":
+        if block_lse.shape[1] != num_heads or block_lse.shape[2] < seq_len:
+            raise ValueError(
+                f"Ring LSE shape {tuple(block_lse.shape)} is incompatible with output "
+                f"shape {tuple(block_out.shape)}; expected (B, H, S) for lse_layout='bhs'."
+            )
+        return block_lse[:, :, :seq_len].transpose(1, 2).unsqueeze(-1)
+    if lse_layout == "bsh":
+        if block_lse.shape[1] < seq_len or block_lse.shape[2] != num_heads:
+            raise ValueError(
+                f"Ring LSE shape {tuple(block_lse.shape)} is incompatible with output "
+                f"shape {tuple(block_out.shape)}; expected (B, S, H) for lse_layout='bsh'."
+            )
+        return block_lse[:, :seq_len, :].unsqueeze(-1)
+
+    raise ValueError(f"Unknown Ring LSE layout {lse_layout!r}; expected 'bhs' or 'bsh'.")
+
+
 def _update_out_and_lse(
     out: torch.Tensor,
     lse: torch.Tensor,
     block_out: torch.Tensor,
     block_lse: torch.Tensor,
+    lse_layout: LseLayout,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     block_out = block_out.to(torch.float32)
-
+    if out.shape != block_out.shape:
-    B, S, H, D = out.shape
+        raise ValueError(
-
+            f"Ring attention block output shape {tuple(block_out.shape)} does not match "
-    # --- Shape Correction Logic for block_lse ---
+            f"accumulated output shape {tuple(out.shape)}."
-    # Goal: block_lse should be (B, S, H, 1) to match out (B, S, H, D)
+        )
-
+    block_lse = _normalize_lse(block_lse, block_out, lse_layout)
-    # Debug info
-    # print(f"DEBUG _update: out={out.shape}, block_lse={block_lse.shape}")
-
-    # Case 0: If block_lse is already 4D, check if it matches
-    if block_lse.dim() == 4:
-        if block_lse.shape[1] == S and block_lse.shape[2] == H:
-            pass  # Good
-        elif block_lse.shape[1] == H and block_lse.shape[2] == S:
-            block_lse = block_lse.transpose(1, 2)
-        elif block_lse.shape[1] == H and block_lse.shape[2] >= S:  # Padding case
-            block_lse = block_lse[:, :, :S, :].transpose(1, 2)
-        # If shape is (B, H, S, 1) but expected (B, S, H, 1) because out is (B, S, H, D)
-        elif block_lse.shape[1] == H and block_lse.shape[2] == S and block_lse.shape[3] == 1:
-            block_lse = block_lse.transpose(1, 2)
-
-    # Case 1: block_lse is 3D (B, H, S) or (B, S, H) or (B, ?, ?)
-    elif block_lse.dim() == 3:
-        # Check for (B, H, S) - Standard SDPA/FA output
-        if block_lse.shape[1] == H and block_lse.shape[2] == S:
-            block_lse = block_lse.transpose(1, 2).unsqueeze(-1)
-
-        # Check for (B, S, H)
-        elif block_lse.shape[1] == S and block_lse.shape[2] == H:
-            block_lse = block_lse.unsqueeze(-1)
-
-        # Check for Padding: (B, H, S_pad) where S_pad >= S
-        elif block_lse.shape[1] == H and block_lse.shape[2] >= S:
-            # print(f"DEBUG: Trimming padding from lse. {block_lse.shape} -> S={S}")
-            block_lse = block_lse[:, :, :S].transpose(1, 2).unsqueeze(-1)
-
-        # Check for weird case: (B, S, H_pad) ? Unlikely for LSE but possible
-        elif block_lse.shape[1] == S and block_lse.shape[2] >= H:
-            block_lse = block_lse[:, :, :H].unsqueeze(-1)
-
-        # Check for flipped weird case: (B, S_pad, H)
-        elif block_lse.shape[1] >= S and block_lse.shape[2] == H:
-            block_lse = block_lse[:, :S, :].unsqueeze(-1)
-
-    # --- Shape Correction for lse (internal state) ---
-    # Ensure lse matches block_lse's corrected shape (B, S, H, 1)
     if lse.shape != block_lse.shape:
-        # If lse was initialized with wrong shape, try to fix it
+        raise ValueError(
-        if lse.dim() == 4 and lse.shape[1] == block_lse.shape[2] and lse.shape[2] == block_lse.shape[1]:
+            f"Ring attention block LSE shape {tuple(block_lse.shape)} does not match "
-            lse = lse.transpose(1, 2)
+            f"accumulated LSE shape {tuple(lse.shape)}."
-        elif lse.shape[1] >= S:  # slice if lse was initialized with padding
+        )
-            lse = lse[:, :S, :, :]
 
-    # Final check
+    out = out - F.sigmoid(block_lse - lse) * (out - block_out)
-    if lse.shape != block_lse.shape:
+    lse = lse - F.logsigmoid(lse - block_lse)
-        # Force broadcast if possible?
-        pass
-
-    try:
-        out = out - F.sigmoid(block_lse - lse) * (out - block_out)
-        lse = lse - F.logsigmoid(lse - block_lse)
-    except RuntimeError as e:
-        logger.error("_update_out_and_lse failed: %s", e)
-        logger.error("out: %s, lse: %s", out.shape, lse.shape)
-        logger.error("block_out: %s, block_lse: %s", block_out.shape, block_lse.shape)
-        # raise e
-        raise e
 
     return out, lse
 
@@ -97,64 +83,22 @@ def update_out_and_lse(
     block_out: torch.Tensor,
     block_lse: torch.Tensor,
     slice_=None,
+    *,
+    lse_layout: LseLayout,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     if out is None:
         if slice_ is not None:
             raise RuntimeError("first update_out_and_lse should not pass slice_ args")
 
         out = block_out.to(torch.float32)
-
+        lse = _normalize_lse(block_lse, out, lse_layout)
-        # Initialize LSE with robust logic (same as _update)
-        B, D1, D2, D3 = out.shape
-
-        S_guess = D1
-        H_guess = D2
-
-        if block_lse.dim() == 3:
-            if block_lse.shape[1] == H_guess and block_lse.shape[2] == S_guess:
-                lse = block_lse.transpose(1, 2).unsqueeze(-1)
-            elif block_lse.shape[1] == S_guess and block_lse.shape[2] == H_guess:
-                lse = block_lse.unsqueeze(-1)
-            elif block_lse.shape[1] == H_guess and block_lse.shape[2] >= S_guess:  # Padding
-                lse = block_lse[:, :, :S_guess].transpose(1, 2).unsqueeze(-1)
-            elif block_lse.shape[1] == S_guess and block_lse.shape[2] >= H_guess:  # Padding/Weird
-                lse = block_lse[:, :, :H_guess].unsqueeze(-1)
-            elif block_lse.shape[1] >= S_guess and block_lse.shape[2] == H_guess:
-                lse = block_lse[:, :S_guess, :].unsqueeze(-1)
-
-            # Reverse case: What if out is (B, H, S, D) so S=D2, H=D1?
-            elif block_lse.shape[1] == D1 and block_lse.shape[2] >= D2:  # Matches (H, S)
-                # Then out is (B, H, S, D). We should transpose out!
-                out = out.transpose(1, 2)
-                lse = block_lse[:, :, :D2].transpose(1, 2).unsqueeze(-1)  # (B, S, H, 1)
-
-            else:
-                # Fallback
-                lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1)
-        else:
-            # Case 0: If block_lse is already 4D, check if it matches
-            if block_lse.dim() == 4:
-                if block_lse.shape[1] == S_guess and block_lse.shape[2] == H_guess:
-                    lse = block_lse
-                elif block_lse.shape[1] == H_guess and block_lse.shape[2] == S_guess:
-                    lse = block_lse.transpose(1, 2)
-                elif block_lse.shape[1] == H_guess and block_lse.shape[2] >= S_guess:  # Padding case
-                    lse = block_lse[:, :, :S_guess, :].transpose(1, 2)
-                elif block_lse.shape[1] == D1 and block_lse.shape[2] >= D2:  # Matches (H, S)
-                    # Then out is (B, H, S, D). We should transpose out!
-                    out = out.transpose(1, 2)
-                    lse = block_lse[:, :, :D2].transpose(1, 2)  # (B, S, H, 1)
-                else:
-                    lse = block_lse
-            else:
-                lse = block_lse
 
     elif slice_ is not None:
         slice_out, slice_lse = out[slice_], lse[slice_]
-        slice_out, slice_lse = _update_out_and_lse(slice_out, slice_lse, block_out, block_lse)
+        slice_out, slice_lse = _update_out_and_lse(slice_out, slice_lse, block_out, block_lse, lse_layout)
         out[slice_], lse[slice_] = slice_out, slice_lse
     else:
-        out, lse = _update_out_and_lse(out, lse, block_out, block_lse)
+        out, lse = _update_out_and_lse(out, lse, block_out, block_lse, lse_layout)
     return out, lse
 
 

vllm_omni/diffusion/attention/backends/ring_flash_attn.py

@@ -95,7 +95,8 @@ def ring_flash_attn_forward(
             if attn_type == AttnType.SPARSE_SAGE:
                 out, lse = block_out, block_lse
             else:
-                out, lse = update_out_and_lse(out, lse, block_out, block_lse)
+                # Ring kernel wrappers canonicalize LSE to (B, H, S).
+                out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs")
 
         if step + 1 != comm.world_size:
             comm.wait()

vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py

@@ -116,7 +116,7 @@ class RingAttentionFunc(torch.autograd.Function):
                     causal=is_causal and step == 0,
                     op_type=op_type,
                 )
-                out, lse = update_out_and_lse(out, lse, block_out, block_lse)
+                out, lse = update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs")
 
             if step + 1 != comm.world_size:
                 comm.wait()

vllm_omni/diffusion/attention/backends/sage_attn.py

@@ -47,7 +47,7 @@ else:
             "SageAttentionBackend is not available. You may install sage-attention"
             " by pip install git+https://github.com/thu-ml/SageAttention.git"
         )
-        raise ImportError
+        sageattn = None
 
 # TODO add sage3 attention backend
 
@@ -92,6 +92,13 @@ class SageAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
+        if attn_metadata is not None and attn_metadata.attn_mask is not None:
+            raise ValueError("SAGE_ATTN does not support attn_mask. Select a mask-capable backend.")
+        if sageattn is None:
+            raise ImportError(
+                "SAGE_ATTN requires sageattention. Install with: "
+                "pip install git+https://github.com/thu-ml/SageAttention.git"
+            )
         output = sageattn(
             query,
             key,

vllm_omni/diffusion/attention/backends/sage_attn3.py

@@ -1,8 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import math
+
 import torch
-import torch.nn.functional as F
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.attention.backends.abstract import (
@@ -65,8 +66,6 @@ class SageAttention3Backend(AttentionBackend):
 
 
 class SageAttention3Impl(AttentionImpl):
-    _warned_gqa_fallback_global: bool = False
-
     def __init__(
         self,
         num_heads: int,
@@ -80,6 +79,12 @@ class SageAttention3Impl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
         self.dropout = extra_impl_args.get("dropout_p", 0.0)
+        expected_scale = head_size**-0.5
+        if not math.isclose(softmax_scale, expected_scale, rel_tol=1e-6):
+            raise ValueError(
+                "SAGE_ATTN_3 does not expose a custom softmax scale; "
+                f"expected {expected_scale} for head_size={head_size}, got {softmax_scale}."
+            )
 
     def forward_cuda(
         self,
@@ -88,6 +93,8 @@ class SageAttention3Impl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
+        if attn_metadata is not None and attn_metadata.attn_mask is not None:
+            raise ValueError("SAGE_ATTN_3 does not support attn_mask. Select a mask-capable backend.")
         query = query.transpose(1, 2).contiguous()
         key = key.transpose(1, 2).contiguous()
         value = value.transpose(1, 2).contiguous()
@@ -98,19 +105,11 @@ class SageAttention3Impl(AttentionImpl):
                     "GQA/MQA requires query heads to be a multiple of KV heads, "
                     f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}"
                 )
-            if not type(self)._warned_gqa_fallback_global:
+            raise NotImplementedError(
-                logger.warning("SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA.")
+                "SAGE_ATTN_3 was explicitly selected but does not support GQA/MQA "
-                type(self)._warned_gqa_fallback_global = True
+                f"(q_heads={query.shape[1]}, kv_heads={key.shape[1]}). Select a GQA-capable backend."
-            output = F.scaled_dot_product_attention(
-                query,
-                key,
-                value,
-                is_causal=self.causal,
-                dropout_p=self.dropout,
-                scale=self.softmax_scale,
-                enable_gqa=True,
             )
-        else:
+
-            output = _sageattn3_blackwell_op(query, key, value, self.causal)
+        output = _sageattn3_blackwell_op(query, key, value, self.causal)
 
         return output.transpose(1, 2).contiguous()

vllm_omni/diffusion/attention/backends/sdpa.py

@@ -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 typing import Literal
 
@@ -56,7 +56,8 @@ class SDPABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
     @classmethod
-    def supports_attention_mask(cls) -> bool:
+    def supports_attention_mask(cls, attention_spec: object | None = None) -> bool:
+        del attention_spec
         return True
 
     @staticmethod

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2025 The Fairseq Authors and the HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,7 +15,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flash_attention_utils.py
+from collections.abc import Callable
 from functools import lru_cache
+from typing import Any
 
 import torch
 import torch.nn.functional as F
@@ -20,6 +25,8 @@ from vllm.logger import init_logger
 
 from vllm_omni.platforms import current_omni_platform
 
+FlashAttnFn = Callable[..., Any]
+
 logger = init_logger(__name__)
 
 
@@ -31,9 +38,27 @@ def _is_blackwell() -> bool:
     return capability is not None and capability.major >= 10
 
 
-# Flash Attention function detection with fallback chain
+@lru_cache(maxsize=1)
-flash_attn_func = None
+def is_flash_attn_4_available() -> bool:
-flash_attn_varlen_func = None
+    """Return whether CuTe FlashAttention-4 (``flash_attn.cute``) can be imported.
+
+    FA2/FA3 wheels are often importable on Blackwell but only ship Hopper kernels.
+    Those must not be treated as a runnable FLASH_ATTN implementation here.
+    """
+    try:
+        from flash_attn.cute import flash_attn_varlen_func  # noqa: F401
+
+        return True
+    except Exception as exc:
+        logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc)
+        return False
+
+
+# Flash Attention implementation discovery. All candidates implement the same
+# logical FLASH_ATTN backend; this never substitutes SDPA at runtime.
+# Bind via aliases so mypy does not treat each candidate import as a redefinition.
+flash_attn_func: FlashAttnFn | None = None
+flash_attn_varlen_func: FlashAttnFn | None = None
 
 if current_omni_platform.is_rocm():
     # ROCm: try Aiter first
@@ -41,7 +66,11 @@ if current_omni_platform.is_rocm():
         from vllm._aiter_ops import is_aiter_found_and_supported
 
         if is_aiter_found_and_supported():
-            from aiter import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+            from aiter import flash_attn_func as _fa_func
+            from aiter import flash_attn_varlen_func as _fa_varlen
+
+            flash_attn_func = _fa_func
+            flash_attn_varlen_func = _fa_varlen
     except (ImportError, ModuleNotFoundError):
         pass
 elif current_omni_platform.is_xpu():
@@ -53,64 +82,84 @@ elif current_omni_platform.is_xpu():
         pass
 elif current_omni_platform.is_musa():
     try:
-        from flash_attn_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+        from flash_attn_interface import flash_attn_func as _fa_func
+        from flash_attn_interface import flash_attn_varlen_func as _fa_varlen
+
+        flash_attn_func = _fa_func
+        flash_attn_varlen_func = _fa_varlen
     except (ImportError, ModuleNotFoundError):
         pass
 else:
-    # CUDA: prefer the CuTe-based FA4 package when present.  This matters on
+    # CUDA implementation candidates, ordered by preference.
-    # Blackwell, where an importable Hopper-only FA3 wheel may otherwise be
+    # Candidate 1: CuTe-based FA4. This matters on Blackwell, where an
-    # selected and fail at launch with "no kernel image is available".
+    # importable Hopper-only FA3 wheel may otherwise be selected and fail
+    # at launch with "no kernel image is available".
     if _is_blackwell():
-        try:
+        if is_flash_attn_4_available():
-            from flash_attn.cute import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+            try:
+                from flash_attn.cute import flash_attn_func as _fa_func
+                from flash_attn.cute import flash_attn_varlen_func as _fa_varlen
 
-            logger.info("Using CuTe FlashAttention-4 on Blackwell")
+                flash_attn_func = _fa_func
-        except Exception as exc:
+                flash_attn_varlen_func = _fa_varlen
-            # Optional FA4 dependencies may be present but ABI-incompatible
+                logger.info("Using CuTe FlashAttention-4 on Blackwell")
-            # (for example, mismatched CUTLASS DSL and Quack builds). Treat
+            except Exception as exc:
-            # that as unavailable so backend discovery can continue.
+                # Optional FA4 dependencies may be present but ABI-incompatible
-            logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc)
+                # (for example, mismatched CUTLASS DSL and Quack builds). Treat
+                # that as unavailable so backend discovery can continue.
+                logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc)
+    else:
+        # Candidate 2: FA3 from fa3-fwd PyPI package.
+        if flash_attn_func is None:
+            try:
+                from fa3_fwd_interface import flash_attn_func as _fa_func
+                from fa3_fwd_interface import flash_attn_varlen_func as _fa_varlen
 
-    # Try FA3 -> FA2 fallback chain.
+                flash_attn_func = _fa_func
-    # Try FA3 from fa3-fwd PyPI package
+                flash_attn_varlen_func = _fa_varlen
-    if flash_attn_func is None:
+            except (ImportError, ModuleNotFoundError):
-        try:
+                pass
-            from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
-        except (ImportError, ModuleNotFoundError):
-            pass
 
-    # Fallback: Try FA3 from flash-attention source build
+        # Candidate 3: FA3 from a flash-attention source build.
-    if flash_attn_func is None:
+        if flash_attn_func is None:
-        try:
+            try:
-            from flash_attn_interface import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+                from flash_attn_interface import flash_attn_func as _fa_func
-        except (ImportError, ModuleNotFoundError):
+                from flash_attn_interface import flash_attn_varlen_func as _fa_varlen
-            pass
 
-    # Fallback: Try FA2 from flash-attn package (try multiple import paths)
+                flash_attn_func = _fa_func
-    if flash_attn_func is None:
+                flash_attn_varlen_func = _fa_varlen
-        try:
+            except (ImportError, ModuleNotFoundError):
-            from flash_attn import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+                pass
-        except (ImportError, ModuleNotFoundError):
-            pass
 
-    if flash_attn_func is None:
+        # Candidate 4: FA2 from the flash-attn package (multiple import paths).
-        try:
+        if flash_attn_func is None:
-            from flash_attn.flash_attn_interface import (  # noqa: F401
+            try:
-                flash_attn_func,
+                from flash_attn import flash_attn_func as _fa_func
-                flash_attn_varlen_func,
+                from flash_attn import flash_attn_varlen_func as _fa_varlen
-            )
-        except (ImportError, ModuleNotFoundError):
-            pass
 
-    # Fallback: Try vLLM's encapsulated flash attention dispatcher
+                flash_attn_func = _fa_func
-    # TODO discuss priority and remove potentially redundant fallbacks
+                flash_attn_varlen_func = _fa_varlen
-    if flash_attn_varlen_func is None:
+            except (ImportError, ModuleNotFoundError):
-        try:
+                pass
-            from vllm.vllm_flash_attn import (  # noqa: F401
+
-                flash_attn_varlen_func,
+        if flash_attn_func is None:
-            )
+            try:
-        except (ImportError, ModuleNotFoundError):
+                from flash_attn.flash_attn_interface import flash_attn_func as _fa_func
-            pass
+                from flash_attn.flash_attn_interface import flash_attn_varlen_func as _fa_varlen
+
+                flash_attn_func = _fa_func
+                flash_attn_varlen_func = _fa_varlen
+            except (ImportError, ModuleNotFoundError):
+                pass
+
+        # Candidate 5: vLLM's encapsulated Flash Attention dispatcher.
+        if flash_attn_varlen_func is None:
+            try:
+                from vllm.vllm_flash_attn import flash_attn_varlen_func as _fa_varlen
+
+                flash_attn_varlen_func = _fa_varlen
+            except (ImportError, ModuleNotFoundError):
+                pass
 
 # If no FA backend available, SDPA backend will be selected at the platform level
 # flash_attn_func and flash_attn_varlen_func will be None
@@ -127,12 +176,7 @@ def is_flash_attn_installed() -> bool:
         # Check for any FA backend: FA4 (flash_attn.cute), FA3
         # (fa3_fwd_interface, flash_attn_interface), or FA2 (flash_attn).
         if _is_blackwell():
-            try:
+            return is_flash_attn_4_available()
-                from flash_attn.cute import flash_attn_varlen_func  # noqa: F401
-
-                return True
-            except Exception as exc:
-                logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc)
 
         # Try FA3 from fa3-fwd PyPI package
         try:
@@ -167,7 +211,7 @@ def is_flash_attn_installed() -> bool:
 
         return True
     except (ImportError, ModuleNotFoundError):
-        logger.warning("No Flash Attention backend found, using pytorch SDPA implementation")
+        logger.warning("No Flash Attention implementation found")
         return False
 
 

vllm_omni/diffusion/attention/layer.py

@@ -98,6 +98,7 @@ class Attention(nn.Module):
         # set_current_diffusion_config(); no env-var re-parsing needed here.
         backend_kwargs: dict | None = None
         self.backend_pref = None
+        self.backend_explicit = False
 
         config = get_current_diffusion_config_or_none()
         attention_config = config.diffusion_attention_config if config is not None else None
@@ -129,22 +130,17 @@ class Attention(nn.Module):
             # FLASH_ATTN is an Omni selector, not a device-specific kernel.
             # vLLM resolves it to CUDA FlashAttention on GPU and Ascend native
             # attention on NPU. A platform may legitimately default dense
-            # attention to SDPA (for example when an optional diffusion FA
+            # attention to CUDNN/SDPA (for example when optional dense FA
-            # package is absent), but a paged layer must advertise the
+            # extras are absent), but a paged layer must advertise the
             # capability so the Worker can register its native cache view.
-            # Resolve the selector only for this marked layer; unmarked dense
+            # Load that selector from the registry: this is not a user-explicit
-            # attention keeps the platform default unchanged.
+            # dense FLASH_ATTN request, so skip platform explicit-backend
-            from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+            # validation (Blackwell would otherwise require CuTe FA4). Formal
+            # paged execution still delegates to vLLM's native paged backend.
+            from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
 
             dense_backend_name = attn_backend_cls.get_name()
-            paged_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN"))
+            attn_backend_cls = DiffusionAttentionBackendEnum.FLASH_ATTN.get_class()
-            attn_backend_cls, spec = get_attn_backend_for_role(
-                role=role,
-                head_size=head_size,
-                attention_config=paged_attention_config,
-                role_category=role_category,
-                allow_trtllm_default=allow_trtllm_default,
-            )
             logger.info(
                 "Resolved marked paged diffusion attention role=%r to %r because the platform default %r "
                 "does not support Scheduler-owned KV",
@@ -161,12 +157,17 @@ class Attention(nn.Module):
                 "TRTLLM_ATTN does not support AllGather-KV sequence parallelism. "
                 "Set --allgather-degree 1 or select another diffusion attention backend."
             )
+        self.attn_spec = spec
         if spec is not None:
             backend_kwargs = spec.backend_kwargs()
             self.backend_pref = spec.backend
+            self.backend_explicit = True
             logger.debug("Attention(role=%s) → backend=%s", role, spec.backend)
         else:
-            logger.debug("Attention(role=%s) → platform default", role)
+            # Propagate the resolved platform default so Ring Attention can
+            # make a compatible automatic selection on the current GPU.
+            self.backend_pref = attn_backend_cls.get_name()
+            logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref)
 
         self.attn_backend = attn_backend_cls
         self.attn_impl_cls = self.attn_backend.get_impl_cls()
@@ -180,8 +181,10 @@ class Attention(nn.Module):
             prefix=prefix,
             backend_kwargs=backend_kwargs,
             role=role,
+            backend_explicit=self.backend_explicit,
         )
-        # Instantiate fallback backend for float32 support
+        # Some model-specific paths call this helper directly. The main backend
+        # dispatch below never switches to it implicitly.
         self.sdpa_fallback = SDPABackend.get_impl_cls()(
             num_heads=num_heads,
             head_size=head_size,
@@ -205,16 +208,13 @@ class Attention(nn.Module):
         if config is not None:
             if config.parallel_config.ring_degree > 1:
                 self.use_ring = True
-                try:
+                sp_group = get_sp_group()
-                    sp_group = get_sp_group()
+                self.ring_pg = sp_group.ring_group
-                    self.ring_pg = sp_group.ring_group
+                self.ring_runner = RingParallelAttention(
-                    self.ring_runner = RingParallelAttention(
+                    sp_group,
-                        sp_group,
+                    attn_backend_pref=self.backend_pref,
-                        attn_backend_pref=self.backend_pref,
+                    attn_backend_explicit=self.backend_explicit,
-                    )
+                )
-                except Exception:
-                    self.use_ring = False
-                    self.ring_runner = None
 
         self.parallel_strategy = build_parallel_attention_strategy(
             scatter_idx=scatter_idx,
@@ -222,7 +222,7 @@ class Attention(nn.Module):
             use_sync=use_sync,
             causal=causal,
         )
-        # Fallback strategy when SP is not active (outside sharded regions)
+        # Local strategy when SP is intentionally inactive outside sharded regions.
         self._no_parallel_strategy = NoParallelAttention()
 
         self.layer_idx: int | None = _try_extract_layer_index(prefix)
@@ -285,14 +285,11 @@ class Attention(nn.Module):
                 )
             platform_key = current_omni_platform.device_name
             if not self.attention.supports_kv_cache_dtype(dtype, platform_key):
-                logger.warning_once(
+                raise ValueError(
-                    "Attention backend %s does not support kv_cache_dtype='%s' on %s. "
+                    f"Attention backend {self.attn_backend.get_name()} does not support "
-                    "KV quantization will be disabled.",
+                    f"kv_cache_dtype={dtype!r} on {platform_key}. Select a compatible "
-                    self.attn_backend.get_name(),
+                    "backend or set diffusion_kv_cache_dtype='auto'."
-                    dtype,
-                    platform_key,
                 )
-                dtype = None
         self._kv_cache_dtype = dtype
         self._kv_cache_skip_steps = getattr(config, "diffusion_kv_cache_skip_step_indices", None)
         self._kv_cache_skip_layers = getattr(config, "diffusion_kv_cache_skip_layer_indices", None)
@@ -470,42 +467,44 @@ class Attention(nn.Module):
         return self.paged_kv_cache_role is not None and self._active_paged_kv_adapter() is not None
 
     def _run_local_attention(self, query, key, value, attn_metadata):
-        self._assert_piecewise_compatible(attn_metadata)
+        self._assert_metadata_compatible(attn_metadata)
-
-        if query.dtype == torch.float32:
-            logger.warning_once(
-                f"Only SDPA supports float32. Overriding user config {type(self.attention)} "
-                f"attention_backend='{self.backend_pref}' to 'sdpa' for dtype={query.dtype}."
-            )
-            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
 
         in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile
-        # NPU only: the startup KV-capacity profile needs tensor shapes, not a
+        # The startup KV-capacity profile needs tensor shapes, not a paged
-        # paged attention result. If dense MindIE-SD is absent, SDPA provides
+        # attention result. If dense FLASH_ATTN deps are absent (NPU MindIE-SD
-        # that profile forward. Formal GPU/NPU paged requests never use this
+        # or CUDA CuTe FA4), SDPA provides that profile forward. Formal paged
-        # branch because their Worker adapter is active.
+        # requests never use this branch because their Worker adapter is active.
+        # Do not silently remap float32 (or other dtypes) to SDPA: a
+        # user-explicit backend must run or raise.
         if (
             self._scheduler_paged_kv
             and self.paged_kv_cache_role is not None
             and in_kv_memory_profile
-            and current_omni_platform.is_npu()
             and self.attn_backend.get_name() == "FLASH_ATTN"
             and not current_omni_platform.supports_diffusion_dense_flash_attention()
         ):
             logger.warning_once(
-                "The startup KV memory profile is using SDPA because MindIE-SD is unavailable. "
+                "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. "
                 "Formal paged requests still use the platform-native paged attention backend."
             )
             return self.sdpa_fallback.forward(query, key, value, attn_metadata)
 
         return self.attention.forward(query, key, value, attn_metadata)
 
-    def _assert_piecewise_compatible(self, attn_metadata: AttentionMetadata | None) -> None:
+    def _assert_metadata_compatible(self, attn_metadata: AttentionMetadata | None) -> None:
-        if attn_metadata is None or attn_metadata.full_attn_spans is None:
+        if attn_metadata is None:
+            return
+        backend_name = self.attn_backend.get_name()
+        if attn_metadata.attn_mask is not None and not self.attn_backend.supports_attention_mask(
+            getattr(self, "attn_spec", None)
+        ):
+            raise ValueError(
+                f"Attention backend '{backend_name}' does not support attn_mask. Select a mask-capable backend."
+            )
+        if attn_metadata.full_attn_spans is None:
             return
         if attn_metadata.attn_mask is not None and attn_metadata.attn_mask.ndim == 4:
             return
-        backend_name = self.attn_backend.get_name()
         if not self.attn_backend.supports_piecewise_spans:
             raise ValueError(
                 f"Attention backend '{backend_name}' does not support "
@@ -515,6 +514,8 @@ class Attention(nn.Module):
             )
 
     def _run_ring_attention(self, query, key, value, attn_metadata):
+        if attn_metadata is not None and attn_metadata.attn_mask is not None:
+            raise ValueError("Ring attention does not support attn_mask; use Ulysses SP or disable mask_sp_padding.")
         skip = getattr(self.attention, "skip", None)
         if skip is not None and getattr(skip, "configured", False):
             raise NotImplementedError(

vllm_omni/diffusion/attention/parallel/factory.py

@@ -18,7 +18,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     get_sequence_parallel_world_size,
     get_sp_group,
 )
-from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 
 logger = init_logger(__name__)
 
@@ -37,32 +37,32 @@ def build_parallel_attention_strategy(
     - Parallel attention selection is handled here, based on distributed config
       and initialized process groups.
     """
-    try:
+    if not is_forward_context_available():
-        cfg = get_forward_context().omni_diffusion_config
-        p = cfg.parallel_config
-    except Exception as e:
-        logger.debug(f"No forward context available for parallel attention strategy: {e}")
         return NoParallelAttention()
+    cfg = get_forward_context().omni_diffusion_config
+    p = cfg.parallel_config
 
     ulysses_degree = getattr(p, "ulysses_degree", 1)
     ring_degree = getattr(p, "ring_degree", 1)
     allgather_degree = getattr(p, "allgather_degree", 1)
     ulysses_a2a_permute = getattr(p, "ulysses_a2a_permute", False)
 
+    sp_configured = ulysses_degree > 1 or ring_degree > 1 or allgather_degree > 1
+    if not sp_configured:
+        return NoParallelAttention()
+
     try:
         sp_group = get_sp_group()
-        # Ensure SP group is initialized and world size > 1
-        if get_sequence_parallel_world_size() <= 1:
-            return NoParallelAttention()
     except Exception as e:
-        # Log warning if SP is configured but group is not available
+        raise RuntimeError(
-        if ulysses_degree > 1 or ring_degree > 1 or allgather_degree > 1:
+            f"SP is configured (ulysses={ulysses_degree}, ring={ring_degree}, "
-            logger.warning(
+            f"allgather={allgather_degree}), but the SP group is unavailable."
-                f"SP configured (ulysses={ulysses_degree}, ring={ring_degree}, "
+        ) from e
-                f"allgather={allgather_degree}) but SP group not available: {e}. "
+    if get_sequence_parallel_world_size() <= 1:
-                f"Falling back to NoParallelAttention. This may cause incorrect results."
+        raise RuntimeError(
-            )
+            f"SP is configured (ulysses={ulysses_degree}, ring={ring_degree}, "
-        return NoParallelAttention()
+            f"allgather={allgather_degree}), but the initialized SP world size is not greater than one."
+        )
 
     if allgather_degree > 1:
         if causal:

vllm_omni/diffusion/attention/parallel/ring.py

@@ -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 __future__ import annotations
 
@@ -10,7 +10,13 @@ import torch
 from vllm.logger import init_logger
 
 # import torch.distributed as dist # Not used directly here, but good practice if needed
-from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_AITER, HAS_FA3, HAS_FA4, HAS_FLASH_ATTN
+from vllm_omni.diffusion.attention.backends.ring.ring_globals import (
+    FA3_SUPPORTED_CUDA_MAJORS,
+    HAS_AITER,
+    HAS_FA3,
+    HAS_FA4,
+    HAS_FLASH_ATTN,
+)
 from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
 from vllm_omni.diffusion.attention.parallel.base import (
     ParallelAttentionContext,
@@ -24,6 +30,41 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _can_use_fa3(device: torch.device) -> bool:
+    """Return whether the installed FA3 kernels support ``device``.
+
+    Importing an extension only proves that its Python module is installed.
+    When the source publishes a supported-major contract, also check it before
+    entering a CUDA launcher that may abort instead of raising an exception.
+    """
+    if not HAS_FA3 or device.type != "cuda":
+        return False
+    if FA3_SUPPORTED_CUDA_MAJORS is None:
+        return True
+    major, _minor = torch.cuda.get_device_capability(device)
+    return major in FA3_SUPPORTED_CUDA_MAJORS
+
+
+def _can_use_fa4(device: torch.device) -> bool:
+    """Return whether the installed FA4 kernels support ``device``."""
+    if not HAS_FA4 or device.type != "cuda":
+        return False
+    major, _minor = torch.cuda.get_device_capability(device)
+    return major >= 10
+
+
+def _can_use_fa2(device: torch.device) -> bool:
+    """Return whether the FA2 backend supports ``device``.
+
+    FA2's CUDA backend supports Ampere, Ada, and Hopper.  Blackwell support is
+    provided by the separate FA4 implementation.
+    """
+    if not HAS_FLASH_ATTN or device.type != "cuda":
+        return False
+    major, _minor = torch.cuda.get_device_capability(device)
+    return major in (8, 9)
+
+
 @dataclass(frozen=True, slots=True)
 class _RingCtx(ParallelAttentionContext):
     """Per-forward context for Ring sequence-parallel attention."""
@@ -46,9 +87,11 @@ class RingParallelAttention:
         self,
         sp_group: SequenceParallelGroupCoordinator,
         attn_backend_pref: str | None = None,
+        attn_backend_explicit: bool = False,
     ) -> None:
         self._sp_group = sp_group
         self.attn_backend_pref = attn_backend_pref
+        self.attn_backend_explicit = attn_backend_explicit
 
     @property
     def enabled(self) -> bool:
@@ -110,14 +153,6 @@ class RingParallelAttention:
         if backend_pref is not None:
             backend_pref = backend_pref.lower()
 
-        # FP32 is not supported by Flash Attention, force SDPA
-        if query.dtype == torch.float32:
-            backend_pref = "sdpa"
-        elif not HAS_FA3 and not HAS_FLASH_ATTN and not HAS_AITER:
-            if backend_pref != "sdpa":
-                logger.warning_once("Flash Attention (FA2/FA3/AITER) is not available! Force enabling SDPA.")
-            backend_pref = "sdpa"
-
         # Extract joint tensors
         joint_key, joint_value = None, None
         joint_strategy = "front"
@@ -127,7 +162,7 @@ class RingParallelAttention:
             if attn_metadata.joint_strategy is not None:
                 joint_strategy = attn_metadata.joint_strategy
 
-        if backend_pref in {"sdpa", "torch", "torch_sdpa"}:
+        def _run_sdpa() -> torch.Tensor:
             from vllm_omni.diffusion.attention.backends.ring_pytorch_attn import ring_pytorch_attn_func
 
             return ring_pytorch_attn_func(
@@ -143,19 +178,72 @@ class RingParallelAttention:
                 joint_strategy=joint_strategy,
             )
 
+        # Ring only implements local Flash/AITER and SDPA kernels. HuggingFace
+        # Hub FA and other non-local backends can use the SDPA ring path only
+        # when they came from automatic platform selection. An explicit
+        # request must never silently substitute local FA4/FA3/FA2.
+        _sdpa_prefs = {
+            "sdpa",
+            "torch",
+            "torch_sdpa",
+        }
+        _non_ring_prefs = {
+            "cudnn_attn",
+            "flashinfer_attn",
+            "trtllm_attn",
+            "sage_attn",
+            "sage_attn_3",
+            "flash_attn_hub",
+            "flash_attn_3_hub",
+        }
+        if backend_pref in _sdpa_prefs:
+            return _run_sdpa()
+        if backend_pref in _non_ring_prefs:
+            if self.attn_backend_explicit:
+                raise ValueError(
+                    f"{self.attn_backend_pref} was explicitly selected, but ring sequence parallelism "
+                    "has no implementation for that backend. Select TORCH_SDPA/FLASH_ATTN or use Ulysses SP."
+                )
+            return _run_sdpa()
+
+        if query.dtype == torch.float32:
+            if self.attn_backend_explicit:
+                raise ValueError(
+                    f"{self.attn_backend_pref} was explicitly selected for ring attention, "
+                    "but its ring kernel does not support float32. Select TORCH_SDPA or use a supported dtype."
+                )
+            return _run_sdpa()
+
+        can_use_fa4 = _can_use_fa4(query.device)
+        can_use_fa3 = _can_use_fa3(query.device)
+        can_use_fa2 = _can_use_fa2(query.device)
+        if not can_use_fa4 and not can_use_fa3 and not can_use_fa2 and not HAS_AITER:
+            if self.attn_backend_explicit:
+                raise RuntimeError(
+                    f"{self.attn_backend_pref} was explicitly selected, but no compatible ring kernel "
+                    f"is available for device {query.device}."
+                )
+            logger.warning_once(
+                "Automatic ring backend selection chose TORCH_SDPA because no compatible "
+                "FA2/FA3/FA4/AITER ring kernel is available for this device."
+            )
+            return _run_sdpa()
+
         from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func
 
         # Prefer FA4 on Blackwell. An importable Hopper-only FA3 wheel can
         # otherwise be selected and fail at launch with "no kernel image".
-        if HAS_FA4 and query.is_cuda and torch.cuda.get_device_capability(query.device)[0] >= 10:
+        if can_use_fa4:
             attn_type = AttnType.FA4
         # Prefer FA3 over FA2 on Ampere/Ada/Hopper. On ROCm, use AITER.
-        elif HAS_FA3:
+        elif can_use_fa3:
             attn_type = AttnType.FA3
         elif HAS_AITER:
             attn_type = AttnType.AITER
-        else:
+        elif can_use_fa2:
             attn_type = AttnType.FA
+        else:
+            raise RuntimeError("No compatible Flash Attention backend is available for ring attention.")
 
         return ring_flash_attn_func(
             query,

vllm_omni/diffusion/attention/selector.py

@@ -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
 """
 Diffusion attention backend selector.
 
@@ -25,6 +25,10 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+# SequenceParallel auto-pad probes mask support before Attention layers exist,
+# so it has no real head_dim. Callers must not treat this sentinel as geometry.
+HEAD_SIZE_UNKNOWN = -1
+
 
 def _load_backend_cls(cls_path: str) -> type[AttentionBackend]:
     """Load a backend class from its fully qualified path.
@@ -151,3 +155,29 @@ def get_attn_backend_for_role(
         source="platform default",
     )
     return backend_cls, None
+
+
+def get_attn_backend_for_capability(
+    role: str,
+    attention_config: AttentionConfig | None = None,
+    role_category: str | None = None,
+) -> type[AttentionBackend]:
+    """Resolve the backend class for capability queries such as mask support.
+
+    SequenceParallel auto-pad does not know ``head_size``. Explicit
+    ``AttentionConfig`` selections are loaded from the registry without
+    geometry validation, so CUDNN_ATTN is not rejected for the unknown
+    sentinel. Platform defaults still consult the platform with
+    ``HEAD_SIZE_UNKNOWN``.
+    """
+    spec = None
+    if attention_config is not None:
+        spec, _ = attention_config.resolve_with_source(
+            role=role,
+            role_category=role_category,
+        )
+    if spec is not None:
+        from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
+
+        return DiffusionAttentionBackendEnum[spec.backend.upper()].get_class()
+    return _cached_get_backend_cls(None, HEAD_SIZE_UNKNOWN)

vllm_omni/diffusion/data.py

@@ -7,7 +7,7 @@ import os
 import random
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field, fields
-from enum import Enum, StrEnum
+from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 import diffusers
@@ -121,6 +121,8 @@ def validate_dlo_host_registration_options(
     hwr_mode: object,
 ) -> float:
     """Validate the optional transport budget without probing CUDA or HWR."""
+    if not isinstance(limit_gib, (int, float, str)):
+        raise TypeError(f"dlo_host_registration_limit_gib must be a number; got {type(limit_gib).__name__}")
     value = float(limit_gib)
     if not math.isfinite(value) or value < 0:
         raise ValueError("dlo_host_registration_limit_gib must be finite and >= 0")
@@ -1682,7 +1684,8 @@ class AttnQuantSpec:
 
     @property
     def enabled(self) -> bool:
-        return self.dtype_qk is not None or self.dtype_vo is not None
+        # Include flashinfer_backend so a variant pin without dtypes is serialized.
+        return self.dtype_qk is not None or self.dtype_vo is not None or self.flashinfer_backend is not None
 
 
 # Backends that select key blocks instead of attending densely, and so accept
@@ -1690,7 +1693,7 @@ class AttnQuantSpec:
 BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"})
 
 
-class RainFusionPrecision(StrEnum):
+class RainFusionPrecision(str, Enum):
     """Execution precision for block-sparse RainFusion (rf_v3) attention.
 
     ``bf16``: no quantization, pure BF16 sparse attention (official baseline).
@@ -1728,11 +1731,11 @@ class BlockSparseSpec:
             raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.")
         if self.end_step < 0:
             raise ValueError(f"block_sparse.end_step must be >= 0; got {self.end_step!r}.")
-        if self.precision not in {p.value for p in RainFusionPrecision}:
+        precision = self.precision.value if isinstance(self.precision, RainFusionPrecision) else self.precision
-            raise ValueError(
+        valid = {member.value for member in RainFusionPrecision}
-                f"block_sparse.precision must be one of {sorted(p.value for p in RainFusionPrecision)}; "
+        if precision not in valid:
-                f"got {self.precision!r}."
+            raise ValueError(f"block_sparse.precision must be one of {sorted(valid)}; got {self.precision!r}.")
-            )
+        self.precision = precision
         self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers)
 
 
@@ -1797,18 +1800,19 @@ class AttentionSpec:
                 kw["target_sparsity"] = ss.target_sparsity
             if ss.disabled_until_timestep:
                 kw["disabled_until_timestep"] = ss.disabled_until_timestep
-        if self.quant is not None and self.quant.enabled:
+        if self.quant is not None:
             q = self.quant
-            quant_kw: dict[str, Any] = {
+            quant_kw: dict[str, Any] = {}
-                "dtype_qk": q.dtype_qk,
+            if q.dtype_qk is not None or q.dtype_vo is not None:
-                "q_block_size": q.q_block_size,
+                quant_kw["dtype_qk"] = q.dtype_qk
-                "k_block_size": q.k_block_size,
+                quant_kw["q_block_size"] = q.q_block_size
-            }
+                quant_kw["k_block_size"] = q.k_block_size
-            if q.dtype_vo is not None:
+                if q.dtype_vo is not None:
-                quant_kw["dtype_vo"] = q.dtype_vo
+                    quant_kw["dtype_vo"] = q.dtype_vo
             if q.flashinfer_backend is not None:
                 quant_kw["flashinfer_backend"] = q.flashinfer_backend
-            kw["quant"] = quant_kw
+            if quant_kw:
+                kw["quant"] = quant_kw
         if self.fastvideo_vsa_topk is not None:
             kw["topk"] = self.fastvideo_vsa_topk
         if self.block_sparse is not None:

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project and the HuggingFace Team.
 # All rights reserved.
 #
@@ -35,6 +36,7 @@ based on the plan specification.
 from __future__ import annotations
 
 import inspect
+from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Any
 
@@ -46,8 +48,10 @@ from vllm_omni.diffusion.distributed.sp_plan import (
     AnySequenceParallelInput,
     SequenceParallelConfig,
     SequenceParallelInput,
+    SequenceParallelInputType,
     SequenceParallelModelPlan,
     SequenceParallelOutput,
+    SequenceParallelOutputType,
     SequenceParallelPartialInput,
 )
 from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard
@@ -69,18 +73,18 @@ class ModuleForwardMetadata:
     """
 
     cached_parameter_indices: dict[str, int] | None = None
-    _cls: type | None = None
+    _cls: type[nn.Module] | None = None
 
     def _get_parameter_from_args_kwargs(
         self,
-        identifier: str,
+        identifier: str | int,
-        args: tuple = (),
+        args: Sequence[Any] = (),
-        kwargs: dict | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> tuple[Any, bool, int | None]:
         """Get a parameter value from args or kwargs by name.
 
         Args:
-            identifier: The parameter name to look up.
+            identifier: The parameter name to look up, or a positional index.
             args: Positional arguments passed to forward.
             kwargs: Keyword arguments passed to forward.
 
@@ -95,6 +99,11 @@ class ModuleForwardMetadata:
         """
         kwargs = kwargs or {}
 
+        if isinstance(identifier, int):
+            if identifier < len(args):
+                return args[identifier], False, identifier
+            return None, False, identifier
+
         # First check kwargs
         if identifier in kwargs:
             return kwargs[identifier], True, None
@@ -164,7 +173,7 @@ class SequenceParallelSplitHook(ModelHook):
 
     def __init__(
         self,
-        metadata: dict[str | int, AnySequenceParallelInput | list[AnySequenceParallelInput]],
+        metadata: SequenceParallelInputType,
         config: SequenceParallelConfig,
     ) -> None:
         super().__init__()
@@ -191,6 +200,8 @@ class SequenceParallelSplitHook(ModelHook):
                 continue
 
             # Get the parameter value
+            if self.module_forward_metadata is None:
+                raise RuntimeError("SequenceParallelSplitHook.initialize_hook must run before forward")
             input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs(
                 name, args_list, kwargs
             )
@@ -200,6 +211,10 @@ class SequenceParallelSplitHook(ModelHook):
 
             # Shard the input
             if isinstance(input_val, torch.Tensor):
+                if not isinstance(spm, (SequenceParallelInput, SequenceParallelPartialInput)):
+                    raise ValueError(
+                        f"Expected SequenceParallelInput for tensor parameter '{name}', got {type(spm).__name__}"
+                    )
                 input_val = self._prepare_sp_input(input_val, spm, args_list, kwargs)
             elif isinstance(input_val, (list, tuple)):
                 # Handle list/tuple of tensors with per-element config
@@ -221,6 +236,8 @@ class SequenceParallelSplitHook(ModelHook):
 
             # Update args or kwargs
             if is_kwarg:
+                if not isinstance(name, str):
+                    raise TypeError(f"Keyword parameter name must be str, got {type(name).__name__}")
                 kwargs[name] = input_val
             elif index is not None and index < len(args_list):
                 args_list[index] = input_val
@@ -270,8 +287,8 @@ class SequenceParallelSplitHook(ModelHook):
     def _resolve_text_len(
         self,
         sp_input: SequenceParallelPartialInput,
-        args: tuple,
+        args: Sequence[Any],
-        kwargs: dict,
+        kwargs: dict[str, Any],
     ) -> int:
         """Resolve text length from the source specification."""
         source = sp_input.text_len_source
@@ -285,6 +302,8 @@ class SequenceParallelSplitHook(ModelHook):
 
         # Try to get from kwargs/args
         try:
+            if self.module_forward_metadata is None:
+                raise RuntimeError("SequenceParallelSplitHook.initialize_hook must run before forward")
             val, _, _ = self.module_forward_metadata._get_parameter_from_args_kwargs(source, args, kwargs)
             if val is None:
                 raise ValueError(f"Parameter '{source}' is None, cannot determine text length.")
@@ -306,8 +325,8 @@ class SequenceParallelSplitHook(ModelHook):
         self,
         x: torch.Tensor,
         sp_input: AnySequenceParallelInput,
-        args: tuple = (),
+        args: Sequence[Any] = (),
-        kwargs: dict | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> torch.Tensor:
         """Shard a tensor according to the input specification."""
         kwargs = kwargs or {}
@@ -382,7 +401,7 @@ class SequenceParallelSplitHook(ModelHook):
         2. Creates an attention mask indicating valid vs padding positions
         3. Stores the mask and padding info in ForwardContext
         """
-        from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role
+        from vllm_omni.diffusion.attention.selector import get_attn_backend_for_capability
         from vllm_omni.diffusion.distributed.parallel_state import (
             get_ring_parallel_world_size,
             get_sequence_parallel_rank,
@@ -408,12 +427,14 @@ class SequenceParallelSplitHook(ModelHook):
             if od_config is not None:
                 attention_config = od_config.diffusion_attention_config
 
-        attn_backend, _ = get_attn_backend_for_role(
+        attn_backend = get_attn_backend_for_capability(
             role="self",
-            head_size=-1,
             attention_config=attention_config,
         )
-        if not attn_backend.supports_attention_mask:
+        attention_spec = None
+        if attention_config is not None:
+            attention_spec, _ = attention_config.resolve_with_source(role="self")
+        if not attn_backend.supports_attention_mask(attention_spec):
             raise ValueError(
                 f"Sequence length ({seq_len}) is not divisible by SP world size ({world_size}). "
                 f"Cannot use {attn_backend.get_name()} which does not support attention_mask. "
@@ -468,12 +489,14 @@ class SequenceParallelGatherHook(ModelHook):
 
     def __init__(
         self,
-        metadata: SequenceParallelOutput | list[SequenceParallelOutput],
+        metadata: SequenceParallelOutputType,
         config: SequenceParallelConfig,
     ) -> None:
         super().__init__()
         if isinstance(metadata, SequenceParallelOutput):
             metadata = [metadata]
+        elif isinstance(metadata, tuple):
+            metadata = list(metadata)
         self.metadata = metadata
         self.config = config
 
@@ -630,6 +653,7 @@ def apply_sequence_parallel(
         logger.debug(f"Applying SP hooks to '{module_id}' ({len(submodule)} module(s))")
 
         for m in submodule:
+            hook: SequenceParallelSplitHook | SequenceParallelGatherHook
             if isinstance(sp_model_plan, dict):
                 # Input specification
                 hook = SequenceParallelSplitHook(sp_model_plan, config)
@@ -638,6 +662,8 @@ def apply_sequence_parallel(
                 # Output specification
                 if isinstance(sp_model_plan, SequenceParallelOutput):
                     sp_model_plan = [sp_model_plan]
+                elif isinstance(sp_model_plan, tuple):
+                    sp_model_plan = list(sp_model_plan)
                 if not all(isinstance(x, SequenceParallelOutput) or x is None for x in sp_model_plan):
                     raise ValueError(f"Expected SequenceParallelOutput elements, got {sp_model_plan}")
                 hook = SequenceParallelGatherHook(sp_model_plan, config)

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -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
 """Diffusers backend adapter for vLLM-Omni.
 
 Provides a black-box wrapper around any 🤗 Diffusers pipeline, enabling
@@ -44,13 +44,21 @@ logger = logging.getLogger(__name__)
 # with the accuracy tests' diffusers reference runs so both sides of a
 # similarity comparison resolve to the same backend on any given image
 # (hub kernels may lack a build variant for the image's torch).
-CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS: list[str] = [
+# Automatic / FLASH_ATTN selection walks the full chain. Explicit Hub
+# options attempt only their matching Diffusers backend names.
+CUDA_FLASH_ATTN_3_HUB_ATTEMPTS: list[str] = [
     "_flash_3_hub",
     "_flash_3_varlen_hub",
-    "_flash_3",
+]
-    "_flash_varlen_3",
+CUDA_FLASH_ATTN_HUB_ATTEMPTS: list[str] = [
     "flash_hub",
     "flash_varlen_hub",
+]
+CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS: list[str] = [
+    *CUDA_FLASH_ATTN_3_HUB_ATTEMPTS,
+    "_flash_3",
+    "_flash_varlen_3",
+    *CUDA_FLASH_ATTN_HUB_ATTEMPTS,
     "flash",
     "flash_varlen",
     "_native_flash",
@@ -293,6 +301,24 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
     # Wrap settings, inputs, and outputs
     # ------------------------------------------------------------------
 
+    @staticmethod
+    def _is_blackwell_gpu() -> bool:
+        """Return True on Blackwell-class CUDA GPUs (sm_10x / sm_11x / sm_12x)."""
+        capability = None
+        try:
+            capability = current_omni_platform.get_device_capability()
+        except Exception:
+            capability = None
+        if capability is not None:
+            return capability.major in (10, 11, 12)
+        if torch.cuda.is_available():
+            try:
+                major, _ = torch.cuda.get_device_capability()
+                return major in (10, 11, 12)
+            except Exception:
+                return False
+        return False
+
     def _set_attention_backend(self) -> None:
         """Set the attention backend.
 
@@ -308,6 +334,20 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
         attention_backend_config = default_spec.backend if default_spec is not None else None
         attention_backend_attempts: list[str] = []
         match attention_backend_config:
+            case "FLASH_ATTN_3_HUB":
+                if self._is_blackwell_gpu():
+                    raise ValueError(
+                        "FLASH_ATTN_3_HUB was explicitly selected for a Diffusers adapter on "
+                        "Blackwell, but no runnable sm_10x+ Flash Attention kernel is available."
+                    )
+                attention_backend_attempts.extend(CUDA_FLASH_ATTN_3_HUB_ATTEMPTS)
+            case "FLASH_ATTN_HUB":
+                if self._is_blackwell_gpu():
+                    raise ValueError(
+                        "FLASH_ATTN_HUB was explicitly selected for a Diffusers adapter on "
+                        "Blackwell, but no runnable sm_10x+ Flash Attention kernel is available."
+                    )
+                attention_backend_attempts.extend(CUDA_FLASH_ATTN_HUB_ATTEMPTS)
             case "FLASH_ATTN" | None:
                 if current_omni_platform.is_rocm():
                     attention_backend_attempts.append("aiter")
@@ -318,17 +358,31 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                         "Unknown diffusers attention backend option for MUSA platform. Falling back to SDPA."
                     )
                     attention_backend_attempts.append("native")
+                elif self._is_blackwell_gpu():
+                    # Installed FA2/FA3 (and HF FA3 hub) wheels are typically Hopper-only.
+                    # set_attention_backend("_flash_3*") can succeed on import, then crash
+                    # at runtime with "no kernel image is available for execution on the device".
+                    if attention_backend_config is not None:
+                        raise ValueError(
+                            f"{attention_backend_config} was explicitly selected for a Diffusers adapter on "
+                            "Blackwell, but no runnable sm_10x+ Flash Attention kernel is available."
+                        )
+                    attention_backend_attempts.append("native")
                 else:
                     attention_backend_attempts.extend(CUDA_FLASH_ATTENTION_BACKEND_ATTEMPTS)
             case "SAGE_ATTN":
-                attention_backend_attempts.extend(["sage_hub", "sage", "sage", "sage_varlen"])
+                attention_backend_attempts.extend(["sage_hub", "sage", "sage_varlen"])
             case "ASCEND":
                 attention_backend_attempts.append("_native_npu")
             case "TORCH_SDPA":
                 attention_backend_attempts.append("native")
+            case "CUDNN_ATTN" | "FLASHINFER_ATTN" | "TRTLLM_ATTN":
+                raise ValueError(
+                    f"{attention_backend_config} was explicitly selected, but Diffusers adapters do not "
+                    "provide a binding for this vLLM-Omni backend. Select TORCH_SDPA or a Diffusers-native backend."
+                )
             case _:
-                logger.warning(f"Invalid attention backend: {attention_backend_config}. Falling back to SDPA.")
+                raise ValueError(f"Invalid Diffusers attention backend: {attention_backend_config}.")
-                attention_backend_attempts.append("native")
 
         attempt_errors: list[str] = []
         set_backend: str | None = None
@@ -340,8 +394,15 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             except Exception as e:
                 attempt_errors.append(str(e))
 
-        # If all attempts fail, fallback to SDPA and warn the user about the failures
+        # Automatic selection may use native SDPA as its final safe choice. An
+        # explicitly configured backend must either be installed or raise.
         if len(attempt_errors) == len(attention_backend_attempts):
+            if attention_backend_config is not None:
+                raise RuntimeError(
+                    f"Failed to set explicitly selected attention backend '{attention_backend_config}' for "
+                    f"diffusers pipeline {self._pipeline.__class__.__name__}. Attempts: "
+                    f"{dict(zip(attention_backend_attempts, attempt_errors))}"
+                )
             self._pipeline.transformer.set_attention_backend("native")
             logger.warning(
                 f"Failed to set attention backend '{attention_backend_config}' for "

vllm_omni/platforms/cuda/platform.py

@@ -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
 
@@ -52,6 +52,21 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
         return True
 
+    @classmethod
+    def has_flash_attn_4(cls) -> bool:
+        """Return whether CuTe FA4 is importable (Blackwell-capable FLASH_ATTN)."""
+        from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_4_available
+
+        return is_flash_attn_4_available()
+
+    @classmethod
+    def supports_diffusion_dense_flash_attention(cls) -> bool:
+        """Dense FLASH_ATTN is FA4 on Blackwell; FA2/FA3 on older CUDA GPUs."""
+        capability = cls.get_device_capability()
+        if capability is not None and capability.major in (10, 11, 12):
+            return cls.has_flash_attn_4()
+        return cls.has_flash_attn_package()
+
     @classmethod
     def get_diffusion_attn_backend_cls(
         cls,
@@ -84,7 +99,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
             is_blackwell = major in (10, 11, 12)
             if is_blackwell and (major, minor) not in _known_blackwell_sms:
                 logger.info(
-                    "Detected Blackwell-class GPU %s (untested variant); routing to CUDNN_ATTN with SDPA fallback.",
+                    "Detected Blackwell-class GPU %s (untested variant); routing to CUDNN_ATTN.",
                     sm_str,
                 )
 
@@ -92,19 +107,24 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         packages_info = PACKAGES_CHECKER.get_packages_info()
         packages_available = packages_info.get("has_flash_attn", False)
 
-        # Both compute capability and packages must be available for FA
+        # FA2/FA3 wheels are often importable on Blackwell but only ship
-        flash_attn_supported = compute_supported and packages_available
+        # Ampere/Ada/Hopper kernels; the first forward then dies with
+        # "no kernel image". FA4 (flash_attn.cute / vllm-omni[fa4]) is the
+        # FLASH_ATTN implementation that actually runs on Blackwell.
+        if is_blackwell:
+            flash_attn_supported = compute_supported and cls.has_flash_attn_4()
+        else:
+            flash_attn_supported = compute_supported and packages_available
 
         # cuDNN 9.5+ ships Blackwell FMHA kernels. If the runtime is older,
-        # the CUDNN_ATTN default would still work via internal fallback but
+        # skip the CUDNN_ATTN default rather than selecting a backend whose
-        # without the tuned Blackwell path, so we skip routing there.
+        # kernel selector may reject the runtime shape.
         cudnn_version = torch.backends.cudnn.version() or 0
         cudnn_blackwell_ready = cudnn_version >= 90500
 
         # FlashInfer edges cuDNN by ~4% at the kernel level on sm_120 but
-        # regresses ~2x at e2e on HV-1.5 because its dense-prefill path can't
+        # regresses ~2x at e2e on HV-1.5 because its dense-prefill path cannot
-        # take 2D attn_masks and the SDPA fallback dispatches to
+        # represent every diffusion attention mask without changing semantics.
-        # EFFICIENT_ATTENTION (~25 ms) instead of the cuDNN mask path (~11 ms).
         # CUDNN_ATTN pins sdpa_kernel([CUDNN_ATTENTION]) directly so masked
         # calls keep the cuDNN path. Blackwell default prefers CUDNN_ATTN;
         # users can opt into FLASHINFER_ATTN explicitly for no-mask workloads.
@@ -120,6 +140,18 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
             # not abort startup — just treat FlashInfer as unavailable.
             logger.debug("FlashInfer probe failed (%s); treating as unavailable", e)
 
+        # FLASHINFER_ATTN needs BatchPrefillWithRaggedKVCacheWrapper, not just a
+        # top-level flashinfer package. Probe the actual symbol so a stub /
+        # partial wheel is not auto-routed then crash at layer init.
+        flashinfer_prefill_available = False
+        if flashinfer_available:
+            try:
+                from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper  # noqa: F401
+
+                flashinfer_prefill_available = True
+            except Exception as e:
+                logger.debug("FlashInfer prefill wrapper probe failed (%s); treating as unavailable", e)
+
         # TRTLLM_ATTN needs the trtllm-gen kernel specifically, not just any FlashInfer
         # wheel. Probe the actual symbol so a released wheel lacking it does not get
         # auto-routed to TRTLLM_ATTN and then crash on the first forward.
@@ -134,58 +166,117 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
 
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
+            # Architecture / package probes for optional kernels must run before
+            # validate_diffusion_attn_backend, which imports the backend module
+            # and would otherwise load sageattention / sageattn3 on unsupported GPUs.
+            if backend_upper == "SAGE_ATTN":
+                sage_supported_sms = {(8, 0), (8, 6), (8, 9), (9, 0), (12, 0), (12, 1)}
+                if compute_capability is None or tuple(compute_capability) not in sage_supported_sms:
+                    raise ValueError(
+                        f"SAGE_ATTN was explicitly selected but does not provide a kernel for {sm_str or 'this GPU'}. "
+                        "Select a compatible backend."
+                    )
+                try:
+                    importlib.import_module("sageattention")
+                except ImportError as e:
+                    raise ImportError(
+                        "SAGE_ATTN was explicitly selected, but the sageattention package is not available."
+                    ) from e
+            if backend_upper == "SAGE_ATTN_3":
+                sage_attn3_supported = compute_capability is not None and compute_capability.major >= 10
+                if not sage_attn3_supported:
+                    raise ValueError(
+                        "SAGE_ATTN_3 was explicitly selected but requires a Blackwell-class GPU "
+                        "with compute capability >= 10.0. Select a compatible backend."
+                    )
+                try:
+                    importlib.import_module("sageattn3")
+                except ImportError as e:
+                    raise ImportError(
+                        "SAGE_ATTN_3 was explicitly selected, but the sageattn3 package is not available. "
+                        "Install SageAttention/sageattention3_blackwell or select a different backend."
+                    ) from e
             cls.validate_diffusion_attn_backend(backend_upper)
+            if backend_upper == "SAGE_ATTN_3" and head_size > 0:
+                sage3_cls = DiffusionAttentionBackendEnum.SAGE_ATTN_3.get_class()
+                if not sage3_cls.supports_head_size(head_size):
+                    raise ValueError(
+                        f"SAGE_ATTN_3 was explicitly selected but head_size={head_size} is unsupported. "
+                        f"Supported head sizes: {sage3_cls.get_supported_head_sizes()}. "
+                        "Select TORCH_SDPA or another backend."
+                    )
             if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
                 try:
                     importlib.import_module("kernels")
                     logger.info("Using HuggingFace kernels-backed attention backend '%s'", backend_upper)
-                except ImportError:
+                except ImportError as e:
-                    if backend_upper == "FLASH_ATTN_HUB":
+                    raise ImportError(
-                        logger.warning(
+                        f"{backend_upper} was explicitly selected, but the HuggingFace `kernels` "
-                            "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN."
+                        "library is not available. Install `kernels` or select a different backend."
-                        )
+                    ) from e
-                        backend_upper = "FLASH_ATTN"
+
-                    elif backend_upper == "FLASH_ATTN_3_HUB":
+            if backend_upper == "FLASH_ATTN_HUB":
-                        logger.warning(
+                fa2_hub_supported = compute_capability is not None and compute_capability.major in (8, 9)
-                            "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN."
+                if not fa2_hub_supported:
-                        )
+                    raise ValueError(
-                        backend_upper = "FLASH_ATTN"
+                        "FLASH_ATTN_HUB was explicitly selected but its current FA2 kernels require "
+                        "an Ampere, Ada, or Hopper GPU (compute capability 8.x/9.x)."
+                    )
 
             if backend_upper == "FLASH_ATTN_3_HUB":
-                fa3_hub_supported = compute_capability is not None and compute_capability.major >= 9
+                fa3_hub_supported = compute_capability is not None and compute_capability.major == 9
                 if not fa3_hub_supported:
-                    logger.warning(
+                    raise ValueError(
-                        "FLASH_ATTN_3_HUB requires a Hopper-class GPU with compute capability >= 9.0. "
+                        "FLASH_ATTN_3_HUB was explicitly selected but its current kernels require a Hopper GPU "
-                        "Falling back to FLASH_ATTN_HUB."
+                        "with compute capability 9.x. Select a compatible backend."
                     )
-                    backend_upper = "FLASH_ATTN_HUB"
 
             if backend_upper == "FLASH_ATTN" and not flash_attn_supported:
-                if not compute_supported:
+                if is_blackwell:
-                    logger.warning(
+                    reason = (
-                        "Flash Attention requires GPU with compute capability >= 8.0. "
+                        "Blackwell requires CuTe FlashAttention-4 "
-                        "Falling back to TORCH_SDPA backend."
+                        "(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only"
                     )
-                elif not packages_available:
+                elif not compute_supported:
-                    logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.")
+                    reason = "compute capability < 8.0"
-                logger.debug("Defaulting to diffusion attention backend SDPA")
+                else:
-                return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+                    reason = "Flash Attention package unavailable"
-            if backend_upper == "SAGE_ATTN_3":
+                raise ValueError(
-                sage_attn3_supported = compute_capability is not None and compute_capability.major >= 10
+                    f"FLASH_ATTN was explicitly selected but is unsupported ({reason}). Select a compatible backend."
-                if not sage_attn3_supported:
+                )
-                    logger.warning(
+            if backend_upper == "FLASHINFER_ATTN" and not flashinfer_available:
-                        "SageAttention3 requires a Blackwell-class GPU with compute capability >= 10.0. "
+                raise ValueError(
-                        "Falling back to TORCH_SDPA backend."
+                    "FLASHINFER_ATTN was explicitly selected, but FlashInfer is unavailable. "
+                    "Install a compatible FlashInfer build or select a different backend."
+                )
+            if backend_upper == "FLASHINFER_ATTN" and not flashinfer_prefill_available:
+                raise ValueError(
+                    "FLASHINFER_ATTN was explicitly selected, but the installed FlashInfer build does not "
+                    "provide BatchPrefillWithRaggedKVCacheWrapper. Install a compatible build."
+                )
+            if backend_upper == "FLASHINFER_ATTN":
+                from vllm_omni.diffusion.attention.backends.flashinfer_attn import FlashInferAttentionBackend
+
+                # head_size <= 0 is the capability-probe sentinel. Skip geometry
+                # validation; Attention construction still checks the real size.
+                if head_size > 0 and not FlashInferAttentionBackend.supports_head_size(head_size):
+                    raise ValueError(
+                        f"FLASHINFER_ATTN was explicitly selected but head_size={head_size} is unsupported. "
+                        f"Supported head sizes: {FlashInferAttentionBackend.get_supported_head_sizes()}. "
+                        "Select TORCH_SDPA or another backend."
                     )
-                    return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
+            if backend_upper == "CUDNN_ATTN":
-                try:
+                from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend
-                    importlib.import_module("sageattn3")
+
-                except ImportError:
+                # head_size <= 0 is the capability-probe sentinel (auto-pad does
-                    logger.warning(
+                # not know head_dim). Skip geometry validation; the real
-                        "SageAttention3 package not available. Install it from "
+                # Attention construction still checks the configured size.
-                        "SageAttention/sageattention3_blackwell. Falling back to TORCH_SDPA backend."
+                if head_size > 0 and not CuDNNAttentionBackend.supports_head_size(head_size):
+                    raise ValueError(
+                        f"CUDNN_ATTN was explicitly selected but head_size={head_size} is unsupported. "
+                        "Blackwell cuDNN FMHA requires head_dim divisible by 8 and no larger than 256. "
+                        "Select FLASHINFER_ATTN or TORCH_SDPA."
                     )
-                    return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
             if backend_upper == "TRTLLM_ATTN":
                 trtllm_attn_supported = compute_capability is not None and compute_capability.major == 10
                 if not trtllm_attn_supported:
@@ -194,9 +285,10 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
                         "Blackwell GPU (SM100 / SM103, compute capability 10.x). Select a "
                         "different --diffusion-attention-backend."
                     )
-                if not flashinfer_available:
+                if not trtllm_gen_available:
                     raise ValueError(
-                        "TRTLLM_ATTN diffusion attention backend requires flashinfer, which is not installed."
+                        "TRTLLM_ATTN was explicitly selected, but the installed FlashInfer build does not "
+                        "provide trtllm_ragged_attention_deepseek. Install a compatible build."
                     )
             backend = DiffusionAttentionBackendEnum[backend_upper]
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
@@ -218,24 +310,40 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
             return DiffusionAttentionBackendEnum.TRTLLM_ATTN.get_path()
 
         if is_blackwell and cudnn_blackwell_ready:
-            logger.info(
+            from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend
-                "Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell %s, cuDNN %d)",
-                sm_str,
-                cudnn_version,
-            )
-            return DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
 
-        if is_blackwell and flashinfer_available:
+            # Unknown head_size is a capability probe (SP auto-pad), not a
+            # real DiT geometry. Keep CUDNN so mask support is reported.
+            if head_size <= 0 or CuDNNAttentionBackend.supports_head_size(head_size):
+                logger.info(
+                    "Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell %s, cuDNN %d, head_dim %d)",
+                    sm_str,
+                    cudnn_version,
+                    head_size,
+                )
+                return DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
             logger.info(
-                "Defaulting to diffusion attention backend FLASHINFER_ATTN (Blackwell %s, cuDNN unavailable)",
+                "Skipping CUDNN_ATTN on Blackwell %s: head_dim %d is outside cuDNN FMHA "
+                "(multiples of 8, 8-256). Automatic selection will use FlashInfer when available, otherwise SDPA.",
                 sm_str,
+                head_size,
             )
-            return DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
+
+        if is_blackwell and flashinfer_prefill_available:
+            from vllm_omni.diffusion.attention.backends.flashinfer_attn import FlashInferAttentionBackend
+
+            if FlashInferAttentionBackend.supports_head_size(head_size):
+                logger.info(
+                    "Defaulting to diffusion attention backend FLASHINFER_ATTN (Blackwell %s, head_dim %d)",
+                    sm_str,
+                    head_size,
+                )
+                return DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
 
         if is_blackwell and not cudnn_blackwell_ready:
             logger.warning(
                 "Detected Blackwell %s but cuDNN %d < 9.5 — no tuned Blackwell FMHA. "
-                "Falling through to FLASH_ATTN / SDPA.",
+                "Automatic backend selection will use FlashInfer when available, otherwise SDPA.",
                 sm_str,
                 cudnn_version,
             )