From 95a93116960e364cc759c49023a4c35284157235 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:39:49 +0800 Subject: [PATCH] [CI] Select mirror_hardwares presets via MIRROR_HW (H100/B200) (#5543) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: lishunyang12 Signed-off-by: [Your Name] 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 Co-authored-by: Cursor Co-authored-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .buildkite/common/ci_mirror_hardwares.yml | 221 ++++++++++++- .buildkite/common/scripts/skip_ci.py | 39 +-- .buildkite/common/scripts/upload_pipeline.py | 300 +++++++++++++++-- .buildkite/cuda/test-merge.yml | 3 + .buildkite/cuda/test-nightly.yml | 13 +- .buildkite/cuda/test-ready.yml | 3 + .buildkite/cuda/test-weekly.yml | 3 + docs/contributing/ci/ci_settings.md | 26 +- docs/contributing/ci/test_system_overview.md | 11 +- .../attention_backends/dense_backends.md | 22 +- .../attention_backends/huggingface_hub.md | 6 +- .../fault_injection_reliability_matrix.md | 18 +- tests/buildkite/test_skip_ci.py | 1 + tests/buildkite/test_upload_pipeline.py | 154 ++++++++- .../perf/scripts/run_diffusion_benchmark.py | 5 +- .../perf/tests/test_minimax_h3_vllm_omni.json | 4 - tests/dfx/perf/tests/test_runner_metadata.py | 1 + .../attention/test_attention_config.py | 313 +++++++++++++++++- .../attention/test_cuda_backend_selection.py | 265 +++++++++++++++ tests/diffusion/attention/test_cudnn_attn.py | 110 +++--- .../attention/test_flashinfer_attn.py | 112 +++++++ tests/diffusion/attention/test_kernels_hub.py | 83 ++++- .../attention/test_ring_backend_selection.py | 199 +++++++++++ tests/diffusion/attention/test_ring_fa4.py | 22 ++ tests/diffusion/attention/test_ring_utils.py | 92 +++++ tests/diffusion/attention/test_sage_attn.py | 36 ++ tests/diffusion/attention/test_sage_attn3.py | 119 +++++-- tests/diffusion/attention/test_selector.py | 44 ++- tests/diffusion/cache/test_cache_dit.py | 49 ++- .../test_diffusers_backend.py | 80 +++++ .../distributed/test_sp_plan_hooks.py | 83 +++++ .../test_boogu_image_transformer.py | 16 +- .../models/ltx2/test_ltx2_transformer.py | 5 +- tests/e2e/accuracy/helpers.py | 56 +++- .../test_hunyuan_image3_pixel_accuracy.py | 50 ++- .../test_wan22_i2v_video_similarity.py | 265 +++------------ .../wan22_i2v_video_similarity_common.py | 14 +- .../e2e/offline_inference/test_gepard_tts.py | 2 +- .../offline_inference/test_minicpmo_4_5.py | 13 +- ...st_qwen3_omni_autoround_w4a16_expansion.py | 12 +- .../test_qwen3_omni_modelopt_nvfp4_w4a4.py | 2 +- tests/e2e/online_serving/test_minicpmo_4_5.py | 8 +- .../test_minicpmo_4_5_expansion.py | 10 +- .../test_nemotron_voicechat_duplex.py | 2 +- .../test_qwen3_omni_expansion.py | 24 +- tests/helpers/__init__.py | 6 +- tests/helpers/mark.py | 63 ++-- tests/helpers/tests/test_mark.py | 25 ++ .../cosyvoice3/test_cosyvoice3_components.py | 83 +++-- .../nightly/buildkite_testcase_statistics.py | 3 + tools/nightly/run_nightly_jobs.sh | 11 +- tools/pre_commit/shellcheck.sh | 6 +- tools/run_jobs_common.sh | 4 +- tools/run_merge_jobs.sh | 2 +- tools/run_ready_jobs.sh | 2 +- .../diffusion/attention/backends/abstract.py | 10 +- .../attention/backends/cudnn_attn.py | 90 +++-- .../attention/backends/flash_attn.py | 29 +- .../attention/backends/flash_attn_hub.py | 48 +-- .../attention/backends/flashinfer_attn.py | 140 +++++--- .../attention/backends/ring/ring_globals.py | 30 +- .../attention/backends/ring/ring_kernels.py | 13 +- .../attention/backends/ring/ring_utils.py | 174 ++++------ .../attention/backends/ring_flash_attn.py | 3 +- .../attention/backends/ring_pytorch_attn.py | 2 +- .../diffusion/attention/backends/sage_attn.py | 9 +- .../attention/backends/sage_attn3.py | 31 +- .../diffusion/attention/backends/sdpa.py | 5 +- .../diffusion/attention/backends/utils/fa.py | 162 +++++---- vllm_omni/diffusion/attention/layer.py | 101 +++--- .../diffusion/attention/parallel/factory.py | 34 +- .../diffusion/attention/parallel/ring.py | 116 ++++++- vllm_omni/diffusion/attention/selector.py | 32 +- vllm_omni/diffusion/data.py | 38 ++- .../diffusion/hooks/sequence_parallel.py | 56 +++- .../pipeline_diffusers_adapter.py | 77 ++++- vllm_omni/platforms/cuda/platform.py | 226 +++++++++---- 77 files changed, 3470 insertions(+), 1077 deletions(-) create mode 100644 tests/diffusion/attention/test_cuda_backend_selection.py create mode 100644 tests/diffusion/attention/test_flashinfer_attn.py create mode 100644 tests/diffusion/attention/test_ring_backend_selection.py create mode 100644 tests/diffusion/attention/test_ring_utils.py create mode 100644 tests/diffusion/attention/test_sage_attn.py diff --git a/.buildkite/common/ci_mirror_hardwares.yml b/.buildkite/common/ci_mirror_hardwares.yml index 2d4aea7d4..e3c9a427a 100644 --- a/.buildkite/common/ci_mirror_hardwares.yml +++ b/.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" diff --git a/.buildkite/common/scripts/skip_ci.py b/.buildkite/common/scripts/skip_ci.py index b452b5adf..8d5fdd5e0 100644 --- a/.buildkite/common/scripts/skip_ci.py +++ b/.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 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 - 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 diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index 63e8780c8..53da6f660 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.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] @@ -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]: - """Replace uploader-only ``mirror_hardwares`` with preset fields from ci_mirror_hardwares.yml.""" - hardware = step.get("mirror_hardwares") - if hardware is None: - return step +@lru_cache(maxsize=1) +def _cuda_mirror_chips() -> tuple[str, ...]: + """Lowercase CUDA SKUs that have at least one ``{chip}_{n}`` preset. - if not isinstance(hardware, str) or not hardware.strip(): + Longer tokens first so ``h100`` is not matched as a prefix of ``h200``. + """ + skus = get_skus_for_platform("cuda") + known = _load_mirror_hardwares() + chips = { + sku.lower() for sku in skus if any(name == sku.lower() or name.startswith(f"{sku.lower()}_") for name in known) + } + return tuple(sorted(chips, key=lambda chip: (-len(chip), chip))) + + +def _read_cards_marks(expr: str) -> tuple[set[int], bool]: + """Parse registered ``cards_n`` / ``not cards_n`` from a pytest ``-m`` expr. + + Returns positive card counts and whether any ``not cards_*`` is present. + """ + counts = get_supported_card_counts() + alt = "|".join(str(n) for n in sorted(counts, key=lambda n: (-len(str(n)), -n))) if counts else "(?!)" + positives: set[int] = set() + has_not_cards = False + for match in re.finditer(rf"\b(?:(not)\s+)?cards_({alt})\b", expr): + if match.group(1): + has_not_cards = True + else: + positives.add(int(match.group(2))) + return positives, has_not_cards + + +def _read_hardware_marks(expr: str) -> set[str]: + """Return lowercase positive CUDA SKUs that have a mirror preset. + + ``not H100`` is ignored. SKUs with no ``{chip}_*`` preset (e.g. H200) are dropped. + """ + skus = sorted(get_skus_for_platform("cuda"), key=lambda s: (-len(s), s)) + if not skus: + return set() + alt = "|".join(re.escape(sku) for sku in skus) + mirror = set(_cuda_mirror_chips()) + chips: set[str] = set() + for match in re.finditer(rf"\b(?:(not)\s+)?({alt})\b", expr): + if match.group(1): + continue + chip = match.group(2).lower() + if chip in mirror: + chips.add(chip) + return chips + + +# CI policy: not pytest-mark registry. +_SUPPORTED_MIRROR_HW_SELECTORS = frozenset({"b200"}) +# Unset MIRROR_HW: match H100 then L4 in -m (skip if neither). +_DEFAULT_INFER_CHIPS = ("h100", "l4") +_PYTEST_MARKER_ARG = re.compile(r"-m\s+(?:\"([^\"]*)\"|'([^']*)'|(\S+))") +_CARDS_CHIP_MAX: Literal["max"] = "max" + + +def _get_mirror_hw_selector() -> str: + """Return lowercase ``MIRROR_HW``, or empty to keep string presets / marker chips. + + A set value must be ``b200`` (case-insensitive). Unknown values fail closed + so a typo cannot silently drop the CUDA pipeline. + """ + selector = os.environ.get("MIRROR_HW", "").strip().lower() + if not selector: + return "" + if selector not in _SUPPORTED_MIRROR_HW_SELECTORS: + allowed = " or ".join(repr(s) for s in sorted(_SUPPORTED_MIRROR_HW_SELECTORS)) + raise ValueError(f"unsupported MIRROR_HW={selector!r}; expected empty or {allowed}") + return selector + + +def _read_pytest_marks(commands: Any) -> tuple[set[str], int | Literal["max"] | None]: + """Read pytest ``-m`` from step commands: CUDA SKU chips and cards count. + + *chips* are lowercase SKUs from positive markers (``not H100`` is ignored). + *cards* is the max positive ``cards_n``, ``"max"`` when only ``not cards_*`` + is present, or ``None``. + """ + chunks: list[str] = [] + + def _collect(value: Any) -> None: + if value is None: + return + if isinstance(value, str): + chunks.append(value) + elif isinstance(value, list): + for part in value: + _collect(part) + else: + chunks.append(str(value)) + + _collect(commands) + text = "\n".join(chunks) + + chips: set[str] = set() + positives: set[int] = set() + has_not_cards = False + for match in _PYTEST_MARKER_ARG.finditer(text): + expr = match.group(1) or match.group(2) or match.group(3) or "" + chips |= _read_hardware_marks(expr) + card_counts, not_cards = _read_cards_marks(expr) + positives.update(card_counts) + has_not_cards = has_not_cards or not_cards + + if positives: + cards: int | Literal["max"] | None = max(positives) + elif has_not_cards: + cards = _CARDS_CHIP_MAX + else: + cards = None + return chips, cards + + +def _compose_mirror_hardware_name( + chips: set[str], + cards: int | Literal["max"] | None, + *, + step_label: str, +) -> str | None: + """Assemble ``{chip}_{n}`` from pytest ``-m`` chips/cards, or None to skip. + + Unset ``MIRROR_HW``: ``H100`` if in ``-m``, else ``L4``, else skip. + Set ``MIRROR_HW``: that chip if in ``-m``, else skip. + ``cards`` is a registered ``cards_n``, ``"max"`` for that chip's highest + existing preset, or None (error after a chip matched — no ``cards_*``). + """ + selector = _get_mirror_hw_selector() + if selector: + chip = selector if selector in chips else None + skip_reason = f"MIRROR_HW={selector!r} not in -m SKUs {sorted(chips)}" + else: + chip = next((c for c in _DEFAULT_INFER_CHIPS if c in chips), None) + skip_reason = f"no H100 or L4 in -m SKUs {sorted(chips)}" + if chip is None: + _log(f"skip {step_label}: {skip_reason}") + return None + if cards is None: raise ValueError( - f"mirror_hardwares must be a non-empty string in step {_get_step_label(step)!r}", + f"step {step_label!r} has a pytest -m SKU marker but no cards_* " + f"(or not cards_*); add cards_n or set mirror_hardwares to force a preset", ) - preset = _load_mirror_hardwares().get(hardware) + count_source = "not cards_*" if cards == _CARDS_CHIP_MAX else f"cards_{cards}" + count = None if cards == _CARDS_CHIP_MAX else cards + registered = get_supported_card_counts() + if count is not None and count not in registered: + supported = ", ".join(f"cards_{n}" for n in sorted(registered)) or "(none)" + raise ValueError( + f"{count_source} in step {step_label!r} is not a registered cards_* mark; supported: {supported}", + ) + known = _load_mirror_hardwares() + if count is None: + available = [n for n in sorted(registered) if f"{chip}_{n}" in known] + if not available: + raise ValueError( + f"{count_source} in step {step_label!r} has no CUDA count presets for {chip}", + ) + count = max(available) + chosen = f"{chip}_{count}" + if chosen not in known: + raise ValueError( + f"{count_source} in step {step_label!r} has no preset {chosen!r}", + ) + _log(f"{step_label}: {count_source} → {chosen!r}") + return chosen + + +def _resolve_mirror_hardware_name(hardware: Any, *, step_label: str) -> str | None: + """Map ``mirror_hardwares`` to a ``ci_mirror_hardwares.yml`` preset, or None to skip. + + The value is a preset name (``h100_2``). Unknown names fail later when the + registry is loaded. CUDA names are omitted when ``MIRROR_HW`` names a + different chip; NPU strings ignore ``MIRROR_HW``. + """ + if isinstance(hardware, bool): + raise ValueError(f"mirror_hardwares must not be a boolean in step {step_label!r}") + if isinstance(hardware, str): + name = hardware.strip() + elif isinstance(hardware, int): + name = str(hardware) + else: + raise ValueError( + f"mirror_hardwares must be a preset string in step {step_label!r}", + ) + if not name: + raise ValueError(f"mirror_hardwares must be a non-empty string in step {step_label!r}") + + selector = _get_mirror_hw_selector() + token = name.lower() + chip = next( + (c for c in _cuda_mirror_chips() if token == c or token.startswith(f"{c}_")), + None, + ) + if selector and chip is not None and chip != selector: + _log( + f"skip {step_label}: preset {name!r} is {chip} hardware; MIRROR_HW={selector!r}", + ) + return None + return name + + +def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None: + """Replace ``mirror_hardwares`` (or infer it from ``-m``) with preset agents/plugins. + + Explicit preset string: used as-is; pytest marks are ignored. + Omitted key: match H100/L4 (or ``MIRROR_HW``) in ``-m``, then compose + ``{chip}_{n}`` from ``cards_n``. No match skips the step. Steps with no + SKU/cards and no ``mirror_hardwares`` are left unchanged (CPU jobs). + + Returns None when the current ``MIRROR_HW`` selector does not apply. + """ + step_label = _get_step_label(step) + has_pool = step.get("agents") is not None or step.get("plugins") is not None or step.get("image") is not None + + if "mirror_hardwares" not in step: + if has_pool: + return step + chips, cards = _read_pytest_marks(step.get("commands")) + if not chips and cards is None: + return step + preset_name = _compose_mirror_hardware_name(chips, cards, step_label=step_label) + else: + if has_pool: + raise ValueError( + f"step {step_label!r} sets mirror_hardwares together with agents/plugins/image; " + f"use mirror_hardwares only", + ) + preset_name = _resolve_mirror_hardware_name( + step.get("mirror_hardwares"), + step_label=step_label, + ) + + if preset_name is None: + return None + + preset = _load_mirror_hardwares().get(preset_name) if preset is None: known = ", ".join(sorted(_load_mirror_hardwares())) raise ValueError( - f"unknown mirror_hardwares {hardware!r} in step {_get_step_label(step)!r}; known: {known}", - ) - - if step.get("agents") is not None or step.get("plugins") is not None or step.get("image") is not None: - raise ValueError( - f"step {_get_step_label(step)!r} sets mirror_hardwares together with agents/plugins/image; " - "use mirror_hardwares only", + f"unknown mirror_hardwares {preset_name!r} in step {step_label!r}; known: {known}", ) expanded = copy.deepcopy(preset) @@ -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: - processed.append( - _expand_mirror_hardwares( - {key: value for key, value in step.items() if key != "source_file_dependencies"}, - ), - ) - else: - processed.append(_expand_mirror_hardwares(step)) + leaf = {key: value for key, value in step.items() if key != "source_file_dependencies"} + expanded = _expand_mirror_hardwares(leaf) + if expanded is None: + continue + processed.append(expanded) return processed @@ -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 - else: - changed_files = ctx.changed_files + changed_files = None if force_all or e2e_only else ctx.changed_files doc = yaml.safe_load(text) if not isinstance(doc, dict): diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index bd6a32db6..f00f30f81 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.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 diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 7c1d5f799..1e6bc8f4a 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.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 - mirror_hardwares: h100_1 + - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and B200 and omni and cards_1" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py - label: ":full_moon: Omni · Function Test with H100 · 2-GPU" timeout_in_minutes: 90 commands: - - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py - mirror_hardwares: h100_2 + - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and B200 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py - label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test" timeout_in_minutes: 50 @@ -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 diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index cb41afbca..7f38f534b 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.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 diff --git a/.buildkite/cuda/test-weekly.yml b/.buildkite/cuda/test-weekly.yml index ae1f60971..b532244ee 100644 --- a/.buildkite/cuda/test-weekly.yml +++ b/.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 diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index c7cb4409d..7348ad183 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/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: ` (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: ` (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--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: diff --git a/docs/contributing/ci/test_system_overview.md b/docs/contributing/ci/test_system_overview.md index b4d3ce121..f432678bc 100644 --- a/docs/contributing/ci/test_system_overview.md +++ b/docs/contributing/ci/test_system_overview.md @@ -136,7 +136,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L4
L4: Purpose, Test Content, Directory Location, Example - Nightly + Nightly / Days before Release
(additional GPU SKUs) GPU @@ -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)). diff --git a/docs/user_guide/diffusion/attention_backends/dense_backends.md b/docs/user_guide/diffusion/attention_backends/dense_backends.md index d7182eb44..2cb7019be 100644 --- a/docs/user_guide/diffusion/attention_backends/dense_backends.md +++ b/docs/user_guide/diffusion/attention_backends/dense_backends.md @@ -21,10 +21,11 @@ vllm-omni serve --diffusion-attention-backend TORCH_SDPA ## `FLASH_ATTN` `FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it -prefers the CuTe-based FlashAttention 4 path and falls back to -FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and -Ampere it is the preferred automatic route when a compatible package is -installed. +is FA4 only (`flash_attn.cute` / `vllm-omni[fa4]`). If FA4 is unavailable, +explicit `FLASH_ATTN` is rejected and automatic selection continues to another +compatible Blackwell backend. Hopper-only FA2/FA3 wheels are not used, even if +they import. On Hopper, Ada, and Ampere it is the preferred automatic route +when a compatible package is installed. ```bash vllm-omni serve --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 -compatible FlashAttention 3 or 2 path. +Blackwell. On Blackwell, `FLASH_ATTN` is FA4 only: Hopper-only FA2/FA3 +wheels are not used, even if they import. If FA4 is missing, explicit +`FLASH_ATTN` is rejected and automatic selection continues to other +Blackwell backends. ## `TRTLLM_ATTN` @@ -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 --diffusion-attention-backend FLASHINFER_ATTN ``` diff --git a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md index 946786b97..846e69fbc 100644 --- a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md +++ b/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 -corresponding local FlashAttention implementation. On CUDA devices older than -Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`. +If `kernels` is unavailable, an explicit `FLASH_ATTN_HUB` or +`FLASH_ATTN_3_HUB` selection raises. Automatic / `FLASH_ATTN` selection in the +Diffusers adapter may still walk hub then local FlashAttention backends. ## Usage diff --git a/docs/user_guide/fault_injection_reliability_matrix.md b/docs/user_guide/fault_injection_reliability_matrix.md index 4a2fa0a16..901c5b171 100644 --- a/docs/user_guide/fault_injection_reliability_matrix.md +++ b/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 `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | | +| ------ | ---------------------------- | ---------- | ---------- | +| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | | +| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | | | No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725)
[#43060](https://github.com/vllm-project/vllm/issues/43060) | -| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released | | -| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released | | -| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly | | -| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed | | +| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released | | +| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released | | +| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly | | +| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed | | | Under load | Send `SIGKILL` to Worker | In-flight requests are hard interrupted (5xx/connection drop); main service becomes unavailable | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | -| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable | | +| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable | | | Under load | Send `SIGKILL` to serve main process | serve is hard-killed and current connections are interrupted; in-flight requests fail; after cleanup there are no residual processes and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | | Under load | Send `SIGTERM` to serve main process | serve stops accepting new requests and executes shutdown flow; in-flight requests fail; no residue remains and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | | Under load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | `Ctrl+C`-style serve shutdown; in-flight requests fail (5xx/connection interruption); service unavailable; after exit there is no residue and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | -| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly | | +| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly | | | Under load | Send `SIGTERM` to all related processes | All processes exit gracefully; in-flight requests fail; service unavailable; GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) | | OOM | Occupy all free GPU memory via an extra process | After OOM injection process starts, GPU memory is continuously saturated; service enters unavailable/degraded state and health check drops to 503; different request types (chat/speech, etc.) fail fast within a fixed time and return 500 (no hanging) | [#4285](https://github.com/vllm-project/vllm-omni/issues/4285) | diff --git a/tests/buildkite/test_skip_ci.py b/tests/buildkite/test_skip_ci.py index ddefc7271..4c90b3ff7 100644 --- a/tests/buildkite/test_skip_ci.py +++ b/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 diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 451b553b1..9011a1aef 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/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() diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 75d1bdc21..bb69391fd 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/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, diff --git a/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json b/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json index a28aab96c..e16b40e91 100644 --- a/tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json +++ b/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, diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py index 2e98dcaf7..d25f66709 100644 --- a/tests/dfx/perf/tests/test_runner_metadata.py +++ b/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) diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 54ad2a32b..b695d59fe 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/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 - produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it). + Regression: ``str(str, Enum)`` later produced ``RainFusionPrecision.MIX`` + instead of ``mix``. ``BlockSparseSpec`` stores ``member.value``. """ from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig from vllm_omni.diffusion.data import RainFusionPrecision @@ -509,6 +516,49 @@ class TestAttentionInitUsesCurrentDiffusionConfig: dense_flash.assert_not_called() sdpa.assert_called_once() + @pytest.mark.parametrize("fa4_available", [False, True]) + def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch, fa4_available): + attention = Attention.__new__(Attention) + attention._scheduler_paged_kv = True + attention.paged_kv_cache_role = "primary" + attention.backend_pref = "FLASH_ATTN" + attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN") + flash_output = torch.ones(1) + sdpa_output = torch.zeros(1) + dense_flash = Mock(return_value=flash_output) + sdpa = Mock(return_value=sdpa_output) + attention.attention = SimpleNamespace(forward=dense_flash) + attention.sdpa_fallback = SimpleNamespace(forward=sdpa) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True) + monkeypatch.setattr( + layer_mod, + "get_forward_context", + lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True), + ) + monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: False) + monkeypatch.setattr(layer_mod.current_omni_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + layer_mod.current_omni_platform, + "supports_diffusion_dense_flash_attention", + lambda: fa4_available, + ) + + result = attention._run_local_attention( + torch.zeros(1, dtype=torch.bfloat16), + torch.zeros(1, dtype=torch.bfloat16), + torch.zeros(1, dtype=torch.bfloat16), + AttentionMetadata(), + ) + + if fa4_available: + assert result is flash_output + dense_flash.assert_called_once() + sdpa.assert_not_called() + else: + assert result is sdpa_output + dense_flash.assert_not_called() + sdpa.assert_called_once() + def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch): attention = Attention.__new__(Attention) attention._scheduler_paged_kv = False @@ -579,7 +629,106 @@ class TestAttentionInitUsesCurrentDiffusionConfig: assert len(calls) == 1 assert attention.attn_backend is _DenseBackend - assert attention.backend_pref is None + # Auto path records the resolved backend name for Ring/SP; it is not explicit. + assert attention.backend_pref == _DenseBackend.get_name() + assert attention.backend_explicit is False + + def test_internal_paged_selector_does_not_require_dense_fa4(self, monkeypatch): + """Blackwell CUDNN default + marked paged layer must not act as explicit FLASH_ATTN.""" + + class _FakeAttentionImpl: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def forward(self, query, key, value, attn_metadata=None): + return query + + class _CuDNNBackend: + supports_paged_kv = False + + @staticmethod + def get_name() -> str: + return "CUDNN_ATTN" + + @staticmethod + def get_impl_cls(): + return _FakeAttentionImpl + + def _fake_get_attn_backend_for_role(*args, attention_config=None, **kwargs): + if attention_config is not None: + spec, _ = attention_config.resolve_with_source(role="self", role_category=None) + if spec is not None and spec.backend.upper() == "FLASH_ATTN": + raise ValueError( + "FLASH_ATTN was explicitly selected but is unsupported " + "(Blackwell requires CuTe FlashAttention-4 " + "(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only). " + "Select a compatible backend." + ) + return _CuDNNBackend, None + + monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role) + monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) + monkeypatch.setattr(FlashAttentionBackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object()) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) + + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(), + diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.PAGED_SCHEDULER, + parallel_config=SimpleNamespace(ring_degree=1), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + with set_current_diffusion_config(od_config): + attention = Attention( + num_heads=4, + head_size=64, + causal=False, + softmax_scale=1.0, + paged_kv_cache_role="primary", + ) + + assert attention.attn_backend is FlashAttentionBackend + assert attention.backend_pref == "FLASH_ATTN" + assert attention.backend_explicit is False + assert attention.attn_spec is None + + def test_user_explicit_flash_attn_still_requires_dense_fa4(self, monkeypatch): + def _fake_get_attn_backend_for_role(*args, attention_config=None, **kwargs): + if attention_config is not None: + spec, _ = attention_config.resolve_with_source(role="self", role_category=None) + if spec is not None and spec.backend.upper() == "FLASH_ATTN": + raise ValueError( + "FLASH_ATTN was explicitly selected but is unsupported " + "(Blackwell requires CuTe FlashAttention-4 " + "(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only). " + "Select a compatible backend." + ) + raise AssertionError("explicit FLASH_ATTN must not fall through to the platform default") + + monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object()) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) + + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")), + diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.PAGED_SCHEDULER, + parallel_config=SimpleNamespace(ring_degree=1), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + with set_current_diffusion_config(od_config), pytest.raises(ValueError, match="explicitly selected"): + Attention( + num_heads=4, + head_size=64, + causal=False, + softmax_scale=1.0, + paged_kv_cache_role="primary", + ) def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch): class _FakeAttentionImpl: @@ -614,9 +763,10 @@ class TestAttentionInitUsesCurrentDiffusionConfig: return _FakeBackend, AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5}) class _FakeRingParallelAttention: - def __init__(self, sp_group, attn_backend_pref=None): + def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False): self.sp_group = sp_group self.attn_backend_pref = attn_backend_pref + self.attn_backend_explicit = attn_backend_explicit monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role) monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) @@ -662,11 +812,134 @@ class TestAttentionInitUsesCurrentDiffusionConfig: assert attn.use_ring is True assert attn.ring_runner is not None assert attn.ring_runner.attn_backend_pref == "TRTLLM_ATTN" + assert attn.ring_runner.attn_backend_explicit is True + + def test_attention_init_forwards_flashinfer_backend_without_dtype(self, monkeypatch): + class _FakeAttentionImpl: + def __init__(self, **kwargs): + self.kwargs = kwargs + + class _FakeBackend: + @staticmethod + def get_name() -> str: + return "FLASHINFER_ATTN" + + @staticmethod + def get_impl_cls(): + return _FakeAttentionImpl + + spec = AttentionSpec(backend="FLASHINFER_ATTN", quant={"flashinfer_backend": "fa2"}) + + monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", lambda **kwargs: (_FakeBackend, spec)) + monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object()) + monkeypatch.setattr(layer_mod, "get_sp_group", lambda: None) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) + + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(default=spec), + parallel_config=SimpleNamespace(ring_degree=1), + diffusion_kv_cache_dtype=None, + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + with set_current_diffusion_config(od_config): + attn = Attention( + num_heads=4, + head_size=64, + causal=False, + softmax_scale=1.0, + role="self", + ) + + assert attn.backend_pref == "FLASHINFER_ATTN" + assert attn.attention.kwargs["backend_kwargs"] == {"quant": {"flashinfer_backend": "fa2"}} + assert attn.attention.kwargs["backend_explicit"] is True + + def test_attention_init_propagates_ring_setup_failure(self, monkeypatch): + class _FakeAttentionImpl: + def __init__(self, **kwargs): + pass + + class _FakeBackend: + @staticmethod + def get_name() -> str: + return "TORCH_SDPA" + + @staticmethod + def get_impl_cls(): + return _FakeAttentionImpl + + monkeypatch.setattr( + layer_mod, + "get_attn_backend_for_role", + lambda **kwargs: (_FakeBackend, AttentionSpec(backend="TORCH_SDPA")), + ) + monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) + monkeypatch.setattr( + layer_mod, + "get_sp_group", + lambda: (_ for _ in ()).throw(RuntimeError("SP group is not initialized")), + ) + + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"), + parallel_config=SimpleNamespace(ring_degree=2), + ) + with set_current_diffusion_config(od_config): + with pytest.raises(RuntimeError, match="SP group is not initialized"): + Attention(num_heads=4, head_size=64, causal=False, softmax_scale=1.0) + + def test_float32_main_dispatch_does_not_override_selected_backend(self): + sentinel = torch.randn(1, 2, 4, 8) + selected_calls = [] + + def _selected_forward(*args): + selected_calls.append(args) + return sentinel + + fake_attention = SimpleNamespace( + attention=SimpleNamespace( + forward=_selected_forward, + ), + sdpa_fallback=SimpleNamespace( + forward=lambda *args: pytest.fail("unexpected SDPA fallback"), + ), + _assert_metadata_compatible=lambda metadata: None, + _scheduler_paged_kv=False, + paged_kv_cache_role=None, + ) + query = torch.randn(1, 2, 4, 8, dtype=torch.float32) + + output = Attention._run_local_attention(fake_attention, query, query, query, None) + + assert output is sentinel + assert len(selected_calls) == 1 + + def test_unsupported_attention_mask_is_rejected_before_backend_call(self): + fake_backend = SimpleNamespace( + get_name=lambda: "NO_MASK_BACKEND", + supports_attention_mask=lambda *args, **kwargs: False, + ) + fake_attention = SimpleNamespace(attn_backend=fake_backend) + metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool)) + + with pytest.raises(ValueError, match="does not support attn_mask"): + Attention._assert_metadata_compatible(fake_attention, metadata) + + def test_ring_attention_rejects_mask_instead_of_ignoring_it(self): + fake_attention = SimpleNamespace(attention=SimpleNamespace(skip=None), ring_runner=None) + query = torch.randn(1, 2, 4, 8) + metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool)) + + with pytest.raises(ValueError, match="Ring attention does not support attn_mask"): + Attention._run_ring_attention(fake_attention, query, query, query, metadata) class TestDiffusionKvCacheQuantization: @staticmethod - def _install_attention_init_stubs(monkeypatch): + def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True): class _FakeAttentionImpl: def __init__(self, **kwargs): self.kwargs = kwargs @@ -674,6 +947,10 @@ class TestDiffusionKvCacheQuantization: def forward(self, query, key, value, attn_metadata=None): return query + @staticmethod + def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool: + return supports_kv_cache_dtype + class _FakeBackend: @staticmethod def get_name() -> str: @@ -683,14 +960,11 @@ class TestDiffusionKvCacheQuantization: def get_impl_cls(): return _FakeAttentionImpl - @staticmethod - def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool: - return True - class _FakeRingParallelAttention: - def __init__(self, sp_group, attn_backend_pref=None): + def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False): self.sp_group = sp_group self.attn_backend_pref = attn_backend_pref + self.attn_backend_explicit = attn_backend_explicit monkeypatch.setattr( layer_mod, @@ -764,3 +1038,22 @@ class TestDiffusionKvCacheQuantization: causal=False, softmax_scale=1.0, ) + + def test_explicit_unsupported_kv_cache_dtype_raises(self, monkeypatch): + self._install_attention_init_stubs(monkeypatch, supports_kv_cache_dtype=False) + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(), + parallel_config=SimpleNamespace(ring_degree=1), + diffusion_kv_cache_dtype="fp8", + diffusion_kv_cache_skip_step_indices=None, + diffusion_kv_cache_skip_layer_indices=None, + ) + + with set_current_diffusion_config(od_config): + with pytest.raises(ValueError, match="does not support kv_cache_dtype='fp8'"): + Attention( + num_heads=4, + head_size=64, + causal=False, + softmax_scale=1.0, + ) diff --git a/tests/diffusion/attention/test_cuda_backend_selection.py b/tests/diffusion/attention/test_cuda_backend_selection.py new file mode 100644 index 000000000..777bf1a7c --- /dev/null +++ b/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 diff --git a/tests/diffusion/attention/test_cudnn_attn.py b/tests/diffusion/attention/test_cudnn_attn.py index 1ff3baafb..5c787a91a 100644 --- a/tests/diffusion/attention/test_cudnn_attn.py +++ b/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( - ("query_sequence_length", "key_sequence_length", "expected_backend"), - [ - (1, 1, SDPBackend.MATH), - (8, 1, SDPBackend.MATH), - (1, 8, SDPBackend.MATH), - (8, 8, SDPBackend.CUDNN_ATTENTION), - ], -) -def test_cudnn_routes_single_token_attention_to_math( - monkeypatch, query_sequence_length, key_sequence_length, expected_backend -): +def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch): + """Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V.""" selected_backends = [] + @contextmanager def fake_sdpa_kernel(backends): - selected_backends.append(backends) - return nullcontext() + selected_backends.append(tuple(backends)) + yield - monkeypatch.setattr( - "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel", - fake_sdpa_kernel, - ) - monkeypatch.setattr( - torch.nn.functional, - "scaled_dot_product_attention", - lambda query, _key, _value, **_kwargs: query, - ) - impl = CuDNNAttentionImpl( - num_heads=2, - head_size=4, - softmax_scale=0.5, - ) - query = torch.randn(1, query_sequence_length, 2, 4) - key = torch.randn(1, key_sequence_length, 2, 4) - value = torch.randn_like(key) + def fake_sdpa(query, key, value, **kwargs): + return query - output = impl.forward_cuda(query, key, value) + monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel) + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) + + impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5) + query = torch.randn(1, 2, 2, 8) + singleton_kv = torch.randn(1, 1, 2, 8) + + output = impl.forward_cuda(query, singleton_kv, singleton_kv) assert output.shape == query.shape - assert selected_backends == [[expected_backend]] + assert selected_backends == [(SDPBackend.MATH,)] + + +def test_explicit_cudnn_rejects_kv_seq_len_one(): + impl = CuDNNAttentionImpl( + num_heads=2, + head_size=8, + softmax_scale=0.5, + backend_explicit=True, + ) + query = torch.randn(1, 2, 2, 8) + singleton_kv = torch.randn(1, 1, 2, 8) + + with pytest.raises(ValueError, match="explicitly selected.*sequence length 1"): + impl.forward_cuda(query, singleton_kv, singleton_kv) + + +def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch): + selected_backends = [] + + @contextmanager + def fake_sdpa_kernel(backends): + selected_backends.append(tuple(backends)) + yield + + def reject_shape(*args, **kwargs): + raise RuntimeError("No available kernel. Aborting execution.") + + monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel) + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", reject_shape) + + impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5) + tensors = torch.randn(1, 2, 2, 8) + + with pytest.raises(RuntimeError, match="No available kernel"): + impl.forward_cuda(tensors, tensors, tensors) + + assert selected_backends == [(SDPBackend.CUDNN_ATTENTION,)] def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch): - observed = {} + observed: dict[str, Any] = {} def fake_sdpa(query, key, value, **kwargs): observed.update(query=query, key=key, value=value, kwargs=kwargs) @@ -105,3 +127,15 @@ def test_cudnn_rejects_invalid_valid_kv_length(): query, AttentionMetadata(extra={"valid_kv_length": 9}), ) + + +@pytest.mark.parametrize("head_size", [8, 64, 128, 256]) +def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size): + assert CuDNNAttentionBackend.supports_head_size(head_size) + assert head_size in CuDNNAttentionBackend.get_supported_head_sizes() + + +@pytest.mark.parametrize("head_size", [0, 7, 12, 320]) +def test_cudnn_backend_rejects_incompatible_head_sizes(head_size): + assert not CuDNNAttentionBackend.supports_head_size(head_size) + assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes() diff --git a/tests/diffusion/attention/test_flashinfer_attn.py b/tests/diffusion/attention/test_flashinfer_attn.py new file mode 100644 index 000000000..bb81c58bf --- /dev/null +++ b/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 diff --git a/tests/diffusion/attention/test_kernels_hub.py b/tests/diffusion/attention/test_kernels_hub.py index 0f22a1244..3f7451ca3 100644 --- a/tests/diffusion/attention/test_kernels_hub.py +++ b/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): - """Test that when kernels package is not available, the platform fallback logic - - routes to local FLASH_ATTN. - """ +def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch): from vllm_omni.platforms.cuda.platform import CudaOmniPlatform # Temporarily hide/remove kernels module if present @@ -38,19 +33,73 @@ def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch): ) monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True}) - # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable - backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) - assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() + with pytest.raises(ImportError, match="explicitly selected"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) + with pytest.raises(ImportError, match="explicitly selected"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) - # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable - backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) - assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() - - # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs kernels_module = types.ModuleType("kernels") monkeypatch.setitem(sys.modules, "kernels", kernels_module) - backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) - assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path() + with pytest.raises(ValueError, match="require.*Hopper GPU"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)), + ) + with pytest.raises(ValueError, match="current FA2 kernels require"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) + with pytest.raises(ValueError, match="require.*Hopper GPU"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_explicit_flash_attention_unavailable_does_not_fallback(monkeypatch: pytest.MonkeyPatch): + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False}) + + with pytest.raises(ValueError, match="explicitly selected but is unsupported"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64) + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True}) + monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: False)) + with pytest.raises(ValueError, match="requires CuTe FlashAttention-4"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_explicit_sage_attention_unsupported_arch_does_not_reach_kernel(monkeypatch: pytest.MonkeyPatch): + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False}) + + with pytest.raises(ValueError, match="does not provide a kernel for sm_103"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN", head_size=64) @pytest.mark.core_model diff --git a/tests/diffusion/attention/test_ring_backend_selection.py b/tests/diffusion/attention/test_ring_backend_selection.py new file mode 100644 index 000000000..a6814ef49 --- /dev/null +++ b/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) diff --git a/tests/diffusion/attention/test_ring_fa4.py b/tests/diffusion/attention/test_ring_fa4.py index b24ce2f20..b465e86a4 100644 --- a/tests/diffusion/attention/test_ring_fa4.py +++ b/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) diff --git a/tests/diffusion/attention/test_ring_utils.py b/tests/diffusion/attention/test_ring_utils.py new file mode 100644 index 000000000..18a3636d8 --- /dev/null +++ b/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)))) diff --git a/tests/diffusion/attention/test_sage_attn.py b/tests/diffusion/attention/test_sage_attn.py new file mode 100644 index 000000000..dde27d9f2 --- /dev/null +++ b/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) diff --git a/tests/diffusion/attention/test_sage_attn3.py b/tests/diffusion/attention/test_sage_attn3.py index 5c46ea5d0..e7d78f704 100644 --- a/tests/diffusion/attention/test_sage_attn3.py +++ b/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) - assert sdpa_calls["enable_gqa"] is True - expected = (query.permute(0, 2, 1, 3) + 1).permute(0, 2, 1, 3) - assert torch.allclose(output, expected) + +def test_sage_attn3_rejects_mask_instead_of_ignoring_it(monkeypatch: pytest.MonkeyPatch): + def fake_kernel(*args, **kwargs): + raise AssertionError("sageattn3_blackwell should not run with an unsupported mask") + + backend_module = load_sage_attn3_module(monkeypatch, fake_kernel) + impl = backend_module.SageAttention3Impl( + num_heads=4, + head_size=64, + softmax_scale=1.0 / 8.0, + causal=False, + ) + query = torch.randn(2, 8, 4, 64) + metadata = AttentionMetadata(attn_mask=torch.ones(2, 8, dtype=torch.bool)) + + with pytest.raises(ValueError, match="does not support attn_mask"): + impl.forward_cuda(query, query, query, metadata) + + +def test_sage_attn3_rejects_custom_softmax_scale(monkeypatch: pytest.MonkeyPatch): + backend_module = load_sage_attn3_module(monkeypatch, lambda *args, **kwargs: None) + + with pytest.raises(ValueError, match="does not expose a custom softmax scale"): + backend_module.SageAttention3Impl( + num_heads=4, + head_size=64, + softmax_scale=1.0, + causal=False, + ) @pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform") @@ -115,10 +132,39 @@ def test_cuda_platform_selects_sage_attn3_alias(monkeypatch: pytest.MonkeyPatch) @pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform") -def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize("head_size", [32, 320]) +def test_cuda_platform_rejects_explicit_sage_attn3_unsupported_head_size( + monkeypatch: pytest.MonkeyPatch, head_size: int +): + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + from vllm_omni.platforms.cuda import platform as cuda_platform_module + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + original_import_module = importlib.import_module + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False}) + monkeypatch.setattr( + cuda_platform_module.importlib, + "import_module", + lambda module_name: object() if module_name == "sageattn3" else original_import_module(module_name), + ) + load_sage_attn3_module(monkeypatch, lambda *args, **kwargs: None) + + with pytest.raises(ValueError, match=f"head_size={head_size} is unsupported"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=head_size) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform") +def test_cuda_platform_rejects_explicit_sage_attn3_on_unsupported_gpu(monkeypatch: pytest.MonkeyPatch): from vllm.platforms.interface import DeviceCapability - from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum from vllm_omni.diffusion.envs import PACKAGES_CHECKER from vllm_omni.platforms.cuda.platform import CudaOmniPlatform @@ -129,6 +175,33 @@ def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch ) monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False}) - backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64) + with pytest.raises(ValueError, match="explicitly selected"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64) - assert backend_path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform") +def test_cuda_platform_rejects_missing_explicit_sage_attn3(monkeypatch: pytest.MonkeyPatch): + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + from vllm_omni.platforms.cuda import platform as cuda_platform_module + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + original_import_module = importlib.import_module + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False}) + + def _missing_sageattn3(module_name): + if module_name == "sageattn3": + raise ImportError("sageattn3 not installed") + return original_import_module(module_name) + + monkeypatch.setattr(cuda_platform_module.importlib, "import_module", _missing_sageattn3) + + with pytest.raises(ImportError, match="explicitly selected"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64) diff --git a/tests/diffusion/attention/test_selector.py b/tests/diffusion/attention/test_selector.py index 79755ac2e..7b16fb754 100644 --- a/tests/diffusion/attention/test_selector.py +++ b/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") diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 165a98cc6..a24595c9b 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/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( - in_channels=16, - out_channels=16, - patch_size=1, - patch_size_t=1, - num_attention_heads=2, - attention_head_dim=8, - cross_attention_dim=16, - audio_in_channels=16, - audio_out_channels=16, - audio_num_attention_heads=2, - audio_attention_head_dim=8, - audio_cross_attention_dim=16, - num_layers=2, - caption_channels=16, - ) + with _force_torch_sdpa(): + model = LTX2VideoTransformer3DModel( + in_channels=16, + out_channels=16, + patch_size=1, + patch_size_t=1, + num_attention_heads=2, + attention_head_dim=8, + cross_attention_dim=16, + audio_in_channels=16, + audio_out_channels=16, + audio_num_attention_heads=2, + audio_attention_head_dim=8, + audio_cross_attention_dim=16, + num_layers=2, + caption_channels=16, + ) # NOTE: This is currently using the LTX2 custom enabler, but the custom # enablers will be consolidated after diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index 6311d494d..c5b71f502 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/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", [ diff --git a/tests/diffusion/distributed/test_sp_plan_hooks.py b/tests/diffusion/distributed/test_sp_plan_hooks.py index 4ef564571..4ae6ec064 100644 --- a/tests/diffusion/distributed/test_sp_plan_hooks.py +++ b/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: diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 19b7ef04c..9daff16fe 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/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. diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index e9bba1f1f..f71b84d3e 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/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): diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py index 832d34557..cc1b7c589 100644 --- a/tests/e2e/accuracy/helpers.py +++ b/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[0-9.]+)") _PSNR_RE = re.compile(r"average:(?P[0-9.]+)") diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py index 16cd9c13c..71dd6f308 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py +++ b/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, - psnr_threshold=_psnr_threshold(), + ssim_threshold=thresholds.ssim, + psnr_threshold=psnr_threshold, ) +@pytest.mark.full_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4) def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> None: model = _model_name() @@ -284,6 +313,7 @@ def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> N _assert_against_baseline(image, "online") +@pytest.mark.full_model @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=4) def test_hunyuan_image3_pixel_accuracy_offline(accuracy_artifact_root: Path) -> None: model = _model_name() diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py index a995a9cf9..771ced5bc 100644 --- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py +++ b/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] - - -def test_parse_video_metadata_extracts_dimensions_and_fps() -> None: - payload = { - "streams": [ - { - "width": 832, - "height": 480, - "avg_frame_rate": "16/1", - "nb_read_frames": "93", - } - ] - } - - metadata = _parse_video_metadata(payload) - - assert metadata["width"] == 832 - assert metadata["height"] == 480 - assert metadata["fps"] == 16.0 - assert metadata["frame_count"] == 93 - - -def test_parse_ssim_summary_extracts_all_score() -> None: - output = """ - [Parsed_ssim_0 @ 000001] SSIM Y:0.971903 (15.512007) U:0.965077 (14.569044) V:0.962414 (14.252637) All:0.968311 (15.035654) - """ - - assert _parse_ssim_score(output) == 0.968311 - - -def test_parse_psnr_summary_extracts_average_score() -> None: - output = """ - [Parsed_psnr_0 @ 000001] PSNR y:32.670157 u:31.844621 v:31.513839 average:32.148004 min:31.563744 max:33.201457 - """ - - assert _parse_psnr_score(output) == 32.148004 - - -def test_build_diffusers_command_uses_python_runner_path(tmp_path: Path) -> None: - runner_path = tmp_path / "run_wan22_i2v_diffusers_cp.py" - command = _build_diffusers_command( - runner_path=runner_path, - image_source=RABBIT_IMAGE_URL, - output_path=tmp_path / "offline.mp4", - metadata_path=tmp_path / "offline.json", - ) - - assert command[:2] == [ - sys.executable, - str(runner_path), - ] - assert "--output" in command - assert "--metadata-output" in command - - -def test_resolve_image_source_prefers_existing_local_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - image_path = tmp_path / "rabbit.png" - Image.new("RGB", (8, 8), color=(255, 128, 64)).save(image_path) - - assert _resolve_image_source(str(image_path)) == str(image_path.resolve()) - - -def test_build_online_image_reference_uses_data_url_for_local_path(tmp_path: Path) -> None: - image_path = tmp_path / "rabbit.png" - Image.new("RGB", (4, 2), color=(10, 20, 30)).save(image_path) - - reference = _build_online_image_reference(str(image_path)) - - assert reference.startswith("data:image/png;base64,") - encoded = reference.split(",", 1)[1] - assert len(b64decode(encoded)) > 0 - - -def test_send_video_request_with_timeout_uses_requested_timeout(monkeypatch: pytest.MonkeyPatch) -> None: - class DummyResponse: - def raise_for_status(self) -> None: - return None - - def json(self) -> dict[str, str]: - return {"id": "video_123"} - - class DummyClient: - def __init__(self) -> None: - self.run_level = "core_model" - self.wait_args: tuple[str, int, int] | None = None - - def _build_url(self, path: str) -> str: - return f"http://localhost:8000{path}" - - def _wait_until_video_completed( - self, - video_id: str, - poll_interval_seconds: int = 2, - timeout_seconds: int = 300, - ) -> None: - self.wait_args = (video_id, poll_interval_seconds, timeout_seconds) - - def _download_video_content(self, video_id: str) -> bytes: - return b"video-bytes" - - monkeypatch.setattr(requests, "post", lambda *args, **kwargs: DummyResponse()) - - client = DummyClient() - result = _send_video_request_with_timeout( - client, - { - "form_data": {"prompt": "test"}, - "image_reference": RABBIT_IMAGE_URL, - }, - timeout_seconds=1200, - ) - - assert result == b"video-bytes" - assert client.wait_args == ("video_123", 2, 1200) - - -def test_online_timeout_defaults_to_1200() -> None: - assert _online_timeout_seconds(None) == 1200 - - -def test_artifact_dir_is_under_repo_result_folder(tmp_path: Path) -> None: - artifact_dir = _artifact_dir(str(tmp_path / "rabbit.png")) - - assert artifact_dir.parent == Path(__file__).parent / "result" - assert artifact_dir.name.startswith("rabbit-") - - -def test_offline_cuda_device_uses_indexed_cuda_device() -> None: - assert _offline_cuda_device() == torch.device("cuda:0") - - -def test_ensure_wan_ftfy_fallback_sets_identity(monkeypatch: pytest.MonkeyPatch) -> None: - from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module - - monkeypatch.delattr(wan_i2v_module, "ftfy", raising=False) - _ensure_wan_ftfy_fallback() - - assert hasattr(wan_i2v_module, "ftfy") - assert isinstance(wan_i2v_module.ftfy, _IdentityFtfy) - assert wan_i2v_module.ftfy.fix_text("abc") == "abc" - - -def test_resize_to_target_matches_requested_dimensions() -> None: - image = Image.new("RGB", (100, 50), color=(1, 2, 3)) - - resized = _resize_to_target(image, width=832, height=480) - - assert resized.size == (832, 480) - - -def test_configure_scheduler_switches_to_unipc() -> None: - class DummyPipe: - def __init__(self) -> None: - self.scheduler = UniPCMultistepScheduler(num_train_timesteps=1000, flow_shift=1.0) - - pipe = DummyPipe() - _configure_scheduler(pipe, flow_shift=5.0) - - assert isinstance(pipe.scheduler, UniPCMultistepScheduler) - assert pipe.scheduler.config.flow_shift == 5.0 - +pytestmark = [pytest.mark.diffusion] REPO_ROOT = Path(__file__).resolve().parents[4] WORKSPACE_ROOT = REPO_ROOT.parent @@ -312,7 +120,7 @@ def _runner_env() -> dict[str, str]: if existing_pythonpath: pythonpath_parts.append(existing_pythonpath) env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts) - env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + return env @@ -324,12 +132,8 @@ def _resolve_image_source(configured: str | None) -> str: return configured -def _artifact_dir(image_source: str) -> Path: - return video_artifact_dir(RESULT_ROOT, image_source) - - def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]: - artifact_dir = _artifact_dir(image_source) + artifact_dir = video_artifact_dir(RESULT_ROOT, image_source) artifact_dir.mkdir(parents=True, exist_ok=True) return ( artifact_dir / "online.mp4", @@ -343,7 +147,7 @@ def _generate_online_video( omni_server, online_client, image_source: str, - online_timeout_seconds: int, + online_timeout_seconds_value: int, ) -> Path: online_path, _, _ = _artifact_paths(image_source) request_config = { @@ -360,12 +164,12 @@ def _generate_online_video( "num_inference_steps": NUM_INFERENCE_STEPS, "seed": SEED, }, - "image_reference": _build_online_image_reference(image_source), + "image_reference": build_online_image_reference(image_source), } - online_video_bytes = _send_video_request_with_timeout( + online_video_bytes = send_video_request_with_timeout( online_client, request_config, - timeout_seconds=_online_timeout_seconds(online_timeout_seconds), + timeout_seconds=online_timeout_seconds(online_timeout_seconds_value), ) online_path.write_bytes(online_video_bytes) return online_path @@ -390,6 +194,7 @@ def _generate_offline_video(*, image_source: str) -> tuple[Path, Path]: @pytest.mark.benchmark +@pytest.mark.full_model @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_wan22_i2v_diffusers_offline_generates_video( wan22_i2v_image_source: str | None, @@ -397,20 +202,21 @@ def test_wan22_i2v_diffusers_offline_generates_video( if not torch.cuda.is_available(): pytest.skip("Wan2.2 I2V diffusers offline test requires CUDA.") - _probe_binary("ffprobe") + probe_binary("ffprobe") if not RUNNER_PATH.exists(): raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}") image_source = _resolve_image_source(wan22_i2v_image_source) - _validate_image_source(image_source) + validate_image_source(image_source) offline_path, offline_metadata_path = _generate_offline_video(image_source=image_source) assert offline_path.exists(), f"Expected offline video artifact at {offline_path}" assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}" - offline_metadata = _probe_video(offline_path) + offline_metadata = probe_video(offline_path) assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES) @pytest.mark.benchmark +@pytest.mark.full_model @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_wan22_i2v_online_serving_generates_video( @@ -422,21 +228,22 @@ def test_wan22_i2v_online_serving_generates_video( if not torch.cuda.is_available() or torch.accelerator.device_count() < 2: pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.") - _probe_binary("ffprobe") + probe_binary("ffprobe") image_source = _resolve_image_source(wan22_i2v_image_source) - _validate_image_source(image_source) + validate_image_source(image_source) online_path = _generate_online_video( omni_server=omni_server, online_client=online_client, image_source=image_source, - online_timeout_seconds=wan22_i2v_online_timeout_seconds, + online_timeout_seconds_value=wan22_i2v_online_timeout_seconds, ) assert online_path.exists(), f"Expected online video artifact at {online_path}" - online_metadata = _probe_video(online_path) + online_metadata = probe_video(online_path) assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES) @pytest.mark.benchmark +@pytest.mark.full_model @hardware_test(res={"cuda": "H100"}, num_cards=2) def test_wan22_i2v_serving_matches_diffusers_video_similarity( wan22_i2v_image_source: str | None, @@ -444,13 +251,13 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity( if not torch.cuda.is_available() or torch.accelerator.device_count() < 2: pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.") - _probe_binary("ffmpeg") - _probe_binary("ffprobe") + probe_binary("ffmpeg") + probe_binary("ffprobe") if not RUNNER_PATH.exists(): raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}") image_source = _resolve_image_source(wan22_i2v_image_source) - _validate_image_source(image_source) + validate_image_source(image_source) online_path, offline_path, offline_metadata_path = _artifact_paths(image_source) if not online_path.exists(): @@ -462,8 +269,8 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity( assert offline_path.exists(), f"Expected offline video artifact at {offline_path}" assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}" - online_metadata = _probe_video(online_path) - offline_metadata = _probe_video(offline_path) + online_metadata = probe_video(online_path) + offline_metadata = probe_video(offline_path) assert online_metadata == offline_metadata, ( f"Video metadata mismatch:\n" f"online={online_metadata}\n" @@ -472,11 +279,17 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity( f"offline_path={offline_path}" ) assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES) + device_profile = resolve_device_profile(profiles=SIMILARITY_THRESHOLDS_BY_DEVICE) + thresholds = resolve_similarity_thresholds(SIMILARITY_THRESHOLDS_BY_DEVICE, device_profile) + print( + f"wan22_i2v similarity thresholds: device_profile={device_profile}, " + f"ssim>={thresholds.ssim:.6f}, psnr>={thresholds.psnr:.6f}" + ) assert_video_similarity_metrics( label="wan22_i2v", online_path=online_path, offline_path=offline_path, - ssim_threshold=SSIM_THRESHOLD, - psnr_threshold=PSNR_THRESHOLD, + ssim_threshold=thresholds.ssim, + psnr_threshold=thresholds.psnr, ) print(f"offline_metadata={offline_metadata_path}") diff --git a/tests/e2e/accuracy/wan22_i2v/wan22_i2v_video_similarity_common.py b/tests/e2e/accuracy/wan22_i2v/wan22_i2v_video_similarity_common.py index 0243754a6..bb6211824 100644 --- a/tests/e2e/accuracy/wan22_i2v/wan22_i2v_video_similarity_common.py +++ b/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度角倾泻而下，在兔脸右侧形成柔和的伦勃朗式阴影，强化了面部立体感和皮毛质感。 diff --git a/tests/e2e/offline_inference/test_gepard_tts.py b/tests/e2e/offline_inference/test_gepard_tts.py index b201842b4..2dec32ffc 100644 --- a/tests/e2e/offline_inference/test_gepard_tts.py +++ b/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 diff --git a/tests/e2e/offline_inference/test_minicpmo_4_5.py b/tests/e2e/offline_inference/test_minicpmo_4_5.py index 5a518a7f7..916d1b2b5 100644 --- a/tests/e2e/offline_inference/test_minicpmo_4_5.py +++ b/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.""" diff --git a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py index b1f17356b..b10d84d1a 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py +++ b/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.""" diff --git a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py index ad6935e71..9502ed5e6 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py +++ b/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. diff --git a/tests/e2e/online_serving/test_minicpmo_4_5.py b/tests/e2e/online_serving/test_minicpmo_4_5.py index cb8995206..290064fc3 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5.py +++ b/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: """ diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py index e7fa50c16..44b959b63 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py +++ b/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: """ diff --git a/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py b/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py index acdffd20d..6f97eaee2 100644 --- a/tests/e2e/online_serving/test_nemotron_voicechat_duplex.py +++ b/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", [ diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 10c5dd8cb..983e2cb7f 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/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: """ diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index 9fa313409..a4e641a57 100644 --- a/tests/helpers/__init__.py +++ b/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 diff --git a/tests/helpers/mark.py b/tests/helpers/mark.py index 18bdbd394..3b07706b2 100644 --- a/tests/helpers/mark.py +++ b/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) - test_resource = getattr(pytest.mark, res) - test_platform_detail = pytest.mark.cuda - marks = [test_resource, test_platform_detail, _cards_mark(num_cards)] + skus = _as_sku_list(res) + for sku in skus: + _require_sku("cuda", sku) + 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 = { - "cuda": _cuda_marks, - "rocm": _rocm_marks, - "xpu": _xpu_marks, - "musa": _musa_marks, - "npu": _npu_marks, - } - builder = builders.get(platform) - if builder is None: - raise ValueError(f"Unsupported platform: {platform}") - marks = builder(res=resource, num_cards=num_cards) + if platform == "cuda": + marks = _cuda_marks(res=resource, num_cards=num_cards) + else: + if not isinstance(resource, str): + raise ValueError(f"{platform} resource must be a string, got {resource!r}") + builders = { + "rocm": _rocm_marks, + "xpu": _xpu_marks, + "musa": _musa_marks, + "npu": _npu_marks, + } + 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): diff --git a/tests/helpers/tests/test_mark.py b/tests/helpers/tests/test_mark.py index d2fed4b6d..8a7770fc1 100644 --- a/tests/helpers/tests/test_mark.py +++ b/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"]}) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 906b27919..4f01022f8 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/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( - dim=256, - depth=2, - heads=4, - dim_head=64, - dropout=0.0, - ff_mult=2, - mel_dim=80, - mu_dim=80, - spk_dim=80, - long_skip_connection=True, - ) + with _force_torch_sdpa(): + return DiT( + dim=256, + depth=2, + heads=4, + dim_head=64, + dropout=0.0, + ff_mult=2, + mel_dim=80, + mu_dim=80, + 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( - num_heads=8, - head_size=64, - causal=False, - softmax_scale=1.0 / 8.0, - ) + with _force_torch_sdpa(): + attn = Attention( + num_heads=8, + head_size=64, + 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) diff --git a/tools/nightly/buildkite_testcase_statistics.py b/tools/nightly/buildkite_testcase_statistics.py index 06f09c791..9093762df 100644 --- a/tools/nightly/buildkite_testcase_statistics.py +++ b/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. diff --git a/tools/nightly/run_nightly_jobs.sh b/tools/nightly/run_nightly_jobs.sh index b5b6b92ce..48bc030ca 100644 --- a/tools/nightly/run_nightly_jobs.sh +++ b/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 diff --git a/tools/pre_commit/shellcheck.sh b/tools/pre_commit/shellcheck.sh index 7794ffe51..f653966dc 100755 --- a/tools/pre_commit/shellcheck.sh +++ b/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) diff --git a/tools/run_jobs_common.sh b/tools/run_jobs_common.sh index 7aa0861e4..0abcc9251 100644 --- a/tools/run_jobs_common.sh +++ b/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 diff --git a/tools/run_merge_jobs.sh b/tools/run_merge_jobs.sh index 9f9ab13b0..45f78a6bc 100644 --- a/tools/run_merge_jobs.sh +++ b/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 diff --git a/tools/run_ready_jobs.sh b/tools/run_ready_jobs.sh index b854faa07..25c890cfb 100644 --- a/tools/run_ready_jobs.sh +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index c0b9a738e..5b5b8a328 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py index ab814f713..ef239a82e 100644 --- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py +++ b/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 - # up to 256 for BF16/FP16. Empty list = "accept any"; the sdpa_kernel - # fallback chain handles shapes cuDNN can't take. - return [] + return list( + range( + CuDNNAttentionBackend._HEAD_SIZE_MULTIPLE, + 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, - # e.g. under torch.compile where Dynamo sees a symbolic head_dim and - # cuDNN's kernel selection fails for some symbolic attention shapes. - # The unpinned dispatcher then picks EFFICIENT/MATH instead of raising. - try: - with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]): - 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, + # Automatic-only: cuDNN FMHA rejects KV sequence length 1 + # ("cudnn SDPA does not support key/value sequence length 1"). Dummy + # warmup and some I2V layers (e.g. LTX-2) hit this when CUDNN_ATTN is + # the platform default. MATH is the only kernel that can run that shape. + # An explicit CUDNN_ATTN request raises instead of silently switching. + if kv_seq_len <= 1: + if self.backend_explicit: + raise ValueError( + "CUDNN_ATTN was explicitly selected but cuDNN FMHA does not " + f"support key/value sequence length {kv_seq_len}. " + "Select TORCH_SDPA or another backend for this shape." ) - except RuntimeError as e: - if "No available kernel" not in str(e): - raise - logger.warning_once( - "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." - ) + backends = [SDPBackend.MATH] + else: + backends = [SDPBackend.CUDNN_ATTENTION] + with sdpa_kernel(backends): output = torch.nn.functional.scaled_dot_product_attention( query, key, diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index a376216fa..cef6a97f5 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/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") - preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1 + input_scale_raw = extra.get("laser_input_scale") + 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 diff --git a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py index 73faa57a9..7f9c78bf7 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py +++ b/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_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None) + self.flash_attn_func: Callable[..., Any] | None = getattr(hub_module, "flash_attn_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_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None) + self.flash_attn_func: Callable[..., Any] | None = getattr(hub_module, "flash_attn_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): diff --git a/vllm_omni/diffusion/attention/backends/flashinfer_attn.py b/vllm_omni/diffusion/attention/backends/flashinfer_attn.py index 188875818..33b4a819e 100644 --- a/vllm_omni/diffusion/attention/backends/flashinfer_attn.py +++ b/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: - # Non-CuTe batch-prefill backends accept boolean custom masks. CuTe-DSL - # falls back to SDPA for nontrivial masks. + def supports_attention_mask(cls, attention_spec: object | None = None) -> bool: + # fa2/fa3 batch-prefill accept boolean custom masks. cute-dsl does not: + # 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 - SDPA, which applies them with the correct softmax semantics. Shape - mismatches also raise ``ValueError``. + additive/float masks and shape mismatches raise ``ValueError`` because + an explicitly selected backend must not silently switch to SDPA. """ 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( - f"non-boolean attn_mask (dtype={mask.dtype}); FlashInfer custom_mask " - "is boolean-only — deferring to SDPA" - ) + raise ValueError(f"non-boolean attn_mask (dtype={mask.dtype}); FlashInfer custom_mask is boolean-only") # 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: - logger.debug("Falling back to SDPA for mask path: %s", error) - return self._sdpa_fallback(query, key, value, attn_metadata) + except ValueError: + if self.backend_explicit: + 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") - return self._sdpa_fallback(query, key, value, attn_metadata) + if self.backend_explicit: + 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) diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py index 25f9a41cc..3ad60fc96 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py index e5f578821..5acf50651 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_utils.py b/vllm_omni/diffusion/attention/backends/ring/ring_utils.py index 67f71562b..b8743fe13 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_utils.py +++ b/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) - - B, S, H, D = out.shape - - # --- Shape Correction Logic for block_lse --- - # Goal: block_lse should be (B, S, H, 1) to match out (B, S, H, D) - - # 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 out.shape != block_out.shape: + raise ValueError( + f"Ring attention block output shape {tuple(block_out.shape)} does not match " + f"accumulated output shape {tuple(out.shape)}." + ) + block_lse = _normalize_lse(block_lse, block_out, lse_layout) if lse.shape != block_lse.shape: - # If lse was initialized with wrong shape, try to fix it - if lse.dim() == 4 and lse.shape[1] == block_lse.shape[2] and lse.shape[2] == block_lse.shape[1]: - lse = lse.transpose(1, 2) - elif lse.shape[1] >= S: # slice if lse was initialized with padding - lse = lse[:, :S, :, :] + raise ValueError( + f"Ring attention block LSE shape {tuple(block_lse.shape)} does not match " + f"accumulated LSE shape {tuple(lse.shape)}." + ) - # Final check - if lse.shape != block_lse.shape: - # 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 + out = out - F.sigmoid(block_lse - lse) * (out - block_out) + lse = lse - F.logsigmoid(lse - block_lse) 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) - - # 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 + lse = _normalize_lse(block_lse, out, lse_layout) 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 diff --git a/vllm_omni/diffusion/attention/backends/ring_flash_attn.py b/vllm_omni/diffusion/attention/backends/ring_flash_attn.py index dd27f88a8..1fce4f1b4 100644 --- a/vllm_omni/diffusion/attention/backends/ring_flash_attn.py +++ b/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() diff --git a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py b/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py index 482cbc9f8..7ab60b784 100644 --- a/vllm_omni/diffusion/attention/backends/ring_pytorch_attn.py +++ b/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() diff --git a/vllm_omni/diffusion/attention/backends/sage_attn.py b/vllm_omni/diffusion/attention/backends/sage_attn.py index 37eab409a..a0369c0f3 100644 --- a/vllm_omni/diffusion/attention/backends/sage_attn.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/backends/sage_attn3.py b/vllm_omni/diffusion/attention/backends/sage_attn3.py index 51ac3a48a..3188eda74 100644 --- a/vllm_omni/diffusion/attention/backends/sage_attn3.py +++ b/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: - logger.warning("SageAttention3 does not support GQA/MQA (Hq != Hkv); falling back to torch SDPA.") - type(self)._warned_gqa_fallback_global = True - output = F.scaled_dot_product_attention( - query, - key, - value, - is_causal=self.causal, - dropout_p=self.dropout, - scale=self.softmax_scale, - enable_gqa=True, + raise NotImplementedError( + "SAGE_ATTN_3 was explicitly selected but does not support GQA/MQA " + f"(q_heads={query.shape[1]}, kv_heads={key.shape[1]}). Select a GQA-capable backend." ) - 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() diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py index aec8eb50b..605a33402 100644 --- a/vllm_omni/diffusion/attention/backends/sdpa.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index d7a3a703f..54a27e217 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/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 -flash_attn_func = None -flash_attn_varlen_func = None +@lru_cache(maxsize=1) +def is_flash_attn_4_available() -> bool: + """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 - # Blackwell, where an importable Hopper-only FA3 wheel may otherwise be - # selected and fail at launch with "no kernel image is available". + # CUDA implementation candidates, ordered by preference. + # Candidate 1: CuTe-based FA4. This matters on Blackwell, where an + # importable Hopper-only FA3 wheel may otherwise be selected and fail + # at launch with "no kernel image is available". if _is_blackwell(): - try: - from flash_attn.cute import flash_attn_func, flash_attn_varlen_func # noqa: F401 + if is_flash_attn_4_available(): + 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") - except Exception as exc: - # Optional FA4 dependencies may be present but ABI-incompatible - # (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) + flash_attn_func = _fa_func + flash_attn_varlen_func = _fa_varlen + logger.info("Using CuTe FlashAttention-4 on Blackwell") + except Exception as exc: + # Optional FA4 dependencies may be present but ABI-incompatible + # (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. - # Try FA3 from fa3-fwd PyPI package - if flash_attn_func is None: - try: - from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 - except (ImportError, ModuleNotFoundError): - pass + flash_attn_func = _fa_func + flash_attn_varlen_func = _fa_varlen + except (ImportError, ModuleNotFoundError): + pass - # Fallback: Try FA3 from flash-attention source build - if flash_attn_func is None: - try: - from flash_attn_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 - except (ImportError, ModuleNotFoundError): - pass + # Candidate 3: FA3 from a flash-attention source build. + if flash_attn_func is None: + try: + from flash_attn_interface import flash_attn_func as _fa_func + from flash_attn_interface import flash_attn_varlen_func as _fa_varlen - # Fallback: Try FA2 from flash-attn package (try multiple import paths) - if flash_attn_func is None: - try: - from flash_attn import flash_attn_func, flash_attn_varlen_func # noqa: F401 - except (ImportError, ModuleNotFoundError): - pass + flash_attn_func = _fa_func + flash_attn_varlen_func = _fa_varlen + except (ImportError, ModuleNotFoundError): + pass - if flash_attn_func is None: - try: - from flash_attn.flash_attn_interface import ( # noqa: F401 - flash_attn_func, - flash_attn_varlen_func, - ) - except (ImportError, ModuleNotFoundError): - pass + # Candidate 4: FA2 from the flash-attn package (multiple import paths). + if flash_attn_func is None: + try: + from flash_attn import flash_attn_func as _fa_func + from flash_attn import flash_attn_varlen_func as _fa_varlen - # Fallback: Try vLLM's encapsulated flash attention dispatcher - # TODO discuss priority and remove potentially redundant fallbacks - if flash_attn_varlen_func is None: - try: - from vllm.vllm_flash_attn import ( # noqa: F401 - flash_attn_varlen_func, - ) - except (ImportError, ModuleNotFoundError): - pass + flash_attn_func = _fa_func + flash_attn_varlen_func = _fa_varlen + except (ImportError, ModuleNotFoundError): + pass + + if flash_attn_func is None: + try: + from flash_attn.flash_attn_interface import flash_attn_func as _fa_func + 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: - 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 is_flash_attn_4_available() # 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 diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 7dca3c1e6..3b2d76ba1 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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 - # package is absent), but a paged layer must advertise the + # attention to CUDNN/SDPA (for example when optional dense FA + # 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 - # attention keeps the platform default unchanged. - from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec + # Load that selector from the registry: this is not a user-explicit + # dense FLASH_ATTN request, so skip platform explicit-backend + # 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, 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, - ) + attn_backend_cls = DiffusionAttentionBackendEnum.FLASH_ATTN.get_class() 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() - self.ring_pg = sp_group.ring_group - self.ring_runner = RingParallelAttention( - sp_group, - attn_backend_pref=self.backend_pref, - ) - except Exception: - self.use_ring = False - self.ring_runner = None + sp_group = get_sp_group() + self.ring_pg = sp_group.ring_group + self.ring_runner = RingParallelAttention( + sp_group, + attn_backend_pref=self.backend_pref, + attn_backend_explicit=self.backend_explicit, + ) 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( - "Attention backend %s does not support kv_cache_dtype='%s' on %s. " - "KV quantization will be disabled.", - self.attn_backend.get_name(), - dtype, - platform_key, + raise ValueError( + f"Attention backend {self.attn_backend.get_name()} does not support " + f"kv_cache_dtype={dtype!r} on {platform_key}. Select a compatible " + "backend or set diffusion_kv_cache_dtype='auto'." ) - 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) - - 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) + self._assert_metadata_compatible(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 - # paged attention result. If dense MindIE-SD is absent, SDPA provides - # that profile forward. Formal GPU/NPU paged requests never use this - # branch because their Worker adapter is active. + # The startup KV-capacity profile needs tensor shapes, not a paged + # attention result. If dense FLASH_ATTN deps are absent (NPU MindIE-SD + # or CUDA CuTe FA4), SDPA provides that profile forward. Formal paged + # 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: - if attn_metadata is None or attn_metadata.full_attn_spans is None: + def _assert_metadata_compatible(self, attn_metadata: AttentionMetadata | None) -> 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( diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py index 5b644ca45..a7e41575b 100644 --- a/vllm_omni/diffusion/attention/parallel/factory.py +++ b/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: - 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}") + if not is_forward_context_available(): 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 - if ulysses_degree > 1 or ring_degree > 1 or allgather_degree > 1: - logger.warning( - f"SP configured (ulysses={ulysses_degree}, ring={ring_degree}, " - f"allgather={allgather_degree}) but SP group not available: {e}. " - f"Falling back to NoParallelAttention. This may cause incorrect results." - ) - return NoParallelAttention() + raise RuntimeError( + f"SP is configured (ulysses={ulysses_degree}, ring={ring_degree}, " + f"allgather={allgather_degree}), but the SP group is unavailable." + ) from e + if get_sequence_parallel_world_size() <= 1: + raise RuntimeError( + f"SP is configured (ulysses={ulysses_degree}, ring={ring_degree}, " + f"allgather={allgather_degree}), but the initialized SP world size is not greater than one." + ) if allgather_degree > 1: if causal: diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py index 74c6fc3fb..d71359ae7 100644 --- a/vllm_omni/diffusion/attention/parallel/ring.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/selector.py b/vllm_omni/diffusion/attention/selector.py index 9d357b4f6..cafe4da46 100644 --- a/vllm_omni/diffusion/attention/selector.py +++ b/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) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 4c9c37c35..d069d95d1 100644 --- a/vllm_omni/diffusion/data.py +++ b/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}: - raise ValueError( - f"block_sparse.precision must be one of {sorted(p.value for p in RainFusionPrecision)}; " - f"got {self.precision!r}." - ) + precision = self.precision.value if isinstance(self.precision, RainFusionPrecision) else self.precision + valid = {member.value for member in RainFusionPrecision} + if precision not in valid: + 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] = { - "dtype_qk": q.dtype_qk, - "q_block_size": q.q_block_size, - "k_block_size": q.k_block_size, - } - if q.dtype_vo is not None: - quant_kw["dtype_vo"] = q.dtype_vo + quant_kw: dict[str, Any] = {} + if q.dtype_qk is not None or q.dtype_vo is not None: + quant_kw["dtype_qk"] = q.dtype_qk + 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: + 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: diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index e8f88c8a9..20800ffc0 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/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, - args: tuple = (), - kwargs: dict | None = None, + identifier: str | int, + args: Sequence[Any] = (), + 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, - kwargs: dict, + args: Sequence[Any], + 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 = (), - kwargs: dict | None = None, + args: Sequence[Any] = (), + 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) diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index 00c2399fa..59140073d 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/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.") - attention_backend_attempts.append("native") + raise ValueError(f"Invalid Diffusers attention backend: {attention_backend_config}.") 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 " diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 4551126a7..c635877a9 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/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 - flash_attn_supported = compute_supported and packages_available + # FA2/FA3 wheels are often importable on Blackwell but only ship + # 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 - # without the tuned Blackwell path, so we skip routing there. + # skip the CUDNN_ATTN default rather than selecting a backend whose + # 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 - # take 2D attn_masks and the SDPA fallback dispatches to - # EFFICIENT_ATTENTION (~25 ms) instead of the cuDNN mask path (~11 ms). + # regresses ~2x at e2e on HV-1.5 because its dense-prefill path cannot + # represent every diffusion attention mask without changing semantics. # 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: - if backend_upper == "FLASH_ATTN_HUB": - logger.warning( - "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN." - ) - backend_upper = "FLASH_ATTN" - elif backend_upper == "FLASH_ATTN_3_HUB": - logger.warning( - "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN." - ) - backend_upper = "FLASH_ATTN" + except ImportError as e: + raise ImportError( + f"{backend_upper} was explicitly selected, but the HuggingFace `kernels` " + "library is not available. Install `kernels` or select a different backend." + ) from e + + if backend_upper == "FLASH_ATTN_HUB": + fa2_hub_supported = compute_capability is not None and compute_capability.major in (8, 9) + if not fa2_hub_supported: + raise ValueError( + "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( - "FLASH_ATTN_3_HUB requires a Hopper-class GPU with compute capability >= 9.0. " - "Falling back to FLASH_ATTN_HUB." + raise ValueError( + "FLASH_ATTN_3_HUB was explicitly selected but its current kernels require a Hopper GPU " + "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: - logger.warning( - "Flash Attention requires GPU with compute capability >= 8.0. " - "Falling back to TORCH_SDPA backend." + if is_blackwell: + reason = ( + "Blackwell requires CuTe FlashAttention-4 " + "(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only" ) - elif not packages_available: - logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.") - logger.debug("Defaulting to diffusion attention backend SDPA") - return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() - if backend_upper == "SAGE_ATTN_3": - sage_attn3_supported = compute_capability is not None and compute_capability.major >= 10 - if not sage_attn3_supported: - logger.warning( - "SageAttention3 requires a Blackwell-class GPU with compute capability >= 10.0. " - "Falling back to TORCH_SDPA backend." + elif not compute_supported: + reason = "compute capability < 8.0" + else: + reason = "Flash Attention package unavailable" + raise ValueError( + f"FLASH_ATTN was explicitly selected but is unsupported ({reason}). Select a compatible backend." + ) + if backend_upper == "FLASHINFER_ATTN" and not flashinfer_available: + raise ValueError( + "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() - try: - importlib.import_module("sageattn3") - except ImportError: - logger.warning( - "SageAttention3 package not available. Install it from " - "SageAttention/sageattention3_blackwell. Falling back to TORCH_SDPA backend." + if backend_upper == "CUDNN_ATTN": + from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend + + # head_size <= 0 is the capability-probe sentinel (auto-pad does + # not know head_dim). Skip geometry validation; the real + # Attention construction still checks the configured size. + 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( - "Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell %s, cuDNN %d)", - sm_str, - cudnn_version, - ) - return DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path() + from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend - 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, )