Compare commits

..
Author SHA1 Message Date
hsliu_ustc 7d41ea6183 Merge remote-tracking branch 'origin/main' into codex-api-server-user-guide
Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>

# Conflicts:
#	docs/user_guide/examples/online_serving/text_to_audio.md
#	examples/online_serving/stable_audio/README.md
2026-09-01 11:23:08 +08:00
Hongsheng Liu a79dcdee5c docs: address API guide review feedback
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-17 13:31:36 +08:00
Hongsheng Liu 159ac2161a docs: standardize API example ports
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 23:34:28 +08:00
Hongsheng Liu 8327f26216 docs: document realtime and specialized API servers
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 21:46:24 +08:00
Hongsheng Liu 5c6ed9b4a9 docs: add API server endpoint guide
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-13 21:28:36 +08:00
259 changed files with 2191 additions and 20136 deletions

No files matched your search

+7 -4
View File
@@ -238,10 +238,13 @@ steps:
grade: Blocking grade: Blocking
timeout_in_minutes: 50 timeout_in_minutes: 50
commands: commands:
- export VLLM_LOGGING_LEVEL=DEBUG - |
- export VLLM_WORKER_MULTIPROC_METHOD=spawn timeout 40m bash -c '
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" export VLLM_LOGGING_LEVEL=DEBUG
- timeout 40m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model and cuda" --run-level "advanced_model" export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
'
# TODO: Bagel test on ROCm is very unstable. @tjtanaa # TODO: Bagel test on ROCm is very unstable. @tjtanaa
# Need to debug before reneable numerical changes across large PRs # Need to debug before reneable numerical changes across large PRs
+30 -432
View File
@@ -1,31 +1,19 @@
# Uploader-only hardware presets for CUDA / NPU test YAML. # Uploader-only hardware presets for CUDA / NPU test YAML.
# Referenced as ``mirror_hardwares: l4_1`` on a step; ``upload_pipeline.py`` expands to # Referenced as ``mirror_hardwares: l4_1`` on a step; ``upload_pipeline.py`` expands to
# ``agents`` (+ optional top-level ``image`` for NPU) + ``plugins`` (+ ``retry`` # ``agents`` (+ optional top-level ``image`` for NPU) + ``plugins``, then strips
# for ``l4_*``), then strips ``mirror_hardwares`` before Buildkite upload. # ``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} # CUDA naming: {chip}_{count}
# l4_1, l4_2, l4_3, l4_4, h100_1 … h100_4, b200_1 … b200_4 # l4_1, l4_4, h100_1 … h100_4
# NPU naming: {queue}_{resource_class} # NPU naming: {queue}_{resource_class}
# a2b3_npu_1 — ascend-a2b3, resource_class npu-1 (A2/B3 CI image) # a2b3_npu_1 — ascend-a2b3, resource_class npu-1 (A2/B3 CI image)
# a2b3_npu_4 — ascend-a2b3, resource_class npu-4 # a2b3_npu_4 — ascend-a2b3, resource_class npu-4
# a3_npu_2 — ascend-a3, resource_class npu-2 (A3 CI image tag …:$BUILDKITE_COMMIT-a3) # a3_npu_2 — ascend-a3, resource_class npu-2
# #
# B200 (b200-k8s): matches vLLM device ``b200-k8s`` from ci-infra # The ``l4_*`` presets run under the ``docker`` plugin, which does not put
# (AgentQueue.B200_K8S + b200_plugin_template). Queue routes the K8s pool; # ``buildkite-agent`` on PATH inside the container the way the ``kubernetes``
# podSpec uses hostNetwork + /raid|/mnt/shared (no H100-style nodeSelector). # presets do, so they set ``mount-buildkite-agent: true``; without it any step
# # calling ``buildkite-agent`` (artifact upload, annotate, meta-data) fails there.
# The ``l4_*`` presets run on the EKS ``l4-k8s`` queue (agent-stack-k8s),
# packed onto g6.12xlarge L4 nodes (``vllm.ci/gpu-pool=l4x4``). GPU / CPU /
# memory / shm requests scale with card count (1 GPU = 10 CPU / 40Gi / 24Gi shm).
# Retry covers lost agents / killed pods, not pytest or OOM (exit 1 / 137).
x-npu-k8s-plugins: &npu_k8s_plugins x-npu-k8s-plugins: &npu_k8s_plugins
- kubernetes: - kubernetes:
@@ -33,224 +21,39 @@ x-npu-k8s-plugins: &npu_k8s_plugins
imagePullSecrets: imagePullSecrets:
- name: swr-secret - name: swr-secret
# Infra-only retry. Omit exit 1 (pytest failed) and 137 (cgroup OOM / SIGKILL
# of the test process). ci-infra ``K8S_RETRY`` still retries exit 1.
x-l4-k8s-retry: &l4_k8s_retry
automatic:
- exit_status: -1
limit: 1
- exit_status: 128
limit: 1
- signal_reason: agent_stop
limit: 1
- signal_reason: agent_refused
limit: 1
mirror_hardwares: mirror_hardwares:
l4_1: l4_1:
agents: agents:
queue: l4-k8s queue: gpu_1_queue
retry: *l4_k8s_retry
plugins: plugins:
- kubernetes: - docker#v5.2.0:
metadata: image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
annotations: always-pull: true
cluster-autoscaler.kubernetes.io/safe-to-evict: "false" propagate-environment: true
karpenter.sh/do-not-disrupt: "true" mount-buildkite-agent: true
podSpec: shm-size: "8gb"
priorityClassName: ci environment:
serviceAccountName: buildkite-gpu-jobs - HF_HOME=/fsx/hf_cache
tolerations: - HF_TOKEN
- key: nvidia.com/gpu
operator: Exists
effect: NoSchedule
nodeSelector:
vllm.ci/gpu-pool: l4x4
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
requests:
nvidia.com/gpu: 1
cpu: "10"
memory: 40Gi
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /fsx/hf_cache
env:
- name: HF_HOME
value: /fsx/hf_cache
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
volumes: volumes:
- name: devshm - /fsx/hf_cache:/fsx/hf_cache
emptyDir:
medium: Memory
sizeLimit: 24Gi
- name: hf-cache
hostPath:
path: /fsx/hf_cache
type: DirectoryOrCreate
l4_2:
agents:
queue: l4-k8s
retry: *l4_k8s_retry
plugins:
- kubernetes:
metadata:
annotations:
cluster-autoscaler.kubernetes.io/safe-to-evict: "false"
karpenter.sh/do-not-disrupt: "true"
podSpec:
priorityClassName: ci
serviceAccountName: buildkite-gpu-jobs
tolerations:
- key: nvidia.com/gpu
operator: Exists
effect: NoSchedule
nodeSelector:
vllm.ci/gpu-pool: l4x4
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
requests:
nvidia.com/gpu: 2
cpu: "20"
memory: 80Gi
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /fsx/hf_cache
env:
- name: HF_HOME
value: /fsx/hf_cache
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
volumes:
- name: devshm
emptyDir:
medium: Memory
sizeLimit: 48Gi
- name: hf-cache
hostPath:
path: /fsx/hf_cache
type: DirectoryOrCreate
l4_3:
agents:
queue: l4-k8s
retry: *l4_k8s_retry
plugins:
- kubernetes:
metadata:
annotations:
cluster-autoscaler.kubernetes.io/safe-to-evict: "false"
karpenter.sh/do-not-disrupt: "true"
podSpec:
priorityClassName: ci
serviceAccountName: buildkite-gpu-jobs
tolerations:
- key: nvidia.com/gpu
operator: Exists
effect: NoSchedule
nodeSelector:
vllm.ci/gpu-pool: l4x4
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 3
requests:
nvidia.com/gpu: 3
cpu: "30"
memory: 120Gi
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /fsx/hf_cache
env:
- name: HF_HOME
value: /fsx/hf_cache
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
volumes:
- name: devshm
emptyDir:
medium: Memory
sizeLimit: 72Gi
- name: hf-cache
hostPath:
path: /fsx/hf_cache
type: DirectoryOrCreate
l4_4: l4_4:
agents: agents:
queue: l4-k8s queue: gpu_4_queue
retry: *l4_k8s_retry
plugins: plugins:
- kubernetes: - docker#v5.2.0:
metadata: image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
annotations: always-pull: true
cluster-autoscaler.kubernetes.io/safe-to-evict: "false" propagate-environment: true
karpenter.sh/do-not-disrupt: "true" mount-buildkite-agent: true
podSpec: shm-size: "8gb"
priorityClassName: ci environment:
serviceAccountName: buildkite-gpu-jobs - HF_HOME=/fsx/hf_cache
tolerations: - HF_TOKEN
- key: nvidia.com/gpu
operator: Exists
effect: NoSchedule
nodeSelector:
vllm.ci/gpu-pool: l4x4
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 4
requests:
nvidia.com/gpu: 4
cpu: "40"
memory: 160Gi
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /fsx/hf_cache
env:
- name: HF_HOME
value: /fsx/hf_cache
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
volumes: volumes:
- name: devshm - /fsx/hf_cache:/fsx/hf_cache
emptyDir:
medium: Memory
sizeLimit: 96Gi
- name: hf-cache
hostPath:
path: /fsx/hf_cache
type: DirectoryOrCreate
h100_1: h100_1:
agents: agents:
@@ -392,211 +195,6 @@ mirror_hardwares:
path: /mnt/hf-cache path: /mnt/hf-cache
type: DirectoryOrCreate 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: a2b3_npu_1:
agents: agents:
queue: "ascend-a2b3" queue: "ascend-a2b3"
+36 -3
View File
@@ -1,6 +1,5 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0 # 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. """Shared Buildkite skip-ci and CI-YAML level targeting logic.
``resolve_ci_decision(changed_files, ...)`` buckets explicit file lists (tests, dry-run). ``resolve_ci_decision(changed_files, ...)`` buckets explicit file lists (tests, dry-run).
@@ -11,6 +10,12 @@ CLI (for bootstrap scripts; exit code is the signal):
python3 skip_ci.py gate <platform> <level> python3 skip_ci.py gate <platform> <level>
Exit 0 when bootstrap should stop (skip-all, or that L2/L3 target is off). 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. Prints ``skip-all`` or ``skip-l23`` on stdout; logs the decision once.
python3 skip_ci.py check-skip-all
Exit 0 when only docs/skip-mark changes → bootstrap skips entire CI upload.
python3 skip_ci.py check-skip-l2-l3 <platform> <level>
Exit 0 when that L2/L3 pipeline should be skipped (CI config YAML-only diffs).
python3 skip_ci.py print-annotate
Print decision message for Buildkite annotation.
""" """
from __future__ import annotations from __future__ import annotations
@@ -253,10 +258,11 @@ class CiDecision:
@dataclass(frozen=True) @dataclass(frozen=True)
class CiContext: class CiContext:
"""Skip-ci decision plus changed files used for pipeline filtering.""" """Skip-ci decision plus git diff inputs used to compute it."""
decision: CiDecision decision: CiDecision
changed_files: list[str] | None changed_files: list[str] | None
diff_range: str | None
def _log(message: str) -> None: def _log(message: str) -> None:
@@ -326,7 +332,7 @@ def resolve_ci_context_from_git() -> CiContext:
diff_range = _resolve_diff_range() diff_range = _resolve_diff_range()
changed_files = _changed_files_for_diff_range(diff_range) changed_files = _changed_files_for_diff_range(diff_range)
decision = resolve_ci_decision(changed_files, diff_range=diff_range) decision = resolve_ci_decision(changed_files, diff_range=diff_range)
return CiContext(decision=decision, changed_files=changed_files) return CiContext(decision=decision, changed_files=changed_files, diff_range=diff_range)
# --- Doc / skip-mark helpers --- # --- Doc / skip-mark helpers ---
@@ -565,6 +571,23 @@ def main() -> int:
gate_parser.add_argument("platform", help="cuda, amd, intel, or npu") gate_parser.add_argument("platform", help="cuda, amd, intel, or npu")
gate_parser.add_argument("level", choices=("l2", "l3"), help="test level") 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() args = parser.parse_args()
decision = resolve_ci_context_from_git().decision decision = resolve_ci_context_from_git().decision
@@ -577,6 +600,16 @@ def main() -> int:
return 0 return 0
return 1 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 return 1
+31 -275
View File
@@ -12,19 +12,8 @@ Bootstrap mode (``bootstrap-upload-steps.yml``):
Test pipeline mode (e.g. test-merge.yml): Test pipeline mode (e.g. test-merge.yml):
- Drop steps whose ``source_file_dependencies`` do not match changed files. - Drop steps whose ``source_file_dependencies`` do not match changed files.
- Expand uploader-only ``mirror_hardwares`` into ``agents`` (+ optional ``image`` - Expand uploader-only ``mirror_hardwares: l4_1`` into ``agents`` (+ optional ``image``
for NPU) + ``plugins`` (see ci_mirror_hardwares.yml). 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: Usage:
python3 upload_pipeline.py [--upload] [--all | --e2e] <pipeline.yml> python3 upload_pipeline.py [--upload] [--all | --e2e] <pipeline.yml>
@@ -36,13 +25,11 @@ from __future__ import annotations
import argparse import argparse
import copy import copy
import os
import re
import subprocess import subprocess
import sys import sys
from functools import lru_cache from functools import lru_cache
from pathlib import Path from pathlib import Path
from typing import Any, Literal from typing import Any
try: try:
import yaml import yaml
@@ -58,14 +45,6 @@ from skip_ci import (
resolve_ci_context_from_git, 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 --- # --- Constants ---
LOG = "upload_pipeline" LOG = "upload_pipeline"
@@ -254,257 +233,32 @@ def _load_mirror_hardwares() -> dict[str, dict[str, Any]]:
return presets return presets
@lru_cache(maxsize=1) def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any]:
def _cuda_mirror_chips() -> tuple[str, ...]: """Replace uploader-only ``mirror_hardwares`` with preset fields from ci_mirror_hardwares.yml."""
"""Lowercase CUDA SKUs that have at least one ``{chip}_{n}`` preset. hardware = step.get("mirror_hardwares")
if hardware is None:
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"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",
)
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 return step
chips, cards = _read_pytest_marks(step.get("commands"))
if not chips and cards is None: if not isinstance(hardware, str) or not hardware.strip():
return step
preset_name = _compose_mirror_hardware_name(chips, cards, step_label=step_label)
else:
if has_pool:
raise ValueError( raise ValueError(
f"step {step_label!r} sets mirror_hardwares together with agents/plugins/image; " f"mirror_hardwares must be a non-empty string in step {_get_step_label(step)!r}",
f"use mirror_hardwares only",
)
preset_name = _resolve_mirror_hardware_name(
step.get("mirror_hardwares"),
step_label=step_label,
) )
if preset_name is None: preset = _load_mirror_hardwares().get(hardware)
return None
preset = _load_mirror_hardwares().get(preset_name)
if preset is None: if preset is None:
known = ", ".join(sorted(_load_mirror_hardwares())) known = ", ".join(sorted(_load_mirror_hardwares()))
raise ValueError( raise ValueError(
f"unknown mirror_hardwares {preset_name!r} in step {step_label!r}; known: {known}", 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",
) )
expanded = copy.deepcopy(preset) expanded = copy.deepcopy(preset)
merged = {key: value for key, value in step.items() if key != "mirror_hardwares"} | expanded return {key: value for key, value in step.items() if key != "mirror_hardwares"} | expanded
# Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out.
if "retry" in step:
merged["retry"] = step["retry"]
return merged
def _match_source_file(changed_files: list[str], prefixes: list[str]) -> bool: def _match_source_file(changed_files: list[str], prefixes: list[str]) -> bool:
@@ -543,7 +297,7 @@ def _process_test_steps(
nested = step.get("steps") nested = step.get("steps")
if nested is not None: if nested is not None:
kept_nested = _process_test_steps(nested, changed_files) kept_nested = _process_test_steps(nested, changed_files)
if not kept_nested: if changed_files is not None and not kept_nested:
_log(f"omit empty group {_get_step_label(step)!r}") _log(f"omit empty group {_get_step_label(step)!r}")
continue continue
new_step = {key: value for key, value in step.items() if key != "source_file_dependencies"} new_step = {key: value for key, value in step.items() if key != "source_file_dependencies"}
@@ -551,11 +305,14 @@ def _process_test_steps(
processed.append(new_step) processed.append(new_step)
continue continue
leaf = {key: value for key, value in step.items() if key != "source_file_dependencies"} if deps is not None:
expanded = _expand_mirror_hardwares(leaf) processed.append(
if expanded is None: _expand_mirror_hardwares(
continue {key: value for key, value in step.items() if key != "source_file_dependencies"},
processed.append(expanded) ),
)
else:
processed.append(_expand_mirror_hardwares(step))
return processed return processed
@@ -581,10 +338,6 @@ def _render_test_pipeline(
e2e_only: bool = False, e2e_only: bool = False,
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Filter steps by PR diff and strip uploader-only ``source_file_dependencies`` metadata.""" """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") steps = doc.get("steps")
if not isinstance(steps, list): if not isinstance(steps, list):
return doc return doc
@@ -614,7 +367,10 @@ def _render_pipeline(
text = path.read_text(encoding="utf-8") text = path.read_text(encoding="utf-8")
ctx = resolve_ci_context_from_git() ctx = resolve_ci_context_from_git()
changed_files = None if force_all or e2e_only else ctx.changed_files if force_all or e2e_only:
changed_files = None
else:
changed_files = ctx.changed_files
doc = yaml.safe_load(text) doc = yaml.safe_load(text)
if not isinstance(doc, dict): if not isinstance(doc, dict):
-38
View File
@@ -1,6 +1,3 @@
# 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: env:
VLLM_WORKER_MULTIPROC_METHOD: spawn VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_USE_FLASHINFER_MOE_FP16: 0 VLLM_USE_FLASHINFER_MOE_FP16: 0
@@ -240,38 +237,3 @@ steps:
pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
" "
mirror_hardwares: h100_1 mirror_hardwares: h100_1
- label: "Diffusion · MiniMax-H3 L3 DLO DP2 and FastH3"
source_file_dependencies:
- tests/e2e/online_serving/minimax_h3/
- tests/helpers/assertions.py
- tests/helpers/mark.py
- tests/helpers/media.py
- tests/helpers/runtime.py
- vllm_omni/diffusion/models/minimax_h3/
- vllm_omni/diffusion/lora/
- vllm_omni/diffusion/offloader/
- vllm_omni/diffusion/distributed/
- vllm_omni/entrypoints/openai/
# Keep four server lifecycles in one Buildkite job, but never in one
# pytest process. DLO uses multiprocessing shared-memory queues; the
# parent's resource_tracker can retain a stale /psm_* entry during
# teardown and strand the next Ref2VA collective.
# A per-file process plus a hard timeout also makes a regression fail
# instead of leaving the merge queue stuck indefinitely.
timeout_in_minutes: 150
commands:
- |
set -eu
export VLLM_IMAGE_FETCH_TIMEOUT=60
if [ -f vllm_omni/diffusion/models/minimax_h3/lora.py ]; then
export VLLM_TEST_MINIMAX_H3_TURBO_LORA=$(python3 -c 'from huggingface_hub import hf_hub_download; print(hf_hub_download(repo_id="lightx2v/Minimax-h3-Turbo", filename="minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"))')
fi
if [ -f vllm_omni/diffusion/models/minimax_h3/fasth3.py ]; then
export VLLM_TEST_MINIMAX_H3_FASTH3_LORA=$(python3 -c 'from huggingface_hub import hf_hub_download; print(hf_hub_download(repo_id="FastVideo/FastVideo-FastH3-4-step-Preview-v1-LoRA", filename="dense-datafree/adapter_model.safetensors"))')
fi
timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
timeout 35m pytest -s -v tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
mirror_hardwares: h100_2
+16 -45
View File
@@ -1,7 +1,3 @@
# 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: env:
VLLM_WORKER_MULTIPROC_METHOD: spawn VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_USE_FLASHINFER_MOE_FP16: 0 VLLM_USE_FLASHINFER_MOE_FP16: 0
@@ -17,14 +13,16 @@ steps:
steps: steps:
- label: ":full_moon: Omni · Function Test with H100 · Single-GPU" - label: ":full_moon: Omni · Function Test with H100 · Single-GPU"
timeout_in_minutes: 120 timeout_in_minutes: 90
commands: commands:
- 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 - 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
- label: ":full_moon: Omni · Function Test with H100 · 2-GPU" - label: ":full_moon: Omni · Function Test with H100 · 2-GPU"
timeout_in_minutes: 90 timeout_in_minutes: 90
commands: commands:
- 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 - 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
- label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test" - label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
timeout_in_minutes: 50 timeout_in_minutes: 50
@@ -32,11 +30,11 @@ steps:
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model" - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model"
mirror_hardwares: h100_1 mirror_hardwares: h100_1
- label: ":full_moon: Omni · Doc Test with L4 · 4-GPU" - label: ":full_moon: Omni · Doc Test with L4"
timeout_in_minutes: 90 timeout_in_minutes: 90
commands: commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/examples/ -m "full_model and omni and L4 and cards_4" --run-level "full_model" - pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model"
mirror_hardwares: l4_4 mirror_hardwares: l4_4
- label: ":full_moon: Omni · Doc Test with H100 · 2-GPU" - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU"
@@ -180,7 +178,7 @@ steps:
build.pull_request.labels includes "nightly-test" build.pull_request.labels includes "nightly-test"
steps: steps:
- label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 2-GPU" - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations)"
timeout_in_minutes: 30 timeout_in_minutes: 30
# NOTE: Single-GPU (core_model) tiny model tests run on the ready label, # NOTE: Single-GPU (core_model) tiny model tests run on the ready label,
# so nightly only adds the multi-GPU parallelism configurations. # so nightly only adds the multi-GPU parallelism configurations.
@@ -190,14 +188,7 @@ steps:
# for the CI, since these tests should pass irrespective of whether or not we use # for the CI, since these tests should pass irrespective of whether or not we use
# the full weights or not. # the full weights or not.
commands: commands:
- pytest -sv tests/model_tests/diffusion -m 'full_model and cuda and cards_2' --run-level "core_model" - timeout 25m pytest -sv tests/model_tests/diffusion -m 'full_model and cuda' --run-level "core_model"
mirror_hardwares: l4_2
- label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 4-GPU"
timeout_in_minutes: 30
# CFG+TP extra_test_groups request 2×2=4 devices (see get_required_device_count).
commands:
- pytest -sv tests/model_tests/diffusion -m 'full_model and cuda and cards_4' --run-level "core_model"
mirror_hardwares: l4_4 mirror_hardwares: l4_4
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test" - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
@@ -229,21 +220,10 @@ steps:
--run-level "full_model" --run-level "full_model"
mirror_hardwares: h100_2 mirror_hardwares: h100_2
- label: ":full_moon: Diffusion X2I(&A&T) · Boogu Edit CFG · 3-GPU" - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · Multi-GPU"
timeout_in_minutes: 120 timeout_in_minutes: 120
commands: commands:
- >- - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and not cards_1" --run-level "full_model" --ignore=tests/e2e/accuracy
pytest -sv
tests/e2e/online_serving/test_boogu_image_edit_expansion.py
-k "test_double_guidance_cfg_parallel and cfg_parallel_3_full"
-m "slow and diffusion and cuda and H100 and cards_3"
--run-level "full_model"
mirror_hardwares: h100_3
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · 4-GPU"
timeout_in_minutes: 120
commands:
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and cards_4" --run-level "full_model" --ignore=tests/e2e/accuracy
mirror_hardwares: l4_4 mirror_hardwares: l4_4
- label: ":full_moon: Diffusion X2I(&A&T) · Doc Test" - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test"
@@ -279,6 +259,7 @@ steps:
timeout_in_minutes: 180 timeout_in_minutes: 180
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- export CACHE_DIT_VERSION=1.5.0 - export CACHE_DIT_VERSION=1.5.0
- | - |
set +e set +e
@@ -381,6 +362,7 @@ steps:
timeout_in_minutes: 180 timeout_in_minutes: 180
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- | - |
set +e 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" 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"
@@ -430,6 +412,7 @@ steps:
timeout_in_minutes: 180 timeout_in_minutes: 180
commands: commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- | - |
set +e 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 pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json
@@ -444,22 +427,10 @@ steps:
depends_on: upload-nightly-pipeline depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
steps: steps:
- label: ":full_moon: Diffusion · Distributed Test with L4 · 2-GPU" - label: ":full_moon: Diffusion · Distributed Test with L4"
timeout_in_minutes: 60 timeout_in_minutes: 60
commands: commands:
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_2' --run-level "full_model" - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
mirror_hardwares: l4_2
- label: ":full_moon: Diffusion · Distributed Test with L4 · 3-GPU"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_3' --run-level "full_model"
mirror_hardwares: l4_3
- label: ":full_moon: Diffusion · Distributed Test with L4 · 4-GPU"
timeout_in_minutes: 60
commands:
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_4' --run-level "full_model"
mirror_hardwares: l4_4 mirror_hardwares: l4_4
- label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU" - label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU"
+10 -16
View File
@@ -1,6 +1,3 @@
# 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: env:
VLLM_WORKER_MULTIPROC_METHOD: spawn VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_USE_FLASHINFER_MOE_FP16: 0 VLLM_USE_FLASHINFER_MOE_FP16: 0
@@ -54,20 +51,17 @@ steps:
- timeout 15m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model" - timeout 15m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
mirror_hardwares: l4_1 mirror_hardwares: l4_1
- label: "Diffusion · Distributed Test · Single-GPU" - label: "Diffusion · Distributed & Attention Test"
source_file_dependencies:
- tests/diffusion/distributed/
- tests/diffusion/attention/test_attention_sp.py
- vllm_omni/diffusion/distributed/
- vllm_omni/diffusion/attention/
- vllm_omni/diffusion/data.py
- vllm_omni/diffusion/models/
commands: commands:
- timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_1' --run-level "core_model" - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda' --run-level "core_model"
mirror_hardwares: l4_1 - timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda' --run-level "core_model"
- label: "Diffusion · Distributed Test · 2-GPU"
commands:
- timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_2' --run-level "core_model"
mirror_hardwares: l4_2
- label: "Diffusion · Distributed & Attention Test · 4-GPU"
commands:
- timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_4' --run-level "core_model"
- timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda and L4 and cards_4' --run-level "core_model"
mirror_hardwares: l4_4 mirror_hardwares: l4_4
- label: "Diffusion · Model Test" - label: "Diffusion · Model Test"
+6 -25
View File
@@ -1,6 +1,3 @@
# CUDA CI L5 (`slow` / reliability / extra coverage).
# Uploaded by Upload Weekly Pipeline: WEEKLY=1 / NON_CRITICAL=1, or PR label `weekly-test`.
env: env:
VLLM_WORKER_MULTIPROC_METHOD: spawn VLLM_WORKER_MULTIPROC_METHOD: spawn
VLLM_USE_FLASHINFER_MOE_FP16: 0 VLLM_USE_FLASHINFER_MOE_FP16: 0
@@ -114,24 +111,17 @@ steps:
if: build.env("NON_CRITICAL") == "1" if: build.env("NON_CRITICAL") == "1"
steps: steps:
# Online Qwen2.5 expansion is 2-card; offline expansion + AutoRound are 4-card. # Qwen2.5 expansion (CUDA 4-card / online 2-card) stays on this l4_4 job.
- label: "Omni · L4 · 2-GPU" - label: "Omni · L4 · Multi-GPU"
timeout_in_minutes: 90 timeout_in_minutes: 90
commands: commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and cards_2" --run-level "core_model" - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and not cards_1" --run-level "core_model"
mirror_hardwares: l4_2
- label: "Omni · L4 · 4-GPU"
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py -m "slow and L4 and omni and cards_4" --run-level "core_model"
- >- - >-
pytest -s -v tests/e2e/ pytest -s -v tests/e2e/
--ignore=tests/e2e/online_serving/test_qwen2_5_omni_expansion.py --ignore=tests/e2e/online_serving/test_qwen2_5_omni_expansion.py
--ignore=tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py --ignore=tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
-m "slow and L4 and omni and cards_4" -m "slow and L4 and omni and not cards_1"
--run-level "full_model" --run-level "full_model"
mirror_hardwares: l4_4 mirror_hardwares: l4_4
@@ -204,21 +194,12 @@ steps:
--run-level "full_model" --run-level "full_model"
mirror_hardwares: l4_1 mirror_hardwares: l4_1
- label: "Diffusion · L4 · 2-GPU" - label: "Diffusion · L4 · Multi-GPU"
timeout_in_minutes: 60 timeout_in_minutes: 60
commands: commands:
- >- - >-
pytest -sv pytest -sv
tests/e2e/ -m "slow and diffusion and L4 and cards_2" tests/e2e/ -m "slow and diffusion and L4 and not cards_1"
--run-level "full_model"
mirror_hardwares: l4_2
- label: "Diffusion · L4 · 4-GPU"
timeout_in_minutes: 60
commands:
- >-
pytest -sv
tests/e2e/ -m "slow and diffusion and L4 and cards_4"
--run-level "full_model" --run-level "full_model"
mirror_hardwares: l4_4 mirror_hardwares: l4_4
@@ -1,458 +0,0 @@
---
name: production-add-diffusion-model
description: Productionize a vLLM-Omni diffusion model after its Day-0 vertical slice works. Use when work requires official API parity and input limits, feature-combination evidence, online FP8, distributed layerwise offload, per-request Cache-DiT, CUDA/ROCm/NPU/XPU recipes, faster or sparse attention, fused operators, USP or disaggregation analysis, continuous batching and abort handling, long-running RPS stability, or accuracy/performance/reliability CI. For initial architecture porting, registry wiring, and basic weight loading, use add-diffusion-model first.
---
# Productionizing a Diffusion Model
## Scope
Use this skill to move a working model integration from Day-0 support to a
measured, fail-closed production deployment. It is deliberately separate from
[`add-diffusion-model`](../add-diffusion-model/SKILL.md); do not modify or copy
that skill when productionizing a model. If registry, pipeline, or basic loader
work is incomplete, read that skill first and finish the Day-0 vertical slice.
MiniMax-H3 [PR #5691](https://github.com/vllm-project/vllm-omni/pull/5691)
is the merged Day-0 case study. Its open optimization
[issue #5700](https://github.com/vllm-project/vllm-omni/issues/5700) is an
actively maintained roadmap and evidence map, not a list of already supported
features. It can lag merged implementation: read its current revision, cited
PRs, target source, and maintained recipes, then leave every unverified row
`not tested` until the target revision has scoped evidence.
Read these references before the corresponding gate:
- [API and hardware recipes](references/api-and-recipes.md)
- [Production feature patterns](references/feature-patterns.md)
- [Performance patterns and code examples](references/performance-patterns.md)
- [Serving, soak, and CI validation](references/production-validation.md)
## Non-negotiable evidence contract
Create a matrix whose row key is at least:
```text
task x API mode x execution mode/capacity x shape/schedule x attention backend x
packed layout x cache policy x quantization x offload mode x topology x
hardware x dtype x output representation/transport
```
Use exactly these states:
| State | Meaning |
|---|---|
| `validated` | The exact row passed correctness and measured deployment tests, with reproducible artifacts. |
| `limited` | A bounded subset passed; the limitation and rejection/fallback behavior are explicit. |
| `unsupported` | A stable architectural/platform restriction is proven and fails early with an actionable error. |
| `not tested` | No adequate evidence exists. This is the default. |
Never infer compatibility from import success, server startup, another task,
another card, or another topology. Record the model/checkpoint revision,
vLLM-Omni commit, commands, request assets and hashes, seed/schedule/shape,
raw output artifacts, environment, and result for every `validated` row.
## Workflow
Follow the gates in order. Keep a dense BF16 single-device path as the
correctness oracle until every advertised fast path has parity evidence.
### Gate 0: Freeze the official contract
Pin the official implementation and checkpoint revision. Build a matrix for:
- constructor/call arguments, defaults, scheduler and sigma schedule;
- task-to-source modality/count/order rules;
- prompt, MIME, bytes, pixels, frames, duration, FPS, steps, seed, guidance,
output count, and output format limits;
- offline API and every public sync/async serving endpoint;
- checkpoint partitions and task selection.
Normalize offline, sync, and async requests through one validation contract.
Reject an invalid task/source combination before downloads, upload persistence,
temporary files, async job creation, or engine submission. Use 400 for invalid
semantics and 413 for payload limits; clean partial resources on every failure.
See [API and hardware recipes](references/api-and-recipes.md).
### Gate 1: Re-prove correctness under production shapes
1. Make loading strict: track expected, loaded, unexpected, and intentionally
ignored tensors. A fused destination is complete only after every source
shard arrives. Missing retained parameters abort startup.
2. Fix reference revision, asset hashes, seed, schedule, and shape. Compare
component outputs, intermediate denoise latents, and final artifacts.
3. Test unequal packed samples. Verify `cu_seqlens`, padding exclusion,
modality/CFG ownership, split/gather boundaries, and one output per request.
4. Test partial process groups with non-member ranks and group size smaller
than world size. Do not construct a private group per component or request.
5. Run negative API cases before positive E2E tests to prove side-effect-free
rejection and cleanup.
### Gate 2: Wire shared fast primitives
Prefer explicit, testable vLLM-Omni/vLLM operators with native fallbacks.
`torch.compile` remains a separately validated regional optimization; it is
not a substitute for correct operator selection or lifecycle-safe boundaries.
#### RMSNorm + RoPE baseline
Use the shared cross-platform operators, preserving the official Q/K norm and
RoPE order:
```python
from vllm_omni.diffusion.layers.norm import RMSNorm
from vllm_omni.diffusion.layers.rope import (
RotaryEmbedding,
apply_rope_to_qk,
)
self.norm_q = RMSNorm(head_dim, eps=eps)
self.norm_k = RMSNorm(head_dim, eps=eps)
self.rope = RotaryEmbedding(is_neox_style=False)
query = self.norm_q(query) # [B, S, Hq, D] or packed [T, Hq, D]
key = self.norm_k(key) # [B, S, Hkv, D] or packed [T, Hkv, D]
query, key = apply_rope_to_qk(
self.rope,
query,
key,
(cos, sin),
)
```
Verify `is_neox_style`, `half_head_dim`, partial rotary dimensions, cos/sin
layout, and packed row ownership against the official model. `False` means
interleaved/GPT-J style; do not assume the default is NeoX. Shared RMSNorm and
RoPE are two dispatched operators, not automatically one fused kernel.
For packed non-interleaved RoPE, use the shared fused boundary when its contract
matches the model:
```python
from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
query, key = fused_qk_norm_rope(
query, # [T, Hq, D]
key, # [T, Hkv, D]
self.norm_q.weight, # [D]
self.norm_k.weight, # [D]
rope_table, # [T, rotary_dim] = [cos | sin]
self.norm_q.variance_epsilon,
)
```
The current CUDA fast path specializes BF16 `head_dim=128`, `rotary_dim=96`;
the public function keeps an eager fallback for unsupported inputs. Verify the
packed row/frequency contract and trace the actual fast path. Do not reshape an
incompatible official RoPE layout merely to enter this kernel.
MiniMax-H3-style partial NeoX RoPE rotates a prefix and passes the suffix
through:
```python
self.q_norm = RMSNorm(head_dim, eps=eps, dtype=torch.bfloat16)
self.k_norm = RMSNorm(head_dim, eps=eps, dtype=torch.bfloat16)
self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)
def apply_partial_rope(self, x, freqs, rot_dim):
x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
cos, sin = torch.cos(freqs).to(x.dtype), torch.sin(freqs).to(x.dtype)
return torch.cat((self.rope(x_rot, cos, sin), x_pass), dim=-1)
```
The partial-RoPE snippet assumes `import torch` and lives on the owning
attention module; adapt names and validated frequency shapes rather than
copying it at module scope.
Call the modules normally so `CustomOp` selects CUDA/HIP/NPU/XPU/native
implementations. Do not call `forward_cuda()` directly. Compare the shared
path with native BF16 at operator, block, denoise-trajectory, and artifact
levels before claiming either accuracy or speed.
#### Faster BF16 attention and fused projections
Use local TP head counts and role-aware shared attention:
```python
from vllm.model_executor.layers.linear import QKVParallelLinear
from vllm_omni.diffusion.attention.layer import Attention
self.qkv_proj = QKVParallelLinear(
hidden_size=hidden_size,
head_size=head_dim,
total_num_heads=num_heads,
total_num_kv_heads=num_kv_heads,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
self.attn = Attention(
num_heads=self.qkv_proj.num_heads,
num_kv_heads=self.qkv_proj.num_kv_heads,
head_size=head_dim,
softmax_scale=head_dim**-0.5,
causal=False,
qkv_layout="BSND",
role="self",
prefix=prefix,
)
```
Prove the selected backend from runtime logs/trace and compare it with dense
SDPA. Packed metadata must express real sample boundaries; do not build masks
that the selected backend ignores. A backend that is faster for one shape or
platform is not a global default. See [performance patterns](references/performance-patterns.md)
for QKV, SwiGLU, AdaLN, sparse attention, redundancy, and fusion examples.
### Gate 3: Add online FP8 without breaking the loader
Route the active config to each component, then pass a stable runtime prefix
and config into every quantizable vLLM linear:
```python
def resolve_component_quant_config(quant_config, component):
return quant_config.resolve(component) if hasattr(quant_config, "resolve") else quant_config
transformer_quant_config = resolve_component_quant_config(
od_config.quantization_config,
"transformer",
)
self.transformer = MyDiT(
od_config,
quant_config=transformer_quant_config,
prefix="transformer",
)
self.qkv_proj = QKVParallelLinear(
hidden_size=hidden_size,
head_size=head_dim,
total_num_heads=num_heads,
total_num_kv_heads=num_kv_heads,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
```
Use `build_quant_config()` for direct construction/testing:
```python
from vllm_omni.quantization import build_quant_config
quant_config = build_quant_config({
"transformer": {"method": "fp8"},
"text_encoder": None,
"vae": None,
"default": None,
})
```
Preserve each parameter's vLLM `weight_loader`; perform checkpoint layout
conversion before invoking it. Account for every QKV/gate-up source shard.
Keep precision-sensitive norm/modulation/embedders in BF16 unless separately
proven, using the current `safe_quant_config` pattern where applicable.
Validate resident DiT FP8 first: prove which modules are quantized, HBM saved,
BF16-vs-FP8 trajectory/artifact quality, latency, and throughput. Text encoder,
VAE, other hardware, TP/cache/DLO combinations are separate rows.
Pre-quantized, pruned, rotated, or adapter-modified checkpoints are separate
model/loading lanes; compare them with the released BF16 model without calling
the result a pure quantization ablation. Read
[`../quantization/SKILL.md`](../quantization/SKILL.md) if present and prefer the
target revision's implementation/docs when sibling guidance differs.
### Gate 4: Make every component offloadable
Declare topology instead of relying on heuristic discovery:
```python
from typing import ClassVar
import torch.nn as nn
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.offloader import OffloadPlan
class MyPipeline(nn.Module, SupportsComponentDiscovery):
_dit_modules: ClassVar[list[str]] = ["transformer"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
_vae_modules: ClassVar[list[str]] = ["vae"]
_resident_modules: ClassVar[list[str]] = []
_offload_plan: ClassVar[OffloadPlan] = OffloadPlan(
on_demand_component_paths=frozenset({"text_encoder", "vae"}),
block_attrs={"transformer": ("blocks",)},
offload_submodules={"token_refiner": "blocks"},
resident_dit_paths=frozenset({"transformer"}),
encoder_block_attrs={"text_encoder": ("encoder.layers",)},
)
```
Add a test that resolves every dotted component/block path and fails if it is
missing or is not the expected module/container; discovery may otherwise warn
and skip. Validate ordinary layerwise offload and distributed layerwise
offload (DLO) independently.
Keep component lifecycle and host-weight storage separate. `OffloadPlan`
declares what can be streamed; the diffusion loader owns any `HostWeightPlan`
and hands the exact prevalidated plan to DLO. Do not make the offloader rescan
checkpoint files or duplicate loader name/shape/dtype decisions. Direct mmap is
a fail-closed optimization: current preflight requires TP1, no HSDP, no online
quantization, complete bindings, and represented transforms. Otherwise the
ordinary loader remains authoritative.
If the target revision supports Host Weight Runtime (HWR), treat it as a
separate, exact-identity startup/cache path rather than generic offload or
zero-copy execution. Qualify `preferred` population/fallback and `required`
consume-only failure independently; verify immutable manifests, final-layout
restore, lease cleanup, corruption/quarantine, and node-local sharing. Current
model contracts may be BF16 no-AllGather-only, so reject AllGather, online
quantization, HSDP, LoRA/adapted weights, and unrecognized load formats unless
the exact target revision explicitly adds them.
On small-HBM cards, record load/materialization peak, resident peak, encode,
each denoise wave, decode, transient per-rank HBM, host PSS, H2D, and collective
time. Test DLO AllGather and `--dlo-no-use-allgather` as different deployments.
Sweep the supported resident-block policies and publish the non-dominated
latency-HBM frontier; a single minimum-memory or fastest point is not the DLO
trade-off. Keep request-specific preparation replica-local. Only weight
materialization may enter the declared DLO collective group; never use WORLD
for per-request work in a DP deployment.
Direct checkpoint mmap preflight is limited to TP1 without HSDP or online
quantization; TP>1 falls back to the ordinary TP-aware loader and can still feed
DLO AllGather or no-AllGather after scoped E2E. HSDP remains incompatible with
DLO AllGather. A target revision may allow finalized per-tensor online FP8
through the ordinary loader and AllGather; keep every other online quantizer
fail-closed, and require model/card/topology E2E before promotion.
Use deploy YAML for DP under `--omni`; vLLM DP CLI flags are rejected.
### Gate 5: Add request-scoped acceleration safely
Server-wide Cache-DiT support does not imply a per-request quality API. Only
add request policy after defining calibrated quality tiers and a safe hook
lifecycle using `SupportsRequestScopedCacheDiT`, `CacheDiTRequestSpec`, and
`RequestScopedCacheDiTRuntime`. Validate alternating and concurrent quality
tiers, real cache hits, refresh/reset, success, error, disconnect, abort, and
the next uncached request. Do not batch different quality policies unless the
batch compatibility key and hook ownership make that safe.
Then test Cache-DiT against FP8, DLO, TP/SP, compile, and batching one
combination at a time. See [production feature patterns](references/feature-patterns.md)
for the lifecycle skeleton and compatibility matrix.
### Gate 6: Optimize only from profiles
Freeze one canonical workload manifest before profiling. Keep lossless runtime
A/Bs, accelerated paths with numerical/quality trade-offs, and production
topology studies in separate result lanes. Reconcile the client boundary as
queue + encode + denoise + video/audio decode + output transport/codec +
residual; do not hide a material residual in another stage.
Profile stage time, GPU kernels, CPU/GPU synchronization, allocations, H2D,
output payload/copies, and collectives before choosing work. Apply and A/B one
change at a time:
1. Remove synchronization and redundant denoise-loop work: hoist static prompt
conditioning/token refinement, RoPE tables, masks/metadata, bounded AdaLN
schedule projections, and reference scans; remove per-layer `.item()` and
unused masks/allocations.
2. Select the fastest correct BF16 attention backend for each role and shape.
3. Prefer explicit fused QKV, gate-up + `SiluAndMul`, RMSNorm/RoPE, AdaLN, and
layout/residual operators with guarded native fallbacks. Preserve every
materialized dtype/rounding boundary consumed by the unfused graph.
4. Enable sparse attention only when platform, shape, topology, realized
sparsity, quality, fallback, and end-to-end speedup are all proven.
5. Profile USP all-to-all count/bytes, packing, contiguous copies, overlap, and
scaling efficiency. Audit replicated cross-attention before using
`skip_sequence_parallel`. Qualify regular and accelerated Ulysses transport
independently; activation logs, JIT/readiness time, workspace growth,
stream ownership, maximum-shape warmup, and numerical drift are gates.
Preserve the checkpoint's Q:KV head ratio when padding GQA for Ulysses, and
validate strided QKV staging rather than assuming packed contiguity.
6. Evaluate text-encoder or VAE disaggregation only if stage share, transfer
volume, reuse/concurrency, and independent scaling justify it. Tiling,
offload, or patch parallelism is not disaggregation.
Do not combine independent optimizations into one performance claim. Preserve
raw before/after runs with fixed workload and warmup policy.
### Gate 7: Make serving interruptible and load-stable
Treat request-mode batching and step-wise execution as separate capabilities:
- request batching requires `supports_request_batch = True` and
`DiffusionRequestBatch -> list[DiffusionOutput]`, exactly one output per
request;
- step execution requires the complete `SupportsStepExecution` contract:
`prepare_encode`, `denoise_step`, `step_scheduler`, and `post_decode`.
Keep all mutable scheduler, RNG, latent, mask, cache, and step state request
scoped. First validate `--step-execution --max-num-seqs 1` and step-boundary
abort; only then test heterogeneous multi-request waves. Step continuous
batching is experimental and is not automatically compatible with Cache-DiT.
Before recommending it, predeclare a useful case and success threshold, then
A/B it against request mode under the same arrival process. Structural support
or a generic bridge is not latency, throughput, or cancellation evidence for a
model.
Test queued and in-flight abort, client disconnect, timeout, OOM, worker error,
and restart. Require one terminal result, idempotent cleanup, no post-decode
after abort, no cache/temp/VRAM/request-ID leak, and a successful next request.
Inject late worker results/exceptions after cancellation and prove the result
pump stays alive instead of raising `InvalidStateError`. For large batched or
long-window video outputs, freeze dtype, range, layout, contiguity, ownership,
payload size, serializer/codec limits, and offline-versus-HTTP semantics.
Validate device-side output preparation, D2H/IPC or shared handles, remote
encoding, event-loop responsiveness, and client materialization as separate
stages, including abort and consumer failure cleanup.
A model/VAE callback that publishes ordered chunks is only a producer contract;
it is not transport, backpressure, cancellation, public streaming, or
time-to-first-frame evidence. Distinguish complete-response faster-than-playback
(`client E2E / output duration <= 1`) from streaming and report first-chunk,
steady cadence, finalization, and complete-artifact latency separately.
Run below-, near-, and above-saturation mixed-RPS soaks and report success/error
rate, throughput, queue time, memory slope, temp growth, and worker health.
Report p50/p95/p99 only from a declared arrival model with enough measured
requests for the claimed percentile. See
[serving and CI validation](references/production-validation.md).
### Gate 8: Publish recipes and four independent CI tracks
For every target vendor/card/topology, publish exact environment, serve/deploy
command, one complete JSON or multipart curl per supported task, expected
output checks, benchmark command, warmup/repeat/concurrency policy, raw metrics,
per-rank HBM/host PSS, accuracy artifact, and limitations. Mark each
CUDA/ROCm/NPU/XPU row `validated`, `unsupported`, or `not tested`; never inherit
support from the platform abstraction or another card. Re-run DLO on each
small-HBM recipe.
Maintain separate CI gates:
1. Function: strict load, API/offline positive and negative contracts.
2. Accuracy: fixed reference revision/assets/seed/schedule and scoped tolerance.
3. Performance: fixed best deployment plus memory-constrained/DLO row, raw JSON,
explicit warmup, enough samples for each claimed statistic, and owned
regression thresholds.
4. Reliability: long mixed-RPS soak, abort/disconnect/fault cleanup, memory and
temporary-resource trend.
## Production Definition of Done
A model is production-ready only when:
- every official advertised task passes API parity, offline/online parity,
strict loading, fixed-reference quality, and negative validation;
- the recommended deployment row passes correctness, isolation, abort/error
cleanup, benchmark, and soak gates on named hardware;
- online FP8 and DLO are each implemented and either validated in scoped rows
or explicitly limited/unsupported with fail-fast behavior;
- per-request cache quality tiers have measured speed/quality frontiers and
request isolation, if advertised;
- CUDA/ROCm/NPU/XPU recipes describe evidence rather than inferred support;
- Function, Accuracy, Performance, and Reliability CI have owners, artifacts,
and actionable failure output;
- accelerated paths identify every precision/quality trade-off, and the raw
offline plus online encoded output contracts are versioned and tested;
- all unverified task/backend/cache/quant/offload/topology/hardware combinations
remain `not tested` and are not presented as supported.
Day-0 support is still a valid milestone, but title, docs, and matrix must say
which representative task passed and which production gates remain.
@@ -1,4 +0,0 @@
interface:
display_name: "Production Add Diffusion Model"
short_description: "Productionize vLLM-Omni diffusion models"
default_prompt: "Use $production-add-diffusion-model to take this integrated diffusion model from Day-0 support to evidence-backed production readiness."
@@ -1,315 +0,0 @@
# API Contract and Hardware Recipes
## Contents
1. [Freeze the upstream contract](#freeze-the-upstream-contract)
2. [Validate before side effects](#validate-before-side-effects)
3. [Task request templates](#task-request-templates)
4. [Serving benchmark template](#serving-benchmark-template)
5. [Hardware recipe contract](#hardware-recipe-contract)
6. [Best deployment selection](#best-deployment-selection)
## Freeze the upstream contract
Pin three revisions independently:
- official source implementation commit or package version;
- checkpoint/model-card revision and partition;
- vLLM-Omni target commit.
Build a table before editing code. One row is one official task/partition.
| Field | Required content |
|---|---|
| Task and partition | Official task name, checkpoint subdirectory, task selection rules |
| Call contract | Constructor and call arguments, types, required/optional, official defaults |
| Prompt/input | Prompt count/length, modality, MIME, order, count, decoded shape |
| Size/time | Width/height alignment, pixels, frames, duration, FPS, start offset |
| Denoise | Scheduler, sigma schedule, steps, CFG/guidance, seed behavior |
| Output | Count, format, codec/sample rate, return type |
| API | Offline mapping, sync endpoint, async endpoint/job lifecycle |
| Rejection | Status code and stable error class/message for every invalid boundary |
Official API parity includes limits and defaults, not just names. Do not add a
generic default that changes a distilled checkpoint's fixed schedule, or accept
an input combination the official task rejects.
For every accepted boundary, add the neighboring rejected values. Examples:
`max_images` and `max_images + 1`, smallest/largest legal duration, aligned and
misaligned dimensions, omitted and illegal steps, one and too many outputs.
## Validate before side effects
Use one normalization/validation path for offline, sync, and async requests.
The ordering must be:
```text
parse -> normalize aliases -> validate task/source matrix -> validate declared
limits -> stream with byte caps -> inspect/decode with frame/pixel/duration caps
-> persist only if needed -> submit to engine
```
Fail before download, persistence, job creation, or engine submit whenever the
available metadata is sufficient. For streamed uploads, enforce:
- total reference count and per-modality count before reading bodies;
- per-file bytes while streaming;
- aggregate bytes across all files;
- decoded frames, duration, pixels, and media type before engine submission.
On rejection, close streams and remove every partial temporary resource. Use
HTTP 400 for invalid task/field combinations and 413 for payload limits. Test
both sync and async paths; an invalid async request must not create a job ID.
Do not trust client MIME alone. Sniff/validate content and keep ordered repeated
multipart fields ordered. Avoid logging prompts, raw references, signed URLs,
or base64 payloads.
## Task request templates
These are endpoint-shape templates, not model recipes. Replace placeholders
with the exact official defaults and validated assets for the target model.
Every production recipe must include all advertised tasks and an expected
output check.
### Text to image: JSON
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/images/generations \
-H 'Content-Type: application/json' \
-d '{
"model": "<model>",
"prompt": "<validated prompt>",
"size": "1024x1024",
"num_inference_steps": 30,
"seed": 42,
"response_format": "b64_json"
}' | jq -er '.data[0].b64_json' | base64 -d > output.png
file output.png
```
### Image edit: multipart
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/images/edits \
-F 'model=<model>' \
-F 'prompt=<validated edit prompt>' \
-F 'image=@/path/to/reference.png;type=image/png' \
-F 'size=1024x1024' \
-F 'num_inference_steps=30' \
-F 'seed=42' \
-F 'output_format=png' \
| jq -er '.data[0].b64_json' | base64 -d > edited.png
file edited.png
```
Use repeated `image` fields only if the model's official input matrix allows
multiple ordered sources. Never describe a file upload as JSON-only.
### Public request fields versus internal sampling fields
Use the field exposed by the target revision's endpoint schema; do not rename
it from the pipeline's internal attribute. For example, the current video JSON
and multipart contracts expose `extra_params`, while serving translates that
object into `DiffusionSamplingParams.extra_args`. Other request paths may
expose canonical `extra_args` or nest it under `extra_body`. Inspect the exact
Pydantic model/FastAPI form dependency and exercise the generated curl before
publishing a recipe. A deprecation warning in a different endpoint is not
evidence that an unrecognized multipart form field works here.
Current endpoint map (re-check it at the target commit):
| Endpoint | Public model-specific field |
|---|---|
| `/v1/videos`, `/v1/videos/sync` | Multipart `extra_params` JSON string |
| `/v1/images/generations` | JSON `extra_params` object |
| `/v1/chat/completions` diffusion | `extra_body.extra_args` |
| `/v1/images/edits` | Declared typed form fields only |
| `/v1/audio/generate` | Declared typed JSON fields only |
Prefer declared top-level fields. Use a generic model-specific envelope only
when the endpoint actually exposes it.
### Text to video: multipart
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
-H 'Accept: video/mp4' \
-F 'model=<model>' \
-F 'prompt=<validated prompt>' \
-F 'size=1280x720' \
-F 'num_frames=81' \
-F 'fps=16' \
-F 'num_inference_steps=30' \
-F 'seed=42' \
-F 'extra_params={"task":"<official-task>"}' \
-o output.mp4
ffprobe -v error -show_entries stream=codec_name,width,height,r_frame_rate \
-of json output.mp4
```
### Image/video to video: multipart
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
-H 'Accept: video/mp4' \
-F 'model=<model>' \
-F 'prompt=<validated prompt>' \
-F 'num_frames=81' \
-F 'fps=16' \
-F 'num_inference_steps=30' \
-F 'seed=42' \
-F 'extra_params={"task":"<official-task>","start_time_seconds":0}' \
-F 'input_reference=@/path/to/reference.png;type=image/png' \
-o i2v.mp4
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
-H 'Accept: video/mp4' \
-F 'model=<model>' \
-F 'prompt=<validated prompt>' \
-F 'num_frames=81' \
-F 'num_inference_steps=30' \
-F 'seed=43' \
-F 'extra_params={"task":"<official-task>","start_time_seconds":0}' \
-F 'input_reference=@/path/to/reference.mp4;type=video/mp4' \
-o v2v.mp4
```
For multiple references, repeat the model's accepted multipart field in
official order. Include complete downloadable or repository fixtures plus
SHA256 hashes. A placeholder-only recipe is not runnable evidence.
### Text to audio: JSON
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/audio/generate \
-H 'Content-Type: application/json' \
-d '{
"model": "<model>",
"input": "<validated audio prompt>",
"audio_length": 10.0,
"num_inference_steps": 100,
"seed": 42,
"response_format": "wav"
}' -o output.wav
ffprobe -v error -show_entries stream=codec_name,sample_rate,channels \
-of json output.wav
```
Use `/v1/audio/generate` for general diffusion audio. Do not substitute the
speech API unless the model is actually served through the speech contract.
### Async video lifecycle
Alongside `/v1/videos/sync`, include `POST /v1/videos`, status polling,
content retrieval, cancellation/abort if public, and cleanup assertions. The
same request must normalize to the same engine sampling parameters in both
modes.
## Serving benchmark template
Run a smoke first, then fixed-concurrency and arrival-rate tests. Video task
names for the harness are `t2v`, `i2v`, `ti2v`, or `v2v`; model-specific task
selection belongs in `--extra-body` when required.
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://127.0.0.1:8091 \
--endpoint /v1/videos \
--model '<model>' \
--dataset random \
--task t2v \
--num-prompts 100 \
--width 1280 --height 720 \
--num-frames 81 --fps 16 \
--num-inference-steps 30 \
--request-rate 0.2 \
--max-concurrency 4 \
--warmup-requests 4 \
--warmup-concurrency 4 \
--extra-body '{"extra_params":"{\"task\":\"<official-task>\"}"}' \
--output-file result.json
```
The video backend currently stringifies each multipart value. Pre-serialize
the nested `extra_params` object as above; passing it as a nested object would
produce Python dict syntax rather than valid JSON. If the target revision
changes the backend to JSON-encode dict/list values, update and smoke-test this
command against that revision.
For i2v/v2v, use the harness's validated dataset/trace path or provide the
documented media source; do not silently benchmark t2v and label it i2v. Keep
prompt, assets, seed, dimensions, frames, steps, guidance, warmup, concurrency,
and arrival process identical across A/B runs.
Keep three benchmark lanes distinct:
1. lossless runtime A/B with the released schedule/precision/attention policy;
2. accelerated-path A/B that changes one adapter, quantization, cache, or
sparsity policy and adds same-seed quality evidence;
3. production-topology studies that vary placement or parallelism and report a
latency/throughput/memory frontier rather than a kernel speedup.
For a fixed single-request A/B, retain every measured run and report median or
mean with range after an explicit warmup; do not infer p95/p99 from a few
repetitions. Tail latency requires a declared arrival process, request count,
concurrency, and enough successful samples. Reconcile client E2E with queue,
encoder, denoise, video/audio decode, D2H/IPC, codec, and residual timings.
For generated media, define `real-time` rather than relying on the label. A
complete-response claim uses `client E2E / validated output duration <= 1` and
does not imply streaming or low time-to-first-frame. Report first chunk/fragment
and cadence separately when a public streaming route exists. Never add gains
across benchmark lanes with different schedules, prompts, seeds, artifacts, or
topologies.
The diffusion serving harness currently covers image/video task families, not
`/v1/audio/generate`. For text-to-audio, use the target model's task-specific
audio benchmark if one exists; otherwise add a serving harness that records
request rate, concurrency, latency, generated-audio seconds per second, error
rate, and output validation. Do not relabel the TTS harness or a video workload
as diffusion-audio evidence.
## Hardware recipe contract
Create a separate row and runnable section for each vendor/card/topology:
| Evidence | Required detail |
|---|---|
| Identity | Vendor, exact SKU, device count, HBM, interconnect |
| Software | Driver, CUDA/ROCm/CANN/oneAPI, framework/container, dependency versions |
| Source | vLLM-Omni commit, official/checkpoint revision and hashes |
| Deployment | dtype, attention backend per role, TP/SP/CFG/HSDP/DP, cache, quant, offload |
| Commands | Complete deploy YAML/serve command, every task curl, benchmark command |
| Method | Warmup exclusion, repetitions or request count, concurrency, arrival distribution, duration, statistic definitions |
| Results | Success rate, p50/p95/p99, RPS/throughput per device, stage latency |
| Memory | Load/materialize, resident, encode, denoise, decode, transient per-rank HBM and host PSS |
| Quality | Metric/tolerance, artifact links/hashes, temporal/audio checks |
| Output | Raw/offline dtype, range, layout and ownership; online codec, payload bytes, transport route and client boundary |
| Limits | `validated`/`limited`/`unsupported`/`not tested`, fallback/rejection behavior |
Use current paths such as `recipes/<vendor>/`, `recipes/README.md`, supported
models, and feature compatibility tables. Top-level documentation must link to
the scoped matrix; it must not use an unqualified checkmark that hides a ROCm,
NPU, XPU, task, topology, or dtype limitation.
DLO evidence is per card, especially on small-HBM devices. A DLO result on one
CUDA card does not establish another CUDA SKU, ROCm, NPU, or XPU support.
## Best deployment selection
Benchmark a correctness-approved candidate matrix, not every theoretical
combination. Choose the recommended row by the production objective—latency,
throughput, cost, or small-HBM fit—and document that objective.
Do not recommend a row until it passes:
- all advertised task/API contracts;
- fixed-reference quality;
- request isolation and abort/error cleanup;
- sustained arrival-rate tests;
- per-rank HBM and host PSS limits;
- hardware-specific accuracy and performance CI.
Keep the most memory-constrained DLO row even when it is not the fastest; it is
a separate production recipe, not a failed latency candidate.
@@ -1,462 +0,0 @@
# Production Feature Patterns
## Contents
1. [Compatibility matrix](#compatibility-matrix)
2. [Online FP8](#online-fp8)
3. [Strict fused-weight loading](#strict-fused-weight-loading)
4. [Distributed layerwise offload](#distributed-layerwise-offload)
5. [Per-request Cache-DiT](#per-request-cache-dit)
6. [Combination gates](#combination-gates)
## Compatibility matrix
Create rows at the granularity below. Add evidence links and a reason for every
non-validated state.
| Task/shape | Execution mode/capacity | Backend/layout | Quant | Cache policy | Offload | Topology | Hardware/dtype | Output contract/transport | State | Evidence |
|---|---|---|---|---|---|---|---|---|---|---|
| `<task>/<shape>/<schedule>` | serial request / 1 | SDPA, padded | BF16 | none | resident | 1 device | `<card>` BF16 | raw float / local | `not tested` | — |
| `<task>/<shape>/<schedule>` | step / N | `<fast>`, packed | BF16 | none | resident | `<TP/SP>` | `<card>` BF16 | encoded / subprocess | `not tested` | — |
| `<task>/<shape>/<schedule>` | step / N | `<fast>`, packed | FP8 | high | DLO no-AG | `<TP/SP>` | `<card>` | device uint8 / SHM | `not tested` | — |
Initialize every row as `not tested`; the table is a work queue, not evidence.
Promote the dense BF16 row to the oracle only after its scoped parity passes.
Add one axis at a time. If the runtime cannot
enforce a known-incompatible combination, add construction/admission
validation before writing a recipe.
## Online FP8
### Component routing
The public builder accepts one method or a longest-prefix component map:
```python
from vllm_omni.quantization import build_quant_config
config = build_quant_config({
"transformer": {
"method": "fp8",
"ignored_layers": ["transformer.final_layer"],
},
"text_encoder": None,
"vae": None,
"default": None,
})
assert config.resolve("transformer.blocks.0.attn.qkv_proj") is not None
assert config.resolve("text_encoder.layers.0.mlp") is None
assert config.resolve("vae.decoder.conv_out") is None
```
Use actual runtime prefixes. Multi-DiT or nested pipelines may not use the
generic names `transformer` and `vae`. A prefix that falls through to a default
is not component validation. Add tests for every included and ignored prefix.
For CLI serving, first validate resident DiT-only FP8:
```bash
vllm serve '<model>' --omni --quantization fp8 --max-num-seqs 1
```
For scoped routing:
```bash
vllm serve '<model>' --omni \
--diffusion-quantization-config \
'{"<dit-runtime-prefix>":{"method":"fp8","ignored_layers":["<exact-linear-prefix>"]},"<vae-runtime-prefix>":null,"default":null}' \
--max-num-seqs 1
```
Do not publish these placeholders. A recipe must contain the target model's
verified prefixes.
### Quantizable module wiring
Pass `quant_config` and a stable prefix through every vLLM linear. Keep
precision-sensitive modulation/norm layers safe unless proven otherwise.
```python
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear,
)
self.qkv_proj = QKVParallelLinear(
hidden_size=hidden_size,
head_size=head_dim,
total_num_heads=num_heads,
total_num_kv_heads=num_kv_heads,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.attn.qkv_proj",
)
self.gate_up_proj = MergedColumnParallelLinear(
hidden_size,
[intermediate_size, intermediate_size],
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.mlp.gate_up_proj",
)
self.act_fn = SiluAndMul()
self.down_proj = RowParallelLinear(
intermediate_size,
hidden_size,
bias=False,
input_is_parallel=True,
quant_config=quant_config,
prefix=f"{prefix}.mlp.down_proj",
)
```
Set required checkpoint parameters to fail rather than silently initialize if
the current loader supports `missing_param_init = "error"`.
### FP8 evidence
For each task/hardware row, capture:
- startup proves the FP8 quant method is attached to intended linear modules;
- no unexpected BF16 fallback and no meta/uninitialized parameters;
- same-seed component, trajectory, and final-artifact comparison with BF16;
- HBM at load, resident, encode, denoise, decode, and peak;
- cold load time; fixed-work warm-latency raw runs; and serving p50/p95/p99 plus
throughput only from a declared arrival load with enough samples;
- named ignored layers and why their BF16 retention is required.
DiT FP8 evidence does not cover the text encoder or VAE. CUDA evidence does not
cover ROCm/NPU/XPU. A pre-quantized checkpoint is a different loading path from
online FP8 and needs its own state. If it is also pruned, rotated, distilled, or
adapter-modified, report that model delta explicitly; BF16-versus-artifact
quality and speed are deployment comparisons, not pure quantization ablations.
## Strict fused-weight loading
Never use `param.data.copy_()` for quantizable vLLM parameters. Transform the
checkpoint layout before invoking the parameter loader so the online quant
loader remains outermost:
```python
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
param = params[target_name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
if target_name.endswith(".qkv_proj.weight"):
packed_qkv = reorder_checkpoint_qkv(checkpoint_weight)
weight_loader(param, packed_qkv)
elif target_name.endswith(".gate_up_proj.weight"):
gate, up = checkpoint_weight.chunk(2, dim=0)
weight_loader(param, gate, 0)
weight_loader(param, up, 1)
else:
weight_loader(param, checkpoint_weight)
```
Track source shards per fused destination:
```python
expected = {
"blocks.0.attn.qkv_proj.weight": {"q", "k", "v"},
"blocks.0.mlp.gate_up_proj.weight": {0, 1},
}
seen = {name: set() for name in expected}
# After mapping/loading each source:
seen[target_name].add(shard_id)
incomplete = {
name: sorted(shards - seen[name], key=str)
for name, shards in expected.items()
if shards - seen[name]
}
if incomplete:
raise RuntimeError(f"incomplete fused checkpoint parameters: {incomplete}")
```
Also compare retained model parameters with the returned loaded set. Explicit
checkpoint ignores need a reason and test; an unknown or missing weight is not
an ignore policy.
Load-time adapter fusion creates a dedicated model, not a request-switchable
LoRA. Pin the base and adapter revisions, validate every low-rank and full-rank
delta target, freeze the adapter-declared schedule/task, and fuse before
sharding only when the loader contract proves identical rank-local weights.
Reject offload or host-weight paths that bypass that fusion rather than serving
the unfused base model silently.
## Distributed layerwise offload
### Declare component topology
Use `SupportsComponentDiscovery` and `OffloadPlan`. Validate all dotted paths
at construction/test time because discovery may otherwise warn and skip.
Treat a future/requested `--layerwise-offload-components` selector as a public
API only if it exists in the target revision. Today, express component topology
with `OffloadPlan` (`on_demand_component_paths`, `block_attrs`,
`encoder_block_attrs`, `resident_dit_paths`, and `offload_submodules`). Do not
invent a CLI flag or claim component-selective DLO from discovery declarations
alone; validate the selector separately when it lands.
```python
from operator import attrgetter
import torch
def assert_offload_path(root, path):
value = attrgetter(path)(root)
if not isinstance(value, torch.nn.Module):
raise TypeError(f"offload path {path!r} is not an nn.Module")
return value
for path in (
*pipeline._dit_modules,
*pipeline._encoder_modules,
*pipeline._vae_modules,
*pipeline._offload_plan.on_demand_component_paths,
):
assert_offload_path(pipeline, path)
```
For `block_attrs`, resolve the DiT first and then each declared block attribute.
Require an indexable block container with at least one `nn.Module`. Validate
`offload_submodules`, `resident_dit_paths`, and `encoder_block_attrs` similarly.
### Keep host storage loader-owned
Current DLO accepts a loader-produced `HostWeightPlan`. Pipeline authors should
not instantiate this plan or add ad hoc mmap capability flags. The loader must
preflight complete runtime/checkpoint bindings, names, shapes, dtypes, persistent
buffers, and represented transforms before it skips ordinary materialization.
The exact accepted plan is consumed once by DLO; a rejected plan records an
observable fallback reason and continues through the ordinary loader.
Direct checkpoint mmap currently requires TP1 without HSDP or online
quantization. TP>1 falls back to the ordinary TP-aware loader and may feed DLO
AllGather or no-AllGather; the current design records a bounded DP2xTP2
AllGather smoke, but that path does not receive direct-mmap host-page sharing.
Per-tensor online FP8 must use the ordinary loader. Current revisions can admit
its finalized weight/scale layout to DLO AllGather, including transposed runtime
weights, while unsupported online quantizers still fail closed. Direct mmap
remains unavailable, so capture the temporary full-model host materialization
before rank sharding.
Host Weight Runtime is a distinct opt-in path. Current final-layout consumers
use no-AllGather DLO and exact model-declared BF16 representations. In
`preferred` mode, consume an exact local hit or canonically load and publish for
a future startup; in `required` mode, consume only and fail on a miss or unusable
artifact. HWR must not interact with DLO AllGather, online quantization, HSDP,
LoRA/adapted weights, or an unrecognized load format unless the target revision
adds an exact producer/restore contract. It is immutable node-local CPU backing,
not zero-copy GPU execution; the DLO transport still owns staging or direct H2D.
Qualify identity inputs, cold population, warm reuse, atomic concurrent
publication, source digests, corruption/quarantine, lease cleanup, capacity,
node-local PSS, H2D, and E2E latency. A host-memory saving that materially
regresses transfer or request latency is a separate memory-first deployment.
### DLO deployment modes
Treat these paths separately:
| Path | Weight/runtime behavior | Required boundaries |
|---|---|---|
| DLO AllGather | Rank-sharded pinned host tensors; reconstruct layer through collective | TP>1 uses ordinary TP-aware loader output; HSDP and unsupported online quantizers are rejected; finalized per-tensor online FP8 is eligible only when recognized by the target revision; concurrent ranks must participate consistently |
| DLO no-AllGather | Loader-approved checkpoint mmap or ordinary runtime tensors; each rank streams a complete block | Direct mmap is TP1/non-HSDP/non-online-quant only; TP/HSDP/online quant use ordinary runtime tensors and require scoped E2E |
| SP + DLO | SP group can supply the DLO collective group | Packed boundaries and collective order require E2E |
Single-stage examples:
```bash
# SP + DLO AllGather candidate; validate exact model/card before publishing.
vllm serve '<model>' --omni \
--enable-distributed-layerwise-offload \
--usp 4 \
--max-num-seqs 1
# TP + DLO rank-local candidate; not a support statement.
vllm serve '<model>' --omni \
--tensor-parallel-size 2 \
--enable-distributed-layerwise-offload \
--dlo-no-use-allgather \
--max-num-seqs 1
```
Under `--omni`, configure DP in a deploy YAML rather than passing vLLM DP CLI
flags:
```yaml
pipeline: <registered-pipeline>
async_chunk: false
data_parallel_size: 4
stages:
- stage_id: 0
devices: "0,1,2,3"
max_num_seqs: 1
enable_distributed_layerwise_offload: true
dlo_use_allgather: true
parallel_config:
tensor_parallel_size: 1
sequence_parallel_size: 1
```
```bash
vllm serve '<model>' --omni --deploy-config /path/to/dlo_dp4.yaml
```
Keep the YAML next to the validated recipe. Match its schema to the target
revision; do not assume a model accepts another pipeline's deploy config.
### DLO validation sequence
1. Resident BF16 reference.
2. Ordinary layerwise offload parity.
3. DLO enable, warmup, same-seed generation, disable/restart.
4. AllGather and no-AllGather separately.
5. Concurrent requests with different prompts and identical collective-affecting
parameters; then reject or schedule incompatible steps/shapes safely.
6. Queued/in-flight abort, error, idle rank, worker exit, and next request.
7. TP, SP, cache, compile, online FP8 one axis at a time.
8. Loader-plan selection/fallback reason, cold/warm storage lifecycle, and
corruption/stale-entry recovery when mmap/cache backing is advertised.
9. Replica-local request preparation plus DLO-group-only weight collectives in
DP; exercise concurrent replicas to catch accidental WORLD rendezvous.
10. Sweep resident block counts; publish raw latency/HBM pairs, dominated
points, and the non-dominated Pareto frontier.
11. Per-rank HBM, host PSS for the full process tree, H2D and collective trace.
On a small-HBM card, measure transient encode/VAE peaks as well as resident DiT
weights. If a non-DiT component still exceeds HBM, combine only independently
validated on-demand component staging, VAE tiling/patching, or another topology;
do not hide the remaining peak behind the phrase “DLO enabled.”
## Per-request Cache-DiT
### Policy and lifecycle skeleton
First verify that the target revision contains the shared protocol/runtime.
If not, follow the nearest current model's lifecycle without inventing a new
public API.
The following is a **serial/exclusive request-mode sketch only**. Pipeline-wide
hooks are mutable; concurrent requests need an admission/batching key and
runner-owned transition serialization before this shape is safe.
```python
import torch.nn as nn
from vllm_omni.diffusion.cache.cachedit import (
CacheDiTBackend,
CacheDiTRequestSpec,
RequestScopedCacheDiTRuntime,
)
class MyPipeline(nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
# ... components ...
self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
self._cache_dit_runtime.adopt(
backend,
installation_key="my_model.generic",
)
def is_cache_dit_enabled(self) -> bool:
return self._cache_dit_runtime.is_enabled
def _cache_spec(self, quality, steps):
if quality == "lossless":
return None
if quality == "high":
return CacheDiTRequestSpec(
installation_key="my_model.high.v1",
cache_config=self._validated_high_config,
num_inference_steps=steps,
)
raise ValueError(f"unsupported quality tier: {quality!r}")
def forward(self, req):
quality, steps = validate_and_resolve_request(req)
self._cache_dit_runtime.prepare(self._cache_spec(quality, steps))
return self._run_validated_request(req)
```
The snippet shows the model-owned transition for exclusive execution. Put
disable/restore handling in the runtime/runner lifecycle that owns success,
exception, disconnect, and asynchronous abort; a `try/finally` inside
`forward()` alone cannot observe every server cancellation boundary.
The `installation_key` must include every policy dimension that changes hook
installation; changes in step count can use runtime refresh. `None` disables
installed hooks. Do not mutate the global startup cache config for a request.
Centralize cleanup in the owning request/runner boundary. An exception, abort,
disconnect, or admission failure must not leave a changed hook policy visible
to the next request. The exact owner depends on the target revision; add a
failure-injection test rather than relying on `forward()` alone to observe
asynchronous abort.
Example model-specific request form after calibration:
```bash
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
-F 'model=<model>' \
-F 'prompt=<prompt>' \
-F 'quality=lossless' \
-F 'num_inference_steps=<validated-steps>' \
-o lossless.mp4
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
-F 'model=<model>' \
-F 'prompt=<prompt>' \
-F 'quality=high' \
-F 'num_inference_steps=<validated-steps>' \
-o high.mp4
```
Do not expose these tier names merely because another model uses them.
### Cache-DiT evidence
For every tier/task/shape/schedule:
- show real cache hits/skipped blocks, not only “backend enabled” logs;
- measure fixed-work speed and quality against lossless/native;
- publish a speed/quality frontier and chosen tier mapping;
- alternate lossless/high/lossless and vary steps to prove restore/refresh;
- interleave tasks and shapes; run concurrent requests only when scheduling
guarantees exclusive compatible hook state;
- inject validation error, denoise error, disconnect, and abort, then prove the
next request starts with its requested policy;
- bound any cached tensors and clear request-local state.
## Combination gates
Start every combination as `not tested`.
| Combination | Default production posture |
|---|---|
| Direct checkpoint mmap + TP>1/HSDP/online quant | Preflight falls back to the ordinary loader; do not claim direct-mmap savings |
| Per-tensor online FP8 + DLO AllGather | Generic compatibility exists for finalized weights/scales; each model/card/topology remains a candidate until E2E, and other online quantizers fail closed |
| Online FP8 + DLO no-AllGather | Candidate only after model/card/topology E2E |
| TP>1 + DLO AllGather | Ordinary TP-aware loader only; bounded DP2xTP2 smoke exists, but each new model/card/topology still needs E2E |
| TP>1 + DLO no-AllGather | Ordinary TP-aware loader only; candidate with limited generic coverage |
| HSDP + DLO AllGather | Reject; it would double-shard HSDP-managed parameters |
| HSDP + DLO no-AllGather | Configuration-compatible candidate with limited E2E coverage |
| Cache-DiT + online FP8 | Separate trajectory, hit-rate, HBM, latency test |
| Cache-DiT + DLO | Verify hooks, streamed blocks, memory peaks, abort cleanup |
| Cache-DiT + step execution | Do not advertise until lifecycle/batching support is explicit |
| Sparse attention + Ring | Verify backend execution; reject silently ignored sparse settings |
| Load-time fused adapter + offload/HWR | Reject unless the selected loader path applies and validates the same fusion before sharding/storage |
| Packed attention + any topology | Unequal-sample boundary parity is mandatory |
| Device-side uint8 preparation + offline caller | Version the dtype/range/layout change; HTTP MP4 parity does not preserve the raw tensor contract |
| Device-side preparation + remote codec | Validate route selection, strided/contiguous layouts, payload bytes, ordering, fallback, and encoded-media parity |
A startup-only test never upgrades a combination. Require request output,
parity/quality, actual fast-path evidence, resource cleanup, and a benchmark on
the named hardware.
@@ -1,481 +0,0 @@
# Performance Patterns and Code Examples
## Contents
1. [Measurement loop](#measurement-loop)
2. [Faster BF16 attention](#faster-bf16-attention)
3. [Sparse attention](#sparse-attention)
4. [Fusion-first operator patterns](#fusion-first-operator-patterns)
5. [Remove redundant computation](#remove-redundant-computation)
6. [USP communication optimization](#usp-communication-optimization)
7. [Text encoder and VAE disaggregation](#text-encoder-and-vae-disaggregation)
8. [Performance acceptance](#performance-acceptance)
## Measurement loop
Do not start with a list of optimizations. Start with a fixed correctness and
performance workload:
```text
model/checkpoint revision + prompt/media hashes + seed + dimensions/frames +
scheduler/sigma/steps/guidance + task + dtype + backend + topology + hardware
```
Capture:
- stage time: load, encode, conditioning, denoise, VAE decode/postprocess;
- GPU trace: kernel duration/count, launch gaps, allocations, H2D, collectives;
- CPU trace: preprocessing, per-layer sync, Python dispatch, temporary files;
- memory: per-rank HBM peaks and process-tree host PSS;
- output: device preparation, D2H/IPC, payload bytes, codec wall/process CPU,
event-loop wait, client materialization, and signed residual;
- serving: queue time, throughput, errors, request mix, and distribution samples.
For a fixed-work single-request A/B, run one explicit warmup policy and at
least three measured repetitions, keep every raw run, and report median or mean
with range. Do not derive p95/p99 from three samples. For tail latency, declare
the arrival model and collect enough successful requests for the percentile and
confidence claimed. Implement one optimization at a time, re-run parity, and
report both local kernel and end-to-end deltas. A kernel win that moves no E2E
metric is not a production speed claim.
Use the pipeline profiler already exposed by the target revision when
available, and use PyTorch/CUDA/ROCm/NPU/XPU profiling tools appropriate to the
platform. Do not force every platform through a CUDA-only workflow.
## Faster BF16 attention
### Shared attention wiring
Use the shared role-aware attention layer and give it the true local head
counts and QKV layout:
```python
import torch
from vllm.model_executor.layers.linear import (
QKVParallelLinear,
RowParallelLinear,
)
from vllm_omni.diffusion.attention.layer import Attention
class SelfAttention(torch.nn.Module):
def __init__(
self,
hidden_size,
head_dim,
num_heads,
num_kv_heads,
quant_config=None,
prefix="",
):
super().__init__()
self.head_dim = head_dim
self.qkv_proj = QKVParallelLinear(
hidden_size=hidden_size,
head_size=head_dim,
total_num_heads=num_heads,
total_num_kv_heads=num_kv_heads,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
self.attention = Attention(
num_heads=self.qkv_proj.num_heads,
num_kv_heads=self.qkv_proj.num_kv_heads,
head_size=head_dim,
softmax_scale=head_dim**-0.5,
causal=False,
qkv_layout="BSND",
role="self",
prefix=prefix,
)
self.out_proj = RowParallelLinear(
num_heads * head_dim,
hidden_size,
bias=False,
input_is_parallel=True,
quant_config=quant_config,
prefix=f"{prefix}.out_proj",
)
def forward(self, x, attn_metadata=None):
qkv, _ = self.qkv_proj(x)
q_size = self.qkv_proj.num_heads * self.head_dim
kv_size = self.qkv_proj.num_kv_heads * self.head_dim
q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
q = q.unflatten(-1, (self.qkv_proj.num_heads, self.head_dim))
k = k.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim))
v = v.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim))
out = self.attention(q, k, v, attn_metadata)
out = out.flatten(-2)
out, _ = self.out_proj(out)
return out
```
Add Q/K norm and RoPE in the official order before the `Attention` call. Keep
checkpoint mapping for the fused QKV and output projection complete.
### Backend qualification
For each role (`self`, cross-attention category, token refiner, etc.), record:
- actual selected backend from logs/trace;
- supported dtype, head size, QKV layout, mask/metadata, packed varlen behavior;
- platform and compute capability;
- TP/SP/Ring/AllGather-KV restrictions;
- fallback behavior and whether fallback changes the claimed performance.
Use unequal packed samples and compare with dense SDPA. Verify `cu_seqlens`,
max sequence length, padding exclusion, modality boundaries, CFG ownership, and
output slicing. Do not pass an `attn_mask` and assume it is consumed. Inspect
the backend contract and trace.
TRTLLM, Flash Attention, cuDNN, NPU varlen, and other backends are independent
qualification rows. The shared `Attention` class does not make every backend
packed-safe or faster automatically. Ring and hybrid paths require their own
boundary parity. AllGather-KV and specific backend combinations may fail fast;
keep that behavior visible in docs.
## Sparse attention
Sparse attention is a backend/model/platform optimization, not a synonym for
low-precision dense attention. Enable it only after dense BF16 parity.
For every candidate task/shape/topology:
1. State the sparsity policy, selected backend/kernel, block/head dimensions,
supported hardware, external kernel/package revision, and fallback.
2. Prove the sparse kernel executed and report realized sparsity; an enabled
flag is insufficient.
3. Compare same-seed per-block outputs, denoise trajectory, and final artifact
with dense attention.
4. Add modality-specific metrics: temporal consistency and audio/video sync for
video, spectral/audio quality for audio, and prompt/image metrics as relevant.
5. Report attention time and full E2E time. Include preprocessing/index-build
overhead, peak memory, and the exact dense control.
6. Validate each task and shape. Do not extrapolate a T2V result to I2V/Ref2V,
or NPU sparse support to CUDA/ROCm/XPU.
If Ring bypasses the selected backend or a sparse setting would be silently
ignored, reject the combination. Prefer a visible dense fallback only when the
recipe labels the row `limited` and does not claim sparse speed.
## Fusion-first operator patterns
Prefer current shared/custom operators with explicit guards and native
fallbacks. Prove parity at operator, block, trajectory, and artifact levels,
then use a trace to prove fewer launches/bytes or lower latency.
### Q/K RMSNorm + RoPE
```python
from vllm_omni.diffusion.layers.norm import RMSNorm
from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk
self.norm_q = RMSNorm(head_dim, eps=eps)
self.norm_k = RMSNorm(head_dim, eps=eps)
# Example only: set these booleans from the official checkpoint contract.
use_neox_layout = False
cos_sin_are_half_dim = True
self.rope = RotaryEmbedding(
is_neox_style=use_neox_layout,
half_head_dim=cos_sin_are_half_dim,
)
q = self.norm_q(q)
k = self.norm_k(k)
q, k = apply_rope_to_qk(self.rope, q, k, (cos, sin))
```
Replace the illustrative booleans with the official layout in real code.
Preserve partial rotary dimensions by splitting the rotated prefix and concatenating
the pass-through suffix. If 3D cos/sin includes a batch axis, the current
shared path assumes the values are shared across the batch; per-sample
different positions require a proven representation/path.
For packed `[tokens, heads, head_dim]` tensors with a non-interleaved
`[cos | sin]` table, the current shared fused boundary is:
```python
from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
q, k = fused_qk_norm_rope(
q,
k,
self.norm_q.weight,
self.norm_k.weight,
rope_table,
self.norm_q.variance_epsilon,
)
```
Its CUDA fast path currently specializes BF16 `head_dim=128` and
`rotary_dim=96`; other valid inputs take the eager implementation. The separate
shared RMSNorm and RotaryEmbedding still provide the general cross-platform
baseline. Treat NPU/MUSA and other vendor fusions as independent qualification
rows, and do not generalize a model-specific Triton kernel without compatible
shape, layout, frequency partition, dtype, and platform evidence.
### Fused QKV projection
Use `QKVParallelLinear` as shown above. Qualification includes:
- `num_heads` and `num_kv_heads` divisibility by TP;
- local split sizes from the layer, never reusing global counts after TP;
- Q/K/V checkpoint order conversion before the vLLM loader;
- all Q, K, and V shards present;
- quantization scale/loader mapping when FP8 or another method is enabled.
### Gate-up projection + fused SwiGLU
```python
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.linear import (
MergedColumnParallelLinear,
RowParallelLinear,
)
self.gate_up_proj = MergedColumnParallelLinear(
input_size=hidden_size,
output_sizes=[intermediate_size, intermediate_size],
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.gate_up_proj",
)
self.act_fn = SiluAndMul()
self.down_proj = RowParallelLinear(
input_size=intermediate_size,
output_size=hidden_size,
bias=False,
input_is_parallel=True,
quant_config=quant_config,
prefix=f"{prefix}.down_proj",
)
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(x)
```
Map both checkpoint `gate_proj` and `up_proj` shards into the merged parameter
and require both. A merged projection followed by `F.silu(gate) * up` is not
the same claim as the `SiluAndMul` fused activation kernel; identify what the
trace actually uses.
### AdaLN
Use a shared layer only when its mathematical/return contract matches the
official model:
```python
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
self.adaln = AdaLayerNorm(
hidden_size,
elementwise_affine=False,
eps=1e-6,
)
x = self.adaln(x, scale[:, None, :], shift[:, None, :])
```
`AdaLayerNormZero`, `AdaLayerNormZeroSingle`, and continuous variants include
different projections, broadcasts, and tuple outputs; they are not drop-in
replacements. A shared AdaLN may have an NPU fast path while using native ops on
CUDA/HIP/XPU. Do not call it a universal fused AdaLN kernel.
For sigma schedules with a small finite set of timesteps, profile whether
precomputing the timestep/AdaLN projection once per request is worthwhile.
Cache keys must include the exact schedule/config/device/dtype and stay bounded.
### Layout, residual, and other candidates
After profiling, inspect:
- bias/dropout/add or gated residual chains;
- norm + linear or norm + modulation boundaries;
- repeated transpose/reshape/contiguous around attention;
- VAE normalization/activation/convolution boundaries;
- MoE routing and fused experts for models that actually use MoE.
For fused residual-plus-norm or modulation chains, preserve the unfused
materialization boundary. If the residual was stored in BF16 before RMSNorm,
round to BF16 and promote again inside the fused kernel; using the unrounded
FP32 temporary changes the denoising trajectory even when the stored residual
looks correct.
Keep DLO materialize/release hooks, process-group collectives, request cache
transitions, and cleanup outside compiled/fused regions unless lifecycle safety
is explicitly tested. Compile only stable regions and retain eager fallback.
## Remove redundant computation
### Hoist request invariants
Audit the denoise loop for work that depends only on the request rather than
the current timestep:
```python
# Before the denoise loop.
condition = self.condition_proj(text_embeddings)
refined_condition = self.token_refiner(condition, packed_metadata)
rope_tables = self.build_rope_tables(position_ids)
attn_metadata = self.build_attention_metadata(layout)
for step, sigma in enumerate(sigmas):
timestep_state = self.prepare_timestep_state(sigma)
noise = self.transformer(
latents,
refined_condition,
rope_tables=rope_tables,
attn_metadata=attn_metadata,
timestep_state=timestep_state,
)
latents = self.scheduler_step(noise, sigma, latents)
```
Candidates include prompt conditioning, token-refiner output, reference
encoding/scans, position IDs/RoPE tables, packed masks/metadata, static CFG
branches, and finite-schedule AdaLN projections. First prove they are invariant
for all official tasks, LoRA/adapters, prompt-update features, shapes, and
schedules. In continuous batching, store request-specific results on request
state rather than the pipeline singleton.
### Eliminate per-layer synchronization
Avoid device-to-host scalar extraction such as `.item()` inside every layer:
```python
# Bad: forces a GPU/CPU synchronization in every attention layer.
max_len = int(cu_seqlens[-1].item())
# Better: compute validated host metadata once during request packing and pass it.
packed = PackedLayout(
cu_seqlens=cu_seqlens,
max_seqlen=max(sample_lengths),
)
for block in self.blocks:
hidden = block(hidden, packed_layout=packed)
```
Trace to verify the synchronization disappears. Do not replace a correct
dynamic value with a stale constant merely to remove `.item()`.
### Remove unused masks and allocations
If packed attention consumes `cu_seqlens` and ignores a dense mask, do not
construct the dense mask. Reuse immutable request metadata and preallocated
buffers when safe. Prefer direct indexed scatter/gather over constructing large
temporary tensors only when duplicate-index semantics and gradients (normally
disabled for inference) remain correct.
Any cross-request cache needs a bounded lifetime and a complete key containing
shape, layout, device, dtype, schedule, task, adapter/LoRA, and any value that
changes the tensor. Otherwise keep it request-scoped.
### VAE and reference work
Profile duplicate VAE decode/postprocess, repeated reference decoding/scans,
unnecessary host-device round trips, and full-frame materialization. Apply VAE
tiling, slicing, patch parallelism, or on-demand staging only with seam,
ordering, color/range, and memory parity tests.
## USP communication optimization
“USP works” is not evidence that USP communication is optimized. For Ulysses,
Ring, hybrid, and AllGather-KV separately, profile:
- collective count, bytes, duration, stream, and overlap with compute;
- Q/K/V and output all-to-all boundaries;
- pack/unpack, padding, metadata broadcast, transpose, and contiguous copies;
- head divisibility or advanced uneven-head support;
- scaling efficiency from 1 to N devices on the same workload/interconnect;
- load balance for unequal packed samples.
Declare `_sp_plan` split/gather boundaries around meaningful module outputs,
not arbitrary Python statements. Keep packed metadata sharded/gathered in the
same ownership model as tokens.
When the target revision offers accelerated Ulysses transport such as
`--ulysses-a2a-permute`, qualify it separately from regular Ulysses. Prove the
strict scatter-heads/gather-sequence layout selected the fast path, record
one-time JIT/readiness cost, prewarm the maximum workspace shape before CUDA
graph capture, keep the grow-only workspace on one stream, and exercise cleanup
before process-group destruction. Compare collective and E2E deltas after the
dense backend and topology are frozen; keep unsupported layouts on an explicit
regular-Ulysses fallback.
For GQA, validate `Q_heads % KV_heads == 0`, pad KV heads first, and derive Q
padding from the original ratio. Exercise both contiguous and row-strided QKV
layouts; a fast copy path must preserve bitwise values and packed boundaries.
Cross-attention sometimes uses replicated text K/V while only video/audio Q is
sequence-sharded. In that specific case, `skip_sequence_parallel=True` can
avoid incorrect or wasteful Q/K/V redistribution:
```python
self.cross_attention = Attention(
num_heads=local_heads,
num_kv_heads=local_kv_heads,
head_size=head_dim,
softmax_scale=head_dim**-0.5,
causal=False,
role="cross",
skip_sequence_parallel=True,
prefix=f"{prefix}.cross_attn",
)
```
Use it only after proving K/V are replicated and Q/output ownership is correct.
Otherwise it disables needed communication. Report whether the optimization
reduces collective bytes/time and improves E2E scaling.
## Text encoder and VAE disaggregation
Disaggregation is a measured architecture decision. First collect:
| Question | Required measurement |
|---|---|
| Is the stage large enough? | Encode/decode share at target concurrency and shape |
| Is transfer economical? | Tensor/media bytes, serialization, connector latency and bandwidth |
| Is there reuse? | Prompt/reference reuse rate, cacheability, fan-out |
| Can it scale independently? | Queue saturation and resource utilization per stage |
| Is the contract stable? | Shape, dtype, mask/tag/position metadata, ordering |
### Text encoder candidate
Define an explicit tensor contract containing hidden states, masks, tags or
position IDs, dtype, shape, task/checkpoint revision, and request ID. Support
disabling the local encoder so the DiT cannot accidentally recompute it. Test
all tasks, prompt lengths, partial TP group membership, retries, backpressure,
abort, connector cleanup, and independent replica scaling.
### VAE candidate
Separate VAE encode and decode contracts if both exist. Define latent shape,
scale/shift, dtype/range, tiling/chunk order, output media ordering, and stage
metadata. Measure latent/media transfer against local decode cost. Test tiling
seams, RNG/determinism where relevant, chunk ordering, codec/postprocess,
backpressure, abort, and worker failure.
Do not label VAE tiling, patch parallelism, compile, or offload as VAE
disaggregation. Likewise, a separate text-encoder TP group is not a separate
stage. If no reusable generic stage/connector exists in the target revision,
produce a feasibility result or RFC and keep the capability `not tested` rather
than inventing a model-local production architecture.
## Performance acceptance
For each accepted optimization, require:
- same fixed workload and correctness artifacts before/after;
- actual fast-path execution proven by logs/trace/counters;
- local metric plus E2E latency/throughput and memory;
- at least three raw fixed-work repetitions after explicit warmup, or a
declared arrival-load sample set large enough for every tail percentile;
- named hardware, backend, topology, dtype, and task scope;
- fallback/rejection tests for unsupported shapes/platforms;
- no regression in abort/error cleanup or long-run memory trend.
Update the feature matrix narrowly. Do not convert a profile-driven direction
into a generic “supported” feature merely because a model-specific PR exists.
@@ -1,464 +0,0 @@
# Production Serving and Validation
## Contents
1. [Execution-mode decision](#execution-mode-decision)
2. [Request-level batching](#request-level-batching)
3. [Step execution and continuous batching](#step-execution-and-continuous-batching)
4. [Abort and cleanup](#abort-and-cleanup)
5. [Mixed-RPS soak](#mixed-rps-soak)
6. [Four CI tracks](#four-ci-tracks)
7. [PR evidence checklist](#pr-evidence-checklist)
## Execution-mode decision
Request batching and step-wise execution solve different problems:
| Mode | Pipeline contract | Scheduling/abort boundary |
|---|---|---|
| Serial request | `forward(DiffusionRequestBatch)` for one request | Whole request; in-denoise preemption is unavailable |
| Request batching | `supports_request_batch = True`; one forward consumes compatible independent requests | Whole batch/request |
| Single-request step | Full `SupportsStepExecution`; `max_num_seqs=1` | Between denoise steps |
| Step continuous batching | Full step contract plus heterogeneous batched-step evidence | Between denoise waves/steps |
Do not expose four stub methods as an opt-out. The current protocol is
structural; if a native pipeline is not step-capable, do not define the four
step methods. If it is capable, implement and test all four.
Treat structural support, a generic request-batch bridge, and model benefit as
three separate claims. Before recommending step execution, name a hypothesis
that request mode cannot satisfy, define a success threshold, and run a matched
request-mode control. If throughput or latency regresses, retain step execution
as functional/limited and keep request mode as the production recommendation.
Choose compatibility keys from every value that changes tensor shapes,
semantics, hooks, or scheduler behavior: task, output shape/count, frames,
steps/schedule, guidance/CFG, dtype/backend, LoRA/adapter, cache quality policy,
and model-specific conditions. Reject or schedule separately rather than
silently using the first request's settings.
## Request-level batching
A request-batched pipeline returns exactly one `DiffusionOutput` per logical
request. This is model-owned pseudocode: `validate_batch_compatibility()` and
`_generate_batch()` are placeholders to implement from the official contract,
not shared APIs to import.
```python
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.worker.request_batch import (
DiffusionRequestBatch,
split_diffusion_output_by_request,
)
class MyPipeline(torch.nn.Module):
supports_request_batch = True
def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
validate_batch_compatibility(req.sampling_params_list, req.prompts)
common = req.sampling_params_list[0]
text_prompts: list[str] | None = [
prompt
if isinstance(prompt, str)
else (prompt.get("prompt") or "")
for prompt in req.prompts
]
prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
req.prompts,
{"prompt_embeds": None, "prompt_mask": None},
)
if prompt_fields["prompt_embeds"] is not None:
text_prompts = None
generators = req.collate_request_generators(
common.num_outputs_per_prompt,
default_generator=None,
)
output = self._generate_batch(
prompt=text_prompts,
prompt_fields=prompt_fields,
generators=generators,
sampling=common,
)
return split_diffusion_output_by_request(
DiffusionOutput(output=output),
req,
num_outputs_per_prompt=common.num_outputs_per_prompt,
)
```
The prompt-field collator only handles tensor fields; extract raw text prompts
and model-specific negative/alias fields separately, as current in-tree batched
pipelines do. Set a raw prompt field to `None` when the matching precomputed
embedding is supplied. The exact aliases and generator device depend on the
model. Use current collators rather than manual `torch.cat` when they cover the
contract; they reject mixed present/missing tensors and incompatible
shape/dtype/device.
Before reading `sampling_params_list[0]`, validate every field that must be
common. Tests must cover:
- two and N requests, independent seeds/generators, one output each;
- mixed prompt lengths and packed/padded boundaries;
- multiple outputs per prompt and correct result slicing;
- compatible versus incompatible shapes, steps, guidance, LoRA, and quality;
- one request failure without cross-request output or state corruption;
- output ordering and request IDs under different completion times.
After the model advertises support, candidate serving is:
```bash
vllm serve '<model>' --omni \
--max-num-seqs 4 \
--request-batch-max-wait-ms 20
```
The wait window trades admission latency for batch formation. Benchmark at 0
and candidate nonzero values under the target arrival process.
## Step execution and continuous batching
### Full protocol skeleton
Use the current shared state types and make every mutable value request-local:
The helpers beginning with `_` and `validate_step_batch()` below are
model-owned placeholders; scheduler initialization is scheduler-specific.
```python
import copy
from collections.abc import Sequence
from typing import Any, ClassVar
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.models.interface import SupportsStepExecution
from vllm_omni.diffusion.worker.input_batch import InputBatch
from vllm_omni.diffusion.worker.utils import StepRequestState
class MyPipeline(torch.nn.Module, SupportsStepExecution):
supports_step_execution: ClassVar[bool] = True
def prepare_encode(
self,
state: StepRequestState,
**kwargs: Any,
) -> StepRequestState:
prepared = self._prepare_request(state.prompt, state.sampling)
state.prompt_embeds = prepared.prompt_embeds
state.latents = prepared.latents
state.timesteps = prepared.timesteps
state.scheduler = copy.deepcopy(self.scheduler)
state.scheduler.set_begin_index(0)
state.step_index = 0
state.extra["task_state"] = prepared.task_state
return state
def denoise_step(
self,
input_batch: InputBatch,
*,
states: Sequence[StepRequestState] | None = None,
**kwargs: Any,
) -> torch.Tensor | None:
validate_step_batch(input_batch, states)
return self.transformer(**build_batched_denoise_inputs(input_batch))
def step_scheduler(
self,
state: StepRequestState,
noise_pred: torch.Tensor,
**kwargs: Any,
) -> None:
timestep = state.current_timestep
if timestep is None:
raise RuntimeError("step_scheduler called without a current timestep")
state.latents = state.scheduler.step(
noise_pred,
timestep,
state.latents,
return_dict=False,
)[0].to(state.latents.dtype)
state.step_index += 1
def post_decode(
self,
state: StepRequestState,
**kwargs: Any,
) -> DiffusionOutput:
return DiffusionOutput(output=self.decode(state.latents))
```
This is a shape guide, not drop-in code: verify the target revision's
`StepRequestState.extra`, scheduler signature, CFG handling, state fields, and
decode contract. Follow a current native pipeline such as Qwen-Image for the
exact revision.
Never use the shared `self.scheduler` to advance a request. Store scheduler,
latents, embeddings, timesteps, masks, RNG/generator, task metadata, cache
policy, adapter, and output assembly state on `StepRequestState` or its
request-local extension. Pipeline fields may hold immutable weights/config and
bounded immutable caches only.
### Qualification sequence
1. Compare request-mode and step-mode at `max_num_seqs=1` for every task using
identical seed/schedule/shape and intermediate latents.
2. Abort queued and in-flight requests at several step boundaries.
3. Increase to two compatible requests; verify independent scheduler/RNG/state
and output order.
4. Increase to the target capacity and mix prompt lengths.
5. Exercise mixed shape, frames, steps/schedule, task, guidance, and quality.
Batch only compatible rows; otherwise split/reject predictably.
6. Inject a failure in encode, denoise, scheduler, and decode independently.
7. Run the production soak and fault matrix.
Candidate commands only after capability validation:
```bash
# Conservative step execution and step-boundary abort.
vllm serve '<model>' --omni \
--step-execution \
--max-num-seqs 1
# Experimental batched-step candidate after multi-request E2E.
vllm serve '<model>' --omni \
--step-execution \
--max-num-seqs 4
```
If an attention backend is only correct for single-request step mode, fail
early or keep the recipe at `max_num_seqs=1`. Do not describe generic runtime
support as model support.
## Abort and cleanup
### Required semantics
For queued and in-flight abort/disconnect:
- emit exactly one terminal aborted/error result;
- do not run `post_decode` after abort unless a documented partial-output mode
explicitly requires it;
- do not increment `step_index` or apply a scheduler result after cancellation;
- remove scheduler/state cache entries and request IDs;
- release latents, prompt/reference tensors, cache hooks, DLO buffers, connector
messages, uploaded/temp files, and output buffers;
- keep other requests alive;
- make cleanup idempotent when disconnect, timeout, and worker error race;
- drop late results/exceptions for already-resolved or cancelled futures
without terminating the shared result-pump thread;
- complete the next valid request with its requested cache/offload policy.
Test cancellation before engine admission, while queued, after encode, during
early/middle/final denoise steps, before decode, and while returning/persisting
output. Include async job cancellation and client socket disconnect.
### Failure injection matrix
| Location | Injection | Assertions |
|---|---|---|
| Validation/upload | bad MIME, count/bytes overflow, truncated media | 400/413, no job/engine submit, no temp leak |
| Load/start | missing fused shard, unsupported topology | startup fails, actionable error, no silent init |
| Encode | exception/OOM/timeout | one error, request state/temp released |
| Denoise | exception/OOM/worker exit | peers handled, hooks/buffers released, health behavior explicit |
| Scheduler | invalid state/shape | no latent/step partial commit |
| Decode/post | exception/codec error | no corrupt successful result, partial output removed |
| Result/IPC | late result after cancel, large split tensor, orphan SHM | pump remains alive, one terminal result, exact ownership cleanup |
| Connector/disagg | timeout/retry/duplicate | idempotency, backpressure, no orphan payload |
After each injection, check server health and issue a known-good request. If
the process must restart, document and test the restart policy; do not claim
transparent recovery.
### Large-output transport
Treat worker-to-engine IPC and engine-to-client transport as separate hops.
For batched video, prove that each logical output does not accidentally retain
or serialize the complete batch storage. Exercise the target revision's shared
memory or artifact-handle path above and below its threshold, including
non-contiguous views, multiple outputs, consumer failure, timeout, abort, and
worker exit. Verify every shared segment/file is released exactly once.
Freeze the media contract at each hop: dtype, value range, shape/layout,
contiguity, request ownership, encoding/color model, payload bytes, and pending
consumers. Device-side float-to-uint8 preparation, D2H/IPC, planar/interleaved
conversion, CPU encode/mux, and delivery off the asyncio event loop are
different optimizations. Validate their route selection and fallbacks
independently. Byte-identical HTTP MP4 output does not prove that a raw offline
tensor kept its dtype or range.
For long-window output, validate segment ordering and audio/video continuity,
then exercise payloads around serializer, shared-memory, HTTP, and codec size
limits. Chunking must reassemble exactly, reject malformed/missing chunks, and
release partial state after abort, timeout, or peer failure.
Qualify chunk production, transport, and public streaming as separate layers.
A VAE callback must define owned representation, frame offsets, final marker,
rank ownership, source/config gate, and collective-safe failure, but it does not
prove SHM/ZMQ delivery, bounded backpressure, cancellation, codec reuse, or
client streaming. Measure time to first owned chunk, first encoded fragment,
steady cadence, mux finalization, and complete client artifact independently.
Measure remote MP4 encoding, serialization/copies, network transfer, and client
materialization separately from denoise/decode. An inference-kernel speedup does
not establish serving throughput if output transport dominates or leaks.
## Mixed-RPS soak
Use a realistic request mix and arrival distribution. Include small/large
shapes, short/long prompts, all advertised tasks, cache tiers, and a small rate
of abort/disconnect/error injections. Run phases below, near, and above measured
saturation so queue/backpressure behavior is visible.
Example serving benchmark phase:
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--base-url http://127.0.0.1:8091 \
--endpoint /v1/videos \
--model '<model>' \
--dataset trace \
--dataset-path /path/to/versioned-production-trace.jsonl \
--task t2v \
--num-prompts 500 \
--request-rate 0.5 \
--max-concurrency 8 \
--warmup-requests 8 \
--warmup-concurrency 8 \
--output-file soak-near-saturation.json
```
Use a task-appropriate trace/media dataset. Run the same phase at rates selected
from an initial saturation curve; a literal `0.5` is illustrative only.
On supported GPU hosts, wrap a test or soak with the repository monitor:
```bash
bash tests/dfx/stability/scripts/resource_monitor.sh run --backend gpu -- \
pytest -s -v tests/dfx/stability/scripts/test_stability_<model>.py -m slow
```
Check the target revision's monitor backends. Do not claim CPU/NPU monitoring
from a reserved CLI choice unless collection is implemented and verified.
Report:
- offered and achieved RPS, success/error/abort rates;
- p50/p95/p99 and queue/service/stage times, with the successful sample count
and arrival process used for each percentile;
- throughput per device and batch/wave utilization;
- per-rank HBM and process-tree PSS trend/slope;
- cache size/hit rate, temp-file/disk growth, open FDs/handles;
- outstanding shared-memory/artifact handles and encoded/network bytes;
- worker/engine health, restarts, OOM/timeouts, backpressure response;
- total duration, request count/mix, raw logs/JSON/monitor artifact hashes.
A short burst or repeated serial curl loop is not a stability test.
## Four CI tracks
Keep the tracks independent so a failure answers a specific question. These
four readiness tracks are not the repository's numbered CI levels. Map them
onto the current L1-L5 taxonomy:
| Readiness track | Repository CI placement |
|---|---|
| Function | L1 unit/logic plus L2 basic online/offline E2E; extend real-model scenarios in L3/L4 |
| Accuracy | L3 real-model gates for important cases; deeper and broader nightly coverage in L4 |
| Performance | Important thresholds in L3 where they fit the time budget; full baselines/regression jobs in nightly L4 |
| Reliability | Cheap rejection/recovery assertions can run earlier, but long stability, fault injection, and reliability suites belong to weekly L5 |
Distributed topology is a test case within the appropriate level, not a
replacement definition for a level.
### 1. Function CI
Cover:
- strict unquantized and advertised quantized loading;
- fused source-shard completeness and prefix routing;
- official task positive boundaries and neighboring negative cases;
- offline, sync, async normalization/parity;
- single-device reference plus every advertised topology smoke;
- output type/count/order and request ID isolation;
- feature rejection paths that must fail early.
### 2. Accuracy CI
Pin official implementation/checkpoint/vLLM-Omni revisions, prompts and media
hashes, seed, scheduler/sigma/steps, dimensions/frames/FPS, guidance, dtype,
backend, and topology. Compare at suitable levels:
- component outputs or selected blocks;
- one or more intermediate denoise latents;
- final image/video/audio artifact metrics;
- temporal consistency and audio/video synchronization where relevant.
Set tolerance per modality/hardware/backend with rationale. Keep failure
artifacts and metadata. A single visual inspection does not define CI.
For quantization trajectory tooling, inspect the target revision's current
`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` and
reuse it when compatible instead of creating an ad hoc comparator.
### 3. Performance CI
Pin workload and best deployment. Include a resident recommended row and a
small-HBM/DLO row when both are production targets. Declare:
- one explicit warmup exclusion and at least three measured fixed-work
repetitions;
- raw result JSON and environment identity;
- median/mean plus range for small fixed-work A/Bs; p50/p95/p99 only for an
arrival-load run with enough samples, plus throughput/RPS, stage time,
HBM/PSS;
- regression metric, threshold, baseline owner, and update procedure;
- correctness/accuracy check linked to the same configuration.
Use the current harness config style, for example:
```bash
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
--test-config-file tests/dfx/perf/tests/test_<model>_vllm_omni.json
```
Verify parser options on the target revision. Do not copy a JSON schema from a
different model without matching its endpoint, task, dataset, and deploy path.
### 4. Reliability CI
Run a bounded CI soak plus a longer scheduled soak. Cover:
- below/near/above-saturation arrival phases;
- queued/in-flight abort and disconnect;
- OOM/timeout and worker/process failure appropriate to the platform;
- DLO/cache/connector cleanup;
- health, fast-fail/backpressure, restart/recovery policy;
- memory/temp/FD trend and successful known-good request after each fault.
Keep reliability thresholds independent from latency thresholds. A fast server
that leaks memory or hangs after cancellation fails production readiness.
## PR evidence checklist
The PR description must include or link:
- official source/checkpoint revisions and API/limit matrix;
- scoped feature compatibility matrix with the four allowed states;
- implementation notes for FP8, DLO, request cache, attention/fusion, batching;
- exact environment, deploy/serve, every task curl, benchmark and soak commands;
- fixed accuracy inputs and artifacts;
- raw before/after performance runs and memory traces;
- output-contract manifests and stage accounting from device preparation
through the complete client artifact;
- abort/disconnect/fault results and cleanup evidence;
- Function/Accuracy/Performance/Reliability CI files and ownership;
- limitations/TODOs that remain `not tested` rather than implied support.
Separate the Day-0 milestone from production readiness. If the PR only closes
some gates, title and description should say which scoped gates it closes.
-4
View File
@@ -37,10 +37,6 @@ self-check.
- [`precheck-pr`](precheck-pr/SKILL.md): self-checks a branch before creating a - [`precheck-pr`](precheck-pr/SKILL.md): self-checks a branch before creating a
PR by validating title format, dead code, simplification opportunities, PR by validating title format, dead code, simplification opportunities,
accuracy and performance claims, and merge readiness accuracy and performance claims, and merge readiness
- [`production-add-diffusion-model`](production-add-diffusion-model/SKILL.md):
takes a working Day-0 diffusion integration through API/limit parity,
feature-combination evidence, quantization/offload, performance, production
serving, hardware recipes, and CI readiness
- [`quantization`](quantization/SKILL.md): guides quantization method selection, - [`quantization`](quantization/SKILL.md): guides quantization method selection,
model integration, checkpoint loading, and quality/performance validation model integration, checkpoint loading, and quality/performance validation
for vLLM-Omni for vLLM-Omni
+3 -5
View File
@@ -16,7 +16,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
*Latest News* 🔥 *Latest News* 🔥
- [2026/08] We released [0.28.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.28.0), featuring production-ready [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) serving on GPU and NPU, a unified AR/DiT paged KV cache runtime, and enhanced realtime full-duplex serving for the [MiniCPM-o series](recipes/OpenBMB/MiniCPM-o-4_5.md).
- [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0). - [2026/08] [VeRL-Omni](https://github.com/verl-project/verl-omni) `v0.2.0` is released: faster diffusion RL powered by vLLM-Omni (request-level/step-wise batching with FA3), rebuilt Qwen3-Omni multimodal training (DPO & GSPO), plus LTX-2.3, Qwen-Image-Edit support and more. See the [release notes](https://github.com/verl-project/verl-omni/releases/tag/v0.2.0).
- [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support. - [2026/08] We released [0.26.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.26.0) - aligned with the vLLM 0.26 release line, featuring [MiniMax H3](recipes/MiniMaxAI/MiniMax-H3.md) joint video/audio generation, an experimental full-duplex realtime runtime for [MiniCPM-o 4.5](recipes/OpenBMB/MiniCPM-o-4_5.md), distributed layerwise diffusion offload, and broader model, hardware, streaming, TTS, and quantization support.
- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support. - [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support.
@@ -58,9 +57,9 @@ vLLM-Omni is flexible and easy to use with:
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including: vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL) - **Omni-modality models** (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
- **TTS models** (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3) - **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2) - **Diffusion models** — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
- **Robot-policy and action models** (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1) - **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
## Getting Started ## Getting Started
@@ -68,7 +67,6 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
- [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/) - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
- [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/) - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
- [Nightly CUDA Docker images](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly) built from the latest `main`
- [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/) - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving - [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
@@ -1,156 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Benchmark diffusion LoRA expand-and-accumulate implementations.
This compares the previous functional expand path with direct accumulation into
the selected output slice. It is a synthetic kernel microbenchmark and does not
load a diffusion model.
Example:
python benchmarks/kernels/benchmark_diffusion_lora_expand.py \
--m 1,16,128,512,4096,16384 \
--in-dim 4096 --out-dim 4096 --rank 64 \
--dtype bfloat16 --packed --warmup 30 \
--small-iterations 500 --large-iterations 100
"""
import argparse
import math
import statistics
from collections.abc import Callable
import torch
MiB = 1024**2
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--m", default="1,16,128,512,4096,16384")
parser.add_argument("--in-dim", type=int, default=4096)
parser.add_argument("--out-dim", type=int, default=4096)
parser.add_argument("--rank", type=int, default=64)
parser.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
parser.add_argument(
"--packed",
action="store_true",
help="Accumulate into the second half of a packed output tensor.",
)
parser.add_argument("--warmup", type=int, default=30)
parser.add_argument("--small-iterations", type=int, default=500)
parser.add_argument("--large-iterations", type=int, default=100)
parser.add_argument("--large-m-threshold", type=int, default=1024)
return parser.parse_args()
def percentile(samples: list[float], probability: float) -> float:
ordered = sorted(samples)
position = probability * (len(ordered) - 1)
lower = math.floor(position)
upper = math.ceil(position)
if lower == upper:
return ordered[lower]
weight = position - lower
return ordered[lower] * (1 - weight) + ordered[upper] * weight
def measure_latency_ms(operation: Callable[[], None], warmup: int, iterations: int) -> list[float]:
for _ in range(warmup):
operation()
torch.accelerator.synchronize()
starts = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)]
ends = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)]
for start, end in zip(starts, ends, strict=True):
start.record()
operation()
end.record()
torch.accelerator.synchronize()
return [start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)]
def measure_extra_peak_mib(operation: Callable[[], None]) -> float:
torch.accelerator.synchronize()
torch.accelerator.reset_peak_memory_stats()
allocated = torch.accelerator.memory_allocated()
operation()
torch.accelerator.synchronize()
peak = torch.accelerator.max_memory_allocated()
return (peak - allocated) / MiB
def print_header(args: argparse.Namespace, dtype: torch.dtype) -> None:
device_name = torch.cuda.get_device_name()
print(f"device={device_name} torch={torch.__version__} dtype={dtype}")
print(f"in_dim={args.in_dim} out_dim={args.out_dim} rank={args.rank} packed={args.packed} warmup={args.warmup}")
print("M iters old_p10 old_p50 old_p90 new_p10 new_p50 new_p90 speedup old_MiB new_MiB max_abs")
def benchmark_shape(m: int, args: argparse.Namespace, dtype: torch.dtype) -> None:
device = torch.device("cuda")
output_width = args.out_dim * (2 if args.packed else 1)
offset = args.out_dim if args.packed else 0
output_slice = slice(offset, offset + args.out_dim)
scale = 1 / math.sqrt(args.in_dim)
x = torch.randn(m, args.in_dim, device=device, dtype=dtype)
lora_a = torch.randn(args.rank, args.in_dim, device=device, dtype=dtype) * scale
lora_b = torch.randn(args.out_dim, args.rank, device=device, dtype=dtype) / math.sqrt(args.rank)
initial_output = torch.randn(m, output_width, device=device, dtype=dtype)
def baseline(output: torch.Tensor) -> None:
buffer = x @ lora_a.t()
delta = buffer @ lora_b.t()
selected = output[:, output_slice]
selected[:] = selected + delta
def optimized(output: torch.Tensor) -> None:
buffer = x @ lora_a.t()
selected = output[:, output_slice]
torch.addmm(selected, buffer, lora_b.t(), out=selected)
reference = initial_output.clone()
candidate = initial_output.clone()
baseline(reference)
optimized(candidate)
torch.testing.assert_close(candidate, reference, rtol=1e-2, atol=1e-2)
max_abs = (candidate - reference).abs().max().item()
old_output = initial_output.clone()
new_output = initial_output.clone()
iterations = args.large_iterations if m >= args.large_m_threshold else args.small_iterations
old_peak = measure_extra_peak_mib(lambda: baseline(old_output))
new_peak = measure_extra_peak_mib(lambda: optimized(new_output))
old_samples = measure_latency_ms(lambda: baseline(old_output), args.warmup, iterations)
new_samples = measure_latency_ms(lambda: optimized(new_output), args.warmup, iterations)
old_p50 = statistics.median(old_samples)
new_p50 = statistics.median(new_samples)
print(
f"{m:<7} {iterations:<6} "
f"{percentile(old_samples, 0.1):>8.5f} {old_p50:>8.5f} {percentile(old_samples, 0.9):>8.5f} "
f"{percentile(new_samples, 0.1):>8.5f} {new_p50:>8.5f} {percentile(new_samples, 0.9):>8.5f} "
f"{old_p50 / new_p50:>8.3f} {old_peak:>8.2f} {new_peak:>8.2f} {max_abs:>8.5f}"
)
def main() -> None:
args = parse_args()
if not torch.cuda.is_available():
raise RuntimeError("This benchmark requires a CUDA GPU")
if args.warmup < 0 or args.small_iterations <= 0 or args.large_iterations <= 0:
raise ValueError("warmup must be non-negative and iteration counts must be positive")
dtype = getattr(torch, args.dtype)
torch.manual_seed(0)
print_header(args, dtype)
with torch.inference_mode():
for m in (int(value) for value in args.m.split(",")):
if m <= 0:
raise ValueError(f"M must be positive, got {m}")
benchmark_shape(m, args, dtype)
if __name__ == "__main__":
main()
-11
View File
@@ -132,14 +132,3 @@ models:
endpoint: /v1/audio/speech endpoint: /v1/audio/speech
task_extra_body: task_extra_body:
voice_design: {} voice_design: {}
# OmniVoice serves plain text-to-speech over /v1/audio/speech; voice cloning
# and voice design are offline-only for this model, so default_voice is the
# only task the server can answer.
k2-fsa/OmniVoice:
supported_tasks: [default_voice]
backend: openai-audio-speech
endpoint: /v1/audio/speech
trust_remote_code: true
task_extra_body:
default_voice: {}
+3 -3
View File
@@ -57,9 +57,9 @@ vLLM-Omni is flexible and easy to use with:
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including: vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL) - Omni-modality models (e.g. Qwen3-Omni, MiniCPM-o 4.5, Cosmos3, HunyuanImage, BAGEL)
- TTS models (e.g. Qwen3-TTS, IndexTTS 2.5, dots.tts, CosyVoice3) - TTS models (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
- Diffusion models — image, video, and audio generation (e.g. MiniMax H3, LTX-2.5, SANA-Video, Wan2.2) - Diffusion models — image, video, and audio generation (e.g. MiniMax H3, Qwen-Image, Wan2.2, FLUX)
- Robot-policy and action models (e.g. π0, GR00T-N1.7, DreamZero-DROID, InternVLA-A1) - Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy)
For more information, checkout the following: For more information, checkout the following:
+1 -1
View File
@@ -2,7 +2,7 @@
This section lists the most common options for running vLLM-Omni. This section lists the most common options for running vLLM-Omni.
For options within a vLLM Engine, please refer to the [vLLM 0.28 configuration guide](https://docs.vllm.ai/en/v0.28.0/configuration/index.html). For options within a vLLM Engine, please refer to the [vLLM 0.26 configuration guide](https://docs.vllm.ai/en/v0.26.0/configuration/index.html).
Each model defines fixed topology in a registered `PipelineConfig` and runtime Each model defines fixed topology in a registered `PipelineConfig` and runtime
overrides in a deploy YAML. overrides in a deploy YAML.
+3 -3
View File
@@ -151,7 +151,7 @@ their keys only.
## Inherited vLLM variables ## Inherited vLLM variables
vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to
the [vLLM 0.28 environment-variable reference](https://docs.vllm.ai/en/v0.28.0/configuration/env_vars.html) the [vLLM 0.26 environment-variable reference](https://docs.vllm.ai/en/v0.26.0/configuration/env_vars.html)
for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm, for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm,
XPU, ModelScope, and FlashInfer workspace settings. XPU, ModelScope, and FlashInfer workspace settings.
@@ -179,7 +179,7 @@ collection, examples, bug reports, or logs.
## Model-specific variables ## Model-specific variables
The audit found 58 variables read by a single model or pipeline family. They are The audit found 56 variables read by a single model or pipeline family. They are
not listed as public usage options here because doing so would turn implementation not listed as public usage options here because doing so would turn implementation
escape hatches into an accidental compatibility contract. escape hatches into an accidental compatibility contract.
@@ -188,7 +188,7 @@ Every audited model-specific name has a migration disposition in the
| Disposition | Count | Required outcome | | Disposition | Count | Required outcome |
| --- | ---: | --- | | --- | ---: | --- |
| Promote | 35 | Move a stable setting into typed stage or model configuration. | | Promote | 33 | Move a stable setting into typed stage or model configuration. |
| Request scope | 6 | Move request-varying behavior into a declared request-option schema. | | Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
| External | 0 | Retain only when a supported third-party library owns the contract. | | External | 0 | Retain only when a supported third-party library owns the contract. |
| Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. | | Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |
+10 -26
View File
@@ -57,11 +57,13 @@ 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. | | 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. | | 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 comparison
| Platform | Bootstrap entry | Test job files | Upload mechanism | Job hardware in YAML | | 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) | omit `mirror_hardwares` (from `-m`) / preset string + `MIRROR_HW` | | **CUDA** | `cuda/pipeline.yml` | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml`, `test-weekly.yml` | `upload_pipeline.py --upload` (expands uploader-only keys) | `mirror_hardwares: <preset>` (string) |
| **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` | | **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) | | **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 | | **Intel** | `intel/scripts/bootstrap-intel-omni.sh` | `intel/pipeline-intel.yml` (steps inline) | Direct `pipeline upload` | Inline `agents.queue` on each step |
@@ -90,23 +92,7 @@ Canonical layout (prefer these paths for new changes):
**Upload:** `upload_pipeline.py --upload` expands uploader-only keys before Buildkite upload. **Upload:** `upload_pipeline.py --upload` expands uploader-only keys before Buildkite upload.
**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**. **Hardware in YAML:** `mirror_hardwares: <preset>` (string)—preset names in [`common/ci_mirror_hardwares.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/ci_mirror_hardwares.yml). Do **not** set `agents` / `plugins` on the same step.
```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** **Conventions**
@@ -117,9 +103,9 @@ Canonical layout (prefer these paths for new changes):
**Adding a job** **Adding a job**
1. Pick the level file (`test-ready.yml` / `test-merge.yml` / `test-nightly.yml` / `test-weekly.yml`). 1. Pick the level file (ready / merge / nightly / weekly).
2. Add a step under the right **group** (usually **E2E Test** for model pytest). 2. Add a step under the right **group** (usually **E2E Test** for model pytest).
3. Set `label`, `commands`, `mirror_hardwares`, `depends_on: upload-ready-pipeline` (or `upload-merge-pipeline` / `upload-nightly-pipeline` / `upload-weekly-pipeline`). 3. Set `label`, `commands`, `mirror_hardwares`, `depends_on: upload-<level>-pipeline`.
4. For L2/L3 E2E, add `source_file_dependencies` (pytest + model + deploy YAML prefixes). 4. For L2/L3 E2E, add `source_file_dependencies` (pytest + model + deploy YAML prefixes).
5. Dry-run: 5. Dry-run:
@@ -428,12 +414,10 @@ settings in `pyproject.toml`: online tests launch the server with
`subprocess.Popen` and stop it with SIGTERM, and without those settings the XML `subprocess.Popen` and stop it with SIGTERM, and without those settings the XML
reflects only the pytest parent process, not the server's code paths. reflects only the pytest parent process, not the server's code paths.
The upload depends on `buildkite-agent` being callable inside the container. The upload depends on `buildkite-agent` being callable inside the container. The
CUDA (`l4_*`, `h100_*`) and NPU (`*_npu_*`) presets all use the `kubernetes` `kubernetes` presets (`h100_*`, `*_npu_*`) provide it; the `docker` ones (`l4_*`)
plugin, which provides the agent. `l4_*` jobs run on the EKS `l4-k8s` queue only do because they set `mount-buildkite-agent: true`. A new docker preset that
(agent-stack-k8s); they no longer use the docker plugin or runs a coverage job needs the same.
`mount-buildkite-agent`. A new **docker** preset that runs a coverage job
still needs `mount-buildkite-agent: true`.
List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's
`source_file_dependencies` — both change what the job measures, so without them a `source_file_dependencies` — both change what the job measures, so without them a
+1 -10
View File
@@ -136,7 +136,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<a href="test_writing_guide.md#l4-level-testing-full-functionality-performance-and-documentation-testing">L4</a><br> <a href="test_writing_guide.md#l4-level-testing-full-functionality-performance-and-documentation-testing">L4</a><br>
L4: Purpose, Test Content, Directory Location, Example L4: Purpose, Test Content, Directory Location, Example
</td> </td>
<td>Nightly / Days before Release<br>(additional GPU SKUs)</td> <td>Nightly</td>
<td>GPU</td> <td>GPU</td>
</tr> </tr>
<tr> <tr>
@@ -175,15 +175,6 @@ Before entering specific testing levels, the project establishes two common spec
## Notes ## 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) ### 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)). 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)).
+55 -162
View File
@@ -1,197 +1,90 @@
# Skip-Softmax # Skip-Softmax
Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. It implements Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
[BLASST](https://arxiv.org/abs/2512.12087) (Dynamic BLocked Attention Sparsity via Softmax and how to pick an operating point — is in
Thresholding), which adds a per-tile skip test to the FlashAttention main loop. This page describes the algorithm, how the user configuration is resolved [TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax).
into the value the kernel consumes, and how checkpoint calibration and timestep gating participate The shared selector contract is documented in
in that resolution. Configuration keys and operating-point guidance are covered in [Diffusion Attention Backend Selection](attention_backend_selection.md). This
[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax); the page explains the algorithm.
backend selection contract is covered in
[Diffusion Attention Backend Selection](attention_backend_selection.md).
## Motivation ## Motivation
In a long attention row (a video DiT can have tens of thousands of keys), the softmax weight In a long attention row (a video DiT can have tens of thousands of keys), the softmax weight
concentrates on a small fraction of the keys; the rest receive near-zero weight and barely move the concentrates on a small fraction of the keys; the rest receive near-zero weight and barely move the
output. Computing softmax and the value-weighted sum over those keys is wasted work. Skip-Softmax output. Computing softmax and the value-weighted sum over those keys is wasted work. Skip-Softmax
detects, per tile of keys, when a tile cannot matter and skips its softmax and its value multiply. detects, per block of keys, when a block cannot matter and skips its softmax and its value
multiply.
It is approximate: a skipped tile still carries a small non-zero contribution, so the mode is It is approximate: a skipped block still carries a small non-zero contribution, so the mode is
opt-in and off by default. opt-in and off by default.
## Algorithm ## The online-softmax pass
![The BLASST algorithm: FlashAttention with a per-tile skip test on the local maximum](../figures/skip_softmax/blasst_algorithm.jpg) Attention is computed in a single streaming pass over the keys, in tiles of 128. Per query row the
FlashAttention computes attention in a single streaming pass over KV tiles. Per query row the
kernel maintains three running values: kernel maintains three running values:
- `m` — the largest score seen so far, - `m` — the largest score seen so far,
- `l` — the running denominator `Σ exp(sⱼ − m)`, - `l` — the running denominator `Σ exp(sⱼ − m)`,
- `O` — the running numerator `Σ exp(sⱼ − m)·vⱼ`, - `O` — the running numerator `Σ exp(sⱼ − m)·vⱼ`,
and returns `O / l` at the end. For each KV tile it computes the scores `QK_j^T`, updates `m`, and and returns `O / l` at the end. For each key tile it computes the tile's scores `Q · K_jᵀ · scale`,
accumulates the tile's contribution into `l` and `O`. Rescaling `l` and `O` when `m` grows keeps the updates `m`, and accumulates that tile's contribution into `l` and `O`. Rescaling `l` and `O` when
dense pass numerically exact. `m` grows keeps the single pass numerically exact.
A tile is a block of query rows against a block of keys, and the kernel processes it as a unit. ## The skip test
BLASST inserts one test between computing a tile's scores and accumulating them. For each query
row `i` in the tile, let `tile_max[i]` be the row's largest score within the tile, `running_max[i]` Once a tile's scores are known, its largest score `tile_max` is compared against the running
the row's running maximum, and `λ` the effective threshold: maximum:
```text ```text
if all(exp(tile_max[i] - running_max[i]) < λ for every row i in the tile): if exp(tile_max − running_max) < threshold:
skip this tile # do not compute its Softmax or PV contribution skip this tile # do not compute its softmax weights or its P·V accumulation
``` ```
`exp(tile_max[i] − running_max[i])` is an upper bound on the softmax weight any key in the tile can `exp(tile_max − running_max)` is an upper bound on the softmax weight any key in the tile can
receive in row `i`: if even the row's best key in this tile is far below the row's current maximum, receive: if even the tile's best key is far below the current maximum, every key in the tile is
the tile is unimportant for that row. Only when this holds for every row in the tile can the negligible, and both the softmax (the exponentials) and the `P·V` accumulation for that tile can be
Softmax and `PV` work for the tile be skipped, in which case its contribution to `l` and `O` is dropped. The tile's contribution to `l` and `O` is simply skipped.
simply omitted. A larger `λ` makes the test more aggressive. The figure states the same test in
log space, `m̃ − m < ln(λ)`, which is how the kernel evaluates it.
The test needs the tile's scores, so `QK_j^T` always runs and only the Softmax and `PV` work can ## What this bounds
be skipped. Since the two matmuls have the same FLOP count, skipping every eligible tile removes at
most roughly half of the attention arithmetic, and the kernel-level speedup stays well under 2×.
## From configuration to the kernel threshold Two properties of the test shape the achievable speedup:
The FlashInfer kernel does not take `λ` directly. It takes a `threshold_scale_factor` and divides - **`Q · K_jᵀ` always runs.** The test needs `tile_max`, which comes from the tile's scores, so the
it by the KV sequence length of the call: score matmul is never skipped — only the softmax and the `P·V` accumulation are. The score matmul
and the value matmul are comparable in cost, so skipping every eligible tile removes at most
roughly half the attention work; the kernel-level speedup is bounded well under 2×, not unbounded.
- **The decision is per tile, not per key.** A tile is skipped only when *all* of its keys are
collectively negligible; a single important key keeps the whole tile. How many tiles actually
qualify depends on the data and rounds down to tile granularity, so `target_sparsity` selects an
operating point on a calibrated curve — it is not a promise that a fixed fraction of tiles is
skipped.
## The threshold
The per-tile threshold is normalized by sequence length:
```text ```text
λ = threshold_scale_factor / max_kv_len threshold = factor / seqlen
``` ```
vLLM-Omni resolves the factor from one of two mutually exclusive user controls: `factor` comes from one of two sources:
| Control | Factor passed to the kernel | Resulting `λ` | - **Calibrated** — `factor = a · exp(b · target_sparsity)`, where `a` and `b` are fit per model (and
| --- | --- | --- | per expert, for a multi-expert model) so that a given `target_sparsity` lands near that fraction of
| `threshold` | `threshold * max_kv_len` | `threshold` | skipped tiles on the calibration data.
| `target_sparsity=s` | `a * exp(b * s)` | `a * exp(b * s) / max_kv_len` | - **Direct** — `factor = skip_softmax.threshold · seqlen`, the calibration-free path; the user sets
the threshold themselves.
`threshold` is therefore `λ` itself and is independent of sequence length: the same value yields
the same per-tile test at any resolution or frame count. `threshold=0` skips no tiles. Because the
left-hand side of the skip test lies in `(0, 1]`, values in `(0, 1)` are the meaningful range; the
schema accepts any finite non-negative value.
`target_sparsity` selects a point on a curve fitted per model by
[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity),
so that `s` lands near that fraction of skipped tiles on the calibration data. The coefficients
`a` and `b` come from the checkpoint; the achieved sparsity on another prompt, shape, or layer can
differ.
### Skip-Softmax calibration config
ModelOpt stores the calibration in each transformer component's `config.json` under
`sparse_attention_config`. The following is taken from the `transformer/config.json` of
[`nvidia/Wan2.2-T2V-A14B-Diffusers-FP8`](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8),
with the `ignore` list abridged:
```json
{
"sparse_attention_config": {
"config_groups": {
"group_0": {
"algorithm": "skip_softmax",
"targets": ["WanAttention"],
"ignore": ["blocks.0.attn1", "blocks.0.attn2", "blocks.1.attn1", "..."],
"threshold_scale_factor": {
"formula": "a * exp(b * target_sparsity)",
"coefficients": {"a": 2822.6209952221557, "b": 4.388250623520793}
},
"target_sparsity": 0.7,
"disabled_until_timestep": 0.93
}
}
}
}
```
A checkpoint with several transformer components carries one such config per component, each
with its own coefficients; in this example, `transformer_2/config.json` holds the low-noise
expert's curve.
vLLM-Omni consumes `formula`, `coefficients`, and `ignore`, and ignores `targets`,
`target_sparsity`, and `disabled_until_timestep`:
- Only the formula `a * exp(b * target_sparsity)` is accepted; a checkpoint with any other formula
is rejected at startup. If a spec requests `target_sparsity` and the checkpoint carries no
calibration, startup fails and the error points to `threshold` as the calibration-free
alternative.
- Each component's curve is applied to that component's attention layers. If the config for a
secondary component such as `transformer_2` cannot be read, that component stays dense and a
warning is logged.
- `ignore` holds fnmatch patterns matched against both the full module name and the name relative
to the component, so `blocks.0.attn1` matches `transformer.blocks.0.attn1` as well as
`transformer_2.blocks.0.attn1`. Matching layers receive no coefficients and therefore stay dense
regardless of the user configuration.
At each attention call, a layer with a `threshold` uses it directly; a layer with
`target_sparsity` and stamped coefficients evaluates the curve; a layer with `target_sparsity` but
no coefficients runs dense. The timestep gate below is applied to the result before it reaches the
kernel.
## Timestep gating ## Timestep gating
The early, high-noise denoising steps set the global structure of the output, and their errors `disabled_until_timestep = D` keeps the mode off during the early, high-noise denoise steps and
propagate through every later step. Keeping just these steps dense markedly improves the quality turns it on once the normalized timestep `t` drops to `t ≤ D` (`t` runs `1.0` → `0.0` over the
of the generated video at a small cost in skipped work. `disabled_until_timestep` implements this: schedule). The early steps set the global structure of the output and their errors propagate through
on each attention call the backend compares the current normalized timestep `t` against the every later step, so keeping them dense costs a few skipped-tile opportunities but protects fidelity.
configured cutoff and passes no skip factor to the kernel while `t` is above it. Skip-Softmax
becomes active on the first step whose `t` is at or below the cutoff and stays active for the rest
of the run, since `t` only decreases.
The default `disabled_until_timestep = 0` is a sentinel rather than a cutoff: the gate is off, the `t` is the scheduler's own timestep divided by `num_train_timesteps`, published by the pipeline via
timestep is never read, and Skip-Softmax runs on every step. Any positive value enables the gate. `DenoiseProgressMixin.record_denoise_step`. It is deliberately not derived from the step index,
`disabled_until_timestep = 1.0` therefore also yields no dense steps on a pipeline that publishes because schedulers space their steps non-uniformly. A pipeline that does not publish a timestep
`t`, but on a pipeline that does not publish it the backend stays dense and logs a warning once stays dense when `disabled_until_timestep` is set.
rather than guessing from the step index.
### Where `t` comes from
The pipeline publishes `t` once per denoising step before running the transformer. It is the
scheduler's own position in `[0, 1]`, decreasing from near `1.0` (pure noise) to `0.0` (clean
sample): scheduler-based pipelines publish the scheduler timestep divided by `num_train_timesteps`,
and rectified-flow pipelines publish the current sigma. `t` is deliberately not the step index
divided by the step count. Schedulers place their steps non-uniformly in `t`, so expressing the
cutoff in `t` ties it to the same noise level across step counts and schedulers.
### Mapping the cutoff to denoising steps
For a run whose published sequence is `t[0], ..., t[N-1]`, the number of dense steps is
```text
dense_steps = count(t[i] > disabled_until_timestep)
skip_softmax_steps = N - dense_steps
```
This count **depends on the schedule**, and schedules differ from model to model: the
scheduler family, the number of steps, and any time shift all change where the steps land in `t`.
The same cutoff can therefore gate a very different fraction of the run on two models. Take
MiniMax-H3 as an example. Its video branch uses a flow-shifted rectified-flow schedule that applies
a shift `s` to `N + 1` uniform positions `u` from `1` to `0`,
```text
t = s * u / (1 + (s - 1) * u)
```
and a large `s` pushes most positions toward `t ≈ 1`. With the default `s = 12` and
`num_inference_steps = 50` (50 sigma points, `N = 49` denoiser forwards), the published sequence
starts `1.000, 0.998, 0.996, 0.995, 0.993, ...` and stays above `0.9` for 28 steps, giving:
| `disabled_until_timestep` | Dense steps | Skip-Softmax steps |
| :---: | ---: | ---: |
| `1.00` | 0 | 49 |
| `0.99` | 6 | 43 |
| `0.97` | 14 | 35 |
| `0.95` | 19 | 30 |
| `0.90` | 28 | 21 |
| `0.86` | 33 | 16 |
These numbers hold only for that schedule. A model with a smaller shift, a different step count,
or a scheduler that is not rectified flow produces a different sequence and a different table; a
request that overrides `flow_shift` or the step count changes it for MiniMax-H3 too.
Choose the cutoff by counting against the schedule actually served, not by reusing a value from
another model.
Binary file not shown.

Before

Width:  |  Height:  |  Size: 562 KiB

+1 -1
View File
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels ### Pre-built wheels
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, 0.26.0, and 0.28.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0; the 0.28.0 wheels are published together with the v0.28.0 release this branch produces. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA" === "NVIDIA CUDA"
@@ -5,7 +5,7 @@
# --8<-- [end:requirements] # --8<-- [end:requirements]
# --8<-- [start:set-up-using-python] # --8<-- [start:set-up-using-python]
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.28.x release line uses vLLM 0.28.x. vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.26.x release line uses vLLM 0.26.x.
!!! note !!! note
PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details. PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See <gh-issue:8420> for more details.
@@ -95,13 +95,6 @@ uv pip install --no-build-isolation --editable .
vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which release of vLLM it is based on. vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which release of vLLM it is based on.
Nightly CUDA images built from the latest `main` are also available:
- [`vllm/vllm-omni:nightly`](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly) is the rolling multi-architecture tag.
- [`vllm/vllm-omni:nightly-<commit>`](https://hub.docker.com/r/vllm/vllm-omni/tags?name=nightly-) provides commit-pinned multi-architecture tags. The newest 14 are retained.
Use a commit-pinned tag when you need a reproducible nightly environment.
Here's an example deployment command that has been verified on 2 x H100's: Here's an example deployment command that has been verified on 2 x H100's:
```bash ```bash
docker run --runtime nvidia --gpus 2 \ docker run --runtime nvidia --gpus 2 \
@@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
```bash ```bash
# vLLM-Ascend image # vLLM-Ascend image
# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.28.0 # Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.26.0
# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 # Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3
docker run --rm \ docker run --rm \
--name vllm-omni-npu \ --name vllm-omni-npu \
--shm-size=1g \ --shm-size=1g \
@@ -28,11 +28,11 @@ docker run --rm \
-v /etc/ascend_install.info:/etc/ascend_install.info \ -v /etc/ascend_install.info:/etc/ascend_install.info \
-v /root/.cache:/root/.cache \ -v /root/.cache:/root/.cache \
-p 8000:8000 \ -p 8000:8000 \
-it quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 bash -it quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 bash
# Inside the container, install vLLM-Omni from source # Inside the container, install vLLM-Omni from source
cd /vllm-workspace cd /vllm-workspace
git clone -b v0.28.0 https://github.com/vllm-project/vllm-omni.git git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git
cd vllm-omni cd vllm-omni
pip install -v -e . --no-build-isolation pip install -v -e . --no-build-isolation
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e . # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -50,11 +50,21 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
# --8<-- [start:installation-main] # --8<-- [start:installation-main]
You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.28.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU. You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
```bash ```bash
# Inside the vLLM-Ascend v0.28.0 container, install vLLM-Omni from main # Pin vLLM and vLLM-Ascend to the v0.26 release line
cd /vllm-workspace git clone -b v0.26.0 https://github.com/vllm-project/vllm.git
cd vllm
VLLM_TARGET_DEVICE=empty pip install -v -e .
cd ..
git clone -b releases/v0.26.0rc https://github.com/vllm-project/vllm-ascend.git
cd vllm-ascend
pip install -v -e .
cd ..
# Install vLLM-Omni from the latest main branch
git clone https://github.com/vllm-project/vllm-omni.git git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni cd vllm-omni
pip install -v -e . --no-build-isolation pip install -v -e . --no-build-isolation
@@ -78,8 +88,8 @@ Supported images as following.
Here's an example deployment command that has been verified on 4 x NPUs: Here's an example deployment command that has been verified on 4 x NPUs:
```bash ```bash
# Atlas A2: quay.io/ascend/vllm-omni:v0.28.0 # Atlas A2: quay.io/ascend/vllm-omni:v0.26.0
# Atlas A3: quay.io/ascend/vllm-omni:v0.28.0-a3 # Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3
docker run --rm \ docker run --rm \
--name vllm-omni-a3 \ --name vllm-omni-a3 \
--shm-size=4g \ --shm-size=4g \
@@ -104,7 +114,7 @@ docker run --rm \
-v /etc/ascend_install.info:/etc/ascend_install.info \ -v /etc/ascend_install.info:/etc/ascend_install.info \
-v ~/.cache:/root/.cache \ -v ~/.cache:/root/.cache \
-p 8091:8091 \ -p 8091:8091 \
-it quay.io/ascend/vllm-omni:v0.28.0-a3 bash -it quay.io/ascend/vllm-omni:v0.26.0-a3 bash
``` ```
!!! tip !!! tip
+1 -1
View File
@@ -95,7 +95,7 @@ th {
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — | | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — | | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — | | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) | | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — | | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — | | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
+32 -62
View File
@@ -2,7 +2,7 @@
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include: vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include:
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Native output: 24 kHz; API output can be resampled to 8 kHz. - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz. - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz. - **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz. - **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
@@ -95,7 +95,7 @@ response.stream_to_file("output.wav")
### Endpoint ### Endpoint
```text ```
POST /v1/audio/speech POST /v1/audio/speech
Content-Type: application/json Content-Type: application/json
``` ```
@@ -105,7 +105,7 @@ Content-Type: application/json
#### OpenAI Standard Parameters #### OpenAI Standard Parameters
| Parameter | Type | Default | Description | | Parameter | Type | Default | Description |
| ----------- | ------ | --------- | ------------- | |-----------|------|---------|-------------|
| `input` | string | **required** | The text to synthesize into speech | | `input` | string | **required** | The text to synthesize into speech |
| `model` | string | server's model | Model to use (optional, should match server if specified) | | `model` | string | server's model | Model to use (optional, should match server if specified) |
| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) | | `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
@@ -115,9 +115,8 @@ Content-Type: application/json
#### vLLM-Omni Extension Parameters #### vLLM-Omni Extension Parameters
| Parameter | Type | Default | Description | | Parameter | Type | Default | Description |
| ----------- | ------ | --------- | ------------- | |-----------|------|---------|-------------|
| `task_type` | string | null (inferred) | TTS task type: CustomVoice, VoiceDesign, or Base. For Qwen3-TTS, `ref_audio` or `ref_text` infers Base when this field is omitted, while an uploaded/precomputed voice always selects Base. Other omitted values select CustomVoice; VoiceDesign must be specified explicitly. | | `task_type` | string | "CustomVoice" | TTS task type: CustomVoice, VoiceDesign, or Base |
| `sample_rate` | integer | model native | Target output sample rate. Qwen3-TTS supports 8000 or 24000 Hz; the model remains native 24 kHz internally and is resampled before encoding. |
| `language` | string | "Auto" | Language (see supported languages below) | | `language` | string | "Auto" | Language (see supported languages below) |
| `instructions` | string | "" | Voice style/emotion instructions | | `instructions` | string | "" | Voice style/emotion instructions |
| `max_new_tokens` | integer | 2048 | Maximum tokens to generate | | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
@@ -131,7 +130,7 @@ Content-Type: application/json
#### Voice Clone Parameters (Base task) #### Voice Clone Parameters (Base task)
| Parameter | Type | Default | Description | | Parameter | Type | Default | Description |
| ----------- | ------ | --------- | ------------- | |-----------|------|---------|-------------|
| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`). Local files fold `mtime_ns` and `size` into cache keys to automatically reload on-disk edits; HTTP URLs and base64 URIs remain cached by string locator. | | `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`). Local files fold `mtime_ns` and `size` into cache keys to automatically reload on-disk edits; HTTP URLs and base64 URIs remain cached by string locator. |
| `ref_text` | string | null | Transcript of reference audio | | `ref_text` | string | null | Transcript of reference audio |
| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) | | `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
@@ -205,7 +204,7 @@ The `usage` object on `speech.audio.done` is the same shape returned per item by
### Voices Endpoint ### Voices Endpoint
```text ```
GET /v1/audio/voices GET /v1/audio/voices
``` ```
@@ -230,7 +229,7 @@ Lists available voices for the loaded model.
`uploaded_voices` is always present (empty list when no custom voices have been uploaded). Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time. `uploaded_voices` is always present (empty list when no custom voices have been uploaded). Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
```text ```
POST /v1/audio/voices POST /v1/audio/voices
Content-Type: multipart/form-data Content-Type: multipart/form-data
``` ```
@@ -240,7 +239,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
**Form Parameters:** **Form Parameters:**
| Parameter | Type | Required | Description | | Parameter | Type | Required | Description |
| ----------- | ------ | ---------- | ------------- | |-----------|------|----------|-------------|
| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) | | `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
| `consent` | string | Yes | Consent recording ID | | `consent` | string | Yes | Consent recording ID |
| `name` | string | Yes | Name for the new voice | | `name` | string | Yes | Name for the new voice |
@@ -289,7 +288,7 @@ The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and
Client -> Server: Client -> Server:
| Message | Description | | Message | Description |
| --------- | ------------- | |---------|-------------|
| `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) | | `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) |
| `{"type": "input.text", "text": "..."}` | Text chunk | | `{"type": "input.text", "text": "..."}` | Text chunk |
| `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open | | `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open |
@@ -298,7 +297,7 @@ Client -> Server:
Server -> Client: Server -> Client:
| Message | Description | | Message | Description |
| --------- | ------------- | |---------|-------------|
| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input | | `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) | | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
| `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input | | `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
@@ -312,17 +311,17 @@ text, emits `session.done`, and then waits on the same connection for the next
utterance, so a client that speaks repeatedly (for example one driven by an utterance, so a client that speaks repeatedly (for example one driven by an
upstream LLM) pays the WebSocket handshake once instead of once per utterance. upstream LLM) pays the WebSocket handshake once instead of once per utterance.
- The session config is sticky. Send `input.text` again straight after * The session config is sticky. Send `input.text` again straight after
`session.done` to reuse it, or send another `session.config` first to change `session.done` to reuse it, or send another `session.config` first to change
voice, format, or reference audio. A `session.config` sent while text is voice, format, or reference audio. A `session.config` sent while text is
still buffered is rejected so no pending input is silently dropped. still buffered is rejected so no pending input is silently dropped.
- An utterance is the flush unit, not a linguistic one: it is whatever text was * An utterance is the flush unit, not a linguistic one: it is whatever text was
buffered when `input.done` arrived, of any length, synthesized as one request. buffered when `input.done` arrived, of any length, synthesized as one request.
`utterance_index` counts those flushes across the connection, so it tells you `utterance_index` counts those flushes across the connection, so it tells you
which `input.done` a frame belongs to. `sentence_index` counts within one which `input.done` a frame belongs to. `sentence_index` counts within one
flush and so pairs with `total_sentences`, which means every utterance reports flush and so pairs with `total_sentences`, which means every utterance reports
`sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer). `sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer).
- End the connection with `session.close`, or by closing the socket. An idle * End the connection with `session.close`, or by closing the socket. An idle
connection is still closed after the server's idle timeout, which now also connection is still closed after the server's idle timeout, which now also
applies to the gap between utterances. applies to the gap between utterances.
@@ -331,9 +330,10 @@ upstream LLM) pays the WebSocket handshake once instead of once per utterance.
All REST API parameters are supported, plus: All REST API parameters are supported, plus:
| Parameter | Type | Default | Description | | Parameter | Type | Default | Description |
| ----------- | ------ | --------- | ------------- | |-----------|------|---------|-------------|
| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket | | `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |
```bash ```bash
DELETE /v1/audio/voices/{name} DELETE /v1/audio/voices/{name}
``` ```
@@ -343,7 +343,7 @@ Delete an uploaded voice sample.
**Path Parameters:** **Path Parameters:**
| Parameter | Type | Required | Description | | Parameter | Type | Required | Description |
| ----------- | ------ | ---------- | ------------- | |-----------|------|----------|-------------|
| `name` | string | Yes | Name of the voice to delete | | `name` | string | Yes | Name of the voice to delete |
**Response Example:** **Response Example:**
@@ -433,7 +433,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
### Upload Voice ### Upload Voice
Upload voice (speaker embedding only): Upload voice (speaker embedding only):
```bash ```bash
curl -X POST http://localhost:8091/v1/audio/voices \ curl -X POST http://localhost:8091/v1/audio/voices \
-F "audio_sample=@/path/to/voice_sample.wav" \ -F "audio_sample=@/path/to/voice_sample.wav" \
@@ -442,7 +441,6 @@ curl -X POST http://localhost:8091/v1/audio/voices \
``` ```
Upload voice with transcript (in-context cloning, higher quality): Upload voice with transcript (in-context cloning, higher quality):
```bash ```bash
curl -X POST http://localhost:8091/v1/audio/voices \ curl -X POST http://localhost:8091/v1/audio/voices \
-F "audio_sample=@/path/to/voice_sample.wav" \ -F "audio_sample=@/path/to/voice_sample.wav" \
@@ -452,7 +450,6 @@ curl -X POST http://localhost:8091/v1/audio/voices \
``` ```
### Use Uploaded Voice ### Use Uploaded Voice
```bash ```bash
curl -X POST http://localhost:8091/v1/audio/speech \ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \ -H "Content-Type: application/json" \
@@ -462,11 +459,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
}' --output cloned.wav }' --output cloned.wav
``` ```
For Qwen3-TTS, an uploaded voice is a Base voice-cloning input. The server
infers `task_type="Base"` when `voice` names an uploaded entry, so the request
must be sent to a Base checkpoint. Built-in presets such as `vivian` and
`ryan` remain CustomVoice speakers and require a CustomVoice checkpoint.
### Voice Storage & Caching ### Voice Storage & Caching
Uploaded voices are persisted to disk as a single `.safetensors` file per voice Uploaded voices are persisted to disk as a single `.safetensors` file per voice
@@ -486,9 +478,6 @@ once.
### Precomputed Custom Voices ### Precomputed Custom Voices
Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup. Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
Qwen3-TTS precomputed voices follow the same Base task and checkpoint-matching
rules as uploaded voices.
Generate a directory containing `custom_voice_manifest.json` plus one Generate a directory containing `custom_voice_manifest.json` plus one
`.safetensors` file per voice, then set the pipeline-wide deploy config field: `.safetensors` file per voice, then set the pipeline-wide deploy config field:
@@ -527,7 +516,7 @@ by `GET /v1/audio/voices`. Valid precomputed voices can be used in
**Configuration (environment variables):** **Configuration (environment variables):**
| Variable | Default | Description | | Variable | Default | Description |
| ---------- | --------- | ------------- | |----------|---------|-------------|
| `SPEAKER_SAMPLES_DIR` | `~/.cache/vllm-omni/speakers` | Directory for persisted uploaded speakers (`.safetensors` files). | | `SPEAKER_SAMPLES_DIR` | `~/.cache/vllm-omni/speakers` | Directory for persisted uploaded speakers (`.safetensors` files). |
| `SPEAKER_MAX_UPLOADED` | `1000` | Maximum number of uploaded speakers kept on disk. Upload requests past the cap return 400. | | `SPEAKER_MAX_UPLOADED` | `1000` | Maximum number of uploaded speakers kept on disk. Upload requests past the cap return 400. |
@@ -539,7 +528,7 @@ The batch endpoint synthesizes multiple texts in a single request, returning all
### Endpoint ### Endpoint
```text ```
POST /v1/audio/speech/batch POST /v1/audio/speech/batch
Content-Type: application/json Content-Type: application/json
``` ```
@@ -547,7 +536,7 @@ Content-Type: application/json
### Request Parameters ### Request Parameters
| Parameter | Type | Default | Description | | Parameter | Type | Default | Description |
| ----------- | ------ | --------- | ------------- | |-----------|------|---------|-------------|
| `items` | array | **required** | List of items to synthesize (1–32) | | `items` | array | **required** | List of items to synthesize (1–32) |
| `model` | string | server's model | Model to use | | `model` | string | server's model | Model to use |
| `voice` | string | null | Default voice for all items | | `voice` | string | null | Default voice for all items |
@@ -564,7 +553,7 @@ Content-Type: application/json
Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set: Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:
| Field | Type | Description | | Field | Type | Description |
| ------- | ------ | ------------- | |-------|------|-------------|
| `input` | string | **required** — text to synthesize | | `input` | string | **required** — text to synthesize |
| `voice` | string | Override voice for this item | | `voice` | string | Override voice for this item |
| `response_format` | string | Override format for this item | | `response_format` | string | Override format for this item |
@@ -696,7 +685,7 @@ for result in response.json()["results"]:
### Configuration ### Configuration
| Parameter | Source | Default | Description | | Parameter | Source | Default | Description |
| ----------- | -------- | --------- | ------------- | |-----------|--------|---------|-------------|
| `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request | | `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request |
All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_num_seqs` and queues the rest — no client-side throttling needed. All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_num_seqs` and queues the rest — no client-side throttling needed.
@@ -719,17 +708,17 @@ The bundled config also sets `initial_codec_chunk_frames: 1`. This emits only th
### Qwen3-TTS ### Qwen3-TTS
| Model | Task Type | Description | | Model | Task Type | Description |
| ------- | ----------- | ------------- | |-------|-----------|-------------|
| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control | | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description | | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio or an uploaded/precomputed voice | | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant | | `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster Base variant for voice cloning, including uploaded/precomputed voices | | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
### Fish Speech S2 Pro ### Fish Speech S2 Pro
| Model | Description | | Model | Description |
| ------- | ------------- | |-------|-------------|
| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. | | `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. |
Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech section of the online TTS hub](../user_guide/examples/online_serving/text_to_speech.md#fish-speech-s2-pro) for details. Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech section of the online TTS hub](../user_guide/examples/online_serving/text_to_speech.md#fish-speech-s2-pro) for details.
@@ -737,31 +726,31 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
### Voxtral TTS ### Voxtral TTS
| Model | Description | | Model | Description |
| ------- | ------------- | |-------|-------------|
| `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. | | `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
### CosyVoice3 ### CosyVoice3
| Model | Description | | Model | Description |
| ------- | ------------- | |-------|-------------|
| `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | Voice cloning from `ref_audio` + `ref_text`. No built-in voice presets — upload a voice or pass `ref_audio`/`ref_text` per request. | | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | Voice cloning from `ref_audio` + `ref_text`. No built-in voice presets — upload a voice or pass `ref_audio`/`ref_text` per request. |
### OmniVoice ### OmniVoice
| Model | Description | | Model | Description |
| ------- | ------------- | |-------|-------------|
| `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. | | `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. |
### VoxCPM2 ### VoxCPM2
| Model | Description | | Model | Description |
| ------- | ------------- | |-------|-------------|
| `openbmb/VoxCPM2` | TTS + voice cloning with built-in speaker presets and uploaded-voice support. Accepts `voice` (preset or uploaded) or `ref_audio` + optional `ref_text`. | | `openbmb/VoxCPM2` | TTS + voice cloning with built-in speaker presets and uploaded-voice support. Accepts `voice` (preset or uploaded) or `ref_audio` + optional `ref_text`. |
### MOSS-TTS-Nano ### MOSS-TTS-Nano
| Model | Description | | Model | Description |
| ------- | ------------- | |-------|-------------|
| `OpenMOSS-Team/MOSS-TTS-Nano` | Voice cloning only. Requires `ref_audio` (or an uploaded `voice`); no built-in voice presets. `ref_text` is accepted but ignored — upstream's `voice_clone` mode does not consume a transcript. | | `OpenMOSS-Team/MOSS-TTS-Nano` | Voice cloning only. Requires `ref_audio` (or an uploaded `voice`); no built-in voice presets. `ref_text` is accepted but ignored — upstream's `voice_clone` mode does not consume a transcript. |
## Error Responses ## Error Responses
@@ -781,19 +770,6 @@ Invalid parameters:
} }
``` ```
Qwen3-TTS task/checkpoint mismatch:
```json
{
"error": {
"message": "Qwen3-TTS CustomVoice checkpoint does not support task_type='Base'. Use task_type='CustomVoice' or load the matching Base checkpoint.",
"type": "BadRequestError",
"param": null,
"code": 400
}
}
```
### 404 Not Found ### 404 Not Found
Model not found: Model not found:
@@ -814,15 +790,10 @@ Model not found:
### "TTS model did not produce audio output" ### "TTS model did not produce audio output"
Ensure you're using the correct model variant for your task type: Ensure you're using the correct model variant for your task type:
- CustomVoice task → CustomVoice model - CustomVoice task → CustomVoice model
- VoiceDesign task → VoiceDesign model - VoiceDesign task → VoiceDesign model
- Base task → Base model - Base task → Base model
Uploaded and precomputed Qwen3-TTS voices select the Base task, even when the
client supplies a different `task_type`; serve them with a Base checkpoint.
Built-in Qwen3-TTS presets remain CustomVoice voices.
### Server Not Running ### Server Not Running
```bash ```bash
@@ -833,7 +804,6 @@ curl http://localhost:8091/v1/audio/voices
### Out of Memory ### Out of Memory
If you encounter OOM errors: If you encounter OOM errors:
1. Use smaller model variant: `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` 1. Use smaller model variant: `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`
2. Reduce `--gpu-memory-utilization` 2. Reduce `--gpu-memory-utilization`
@@ -852,7 +822,7 @@ final-output drain to a condition-variable wakeup at the same time.
**Configuration (environment variables):** **Configuration (environment variables):**
| Variable | Default | Description | | Variable | Default | Description |
| --- | --- | --- | |----------|---------|-------------|
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. | | `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |
Set it on the process that runs the orchestrator (stage 0 of an omni Set it on the process that runs the orchestrator (stage 0 of an omni
@@ -20,7 +20,7 @@ For the internal selector, registry, and platform contract, see
| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) | | Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | | Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) | | Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
| Use FastVideo VSA with FastWan2.2-TI2V-5B or FastH3 MiniMax-H3 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) | | Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
## Backend options ## Backend options
@@ -36,7 +36,7 @@ For the internal selector, registry, and platform contract, see
| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | | `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | | `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) | | `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B and FastH3 MiniMax-H3 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) | | `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
## Configuration ## Configuration
@@ -1,9 +1,8 @@
# Dense Attention Backends # Dense Attention Backends
Use these backends to run dense attention. Optional quantized or sparse modes Use these backends when you want dense attention without backend-specific
remain inactive unless they are explicitly configured. Start with sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark
`TORCH_SDPA` for correctness comparisons, then benchmark the fastest compatible the fastest compatible kernel for your model, shape, and hardware.
kernel for your model, shape, and hardware.
For selection precedence and per-role configuration, see the For selection precedence and per-role configuration, see the
[attention backend overview](../attention_backends.md). [attention backend overview](../attention_backends.md).
@@ -21,11 +20,10 @@ vllm-omni serve <model> --diffusion-attention-backend TORCH_SDPA
## `FLASH_ATTN` ## `FLASH_ATTN`
`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it `FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it
is FA4 only (`flash_attn.cute` / `vllm-omni[fa4]`). If FA4 is unavailable, prefers the CuTe-based FlashAttention 4 path and falls back to
explicit `FLASH_ATTN` is rejected and automatic selection continues to another FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
compatible Blackwell backend. Hopper-only FA2/FA3 wheels are not used, even if Ampere it is the preferred automatic route when a compatible package is
they import. On Hopper, Ada, and Ampere it is the preferred automatic route installed.
when a compatible package is installed.
```bash ```bash
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
@@ -40,25 +38,8 @@ pip install 'vllm-omni[fa4]'
``` ```
Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on
Blackwell. On Blackwell, `FLASH_ATTN` is FA4 only: Hopper-only FA2/FA3 Blackwell. If the CuTe path is unavailable, the backend falls back to the
wheels are not used, even if they import. If FA4 is missing, explicit compatible FlashAttention 3 or 2 path.
`FLASH_ATTN` is rejected and automatic selection continues to other
Blackwell backends.
## `TRTLLM_ATTN`
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels and is the platform
default on datacenter Blackwell for models that declare a compatible attention
path. Selected without a `quant` or `skip_softmax` block, it runs dense BF16
attention at FA4-level performance.
```bash
vllm serve <model> --omni \
--diffusion-attention-backend TRTLLM_ATTN
```
See [TRTLLM Attention](trtllm.md) for its requirements and for the optional
SAGE quantization and Skip-Softmax modes.
## `CUDNN_ATTN` ## `CUDNN_ATTN`
@@ -92,13 +73,6 @@ DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
option on CUDA platforms and an automatic Blackwell fallback when FlashInfer option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
is installed but cuDNN is too old for `CUDNN_ATTN`. 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 ```bash
vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN
``` ```
@@ -7,18 +7,11 @@ attention only against the selected top-k blocks.
The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope. The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
MiniMax-H3 served with a FastH3 VSA adapter uses a second, H3-specific route
through the same backend. See the
[MiniMax-H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for its geometry,
supported topologies, and commands.
VSA is a CUDA-only, explicitly selected backend. It requires the VSA is a CUDA-only, explicitly selected backend. It requires the
`fastvideo-kernel` package and currently supports non-causal self-attention `fastvideo-kernel` package and currently supports non-causal self-attention
with equal query and key/value sequence lengths. Unsupported shapes, masks, with equal query and key/value sequence lengths. Unsupported shapes, masks,
dtypes, sequence-parallel execution, or kernel failures fall back to dtypes, sequence-parallel execution, or kernel failures fall back to
`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active `TORCH_SDPA` and emit a warning with the reason.
sequence-parallel context is one of those fallbacks; the H3 route supports pure
Ulysses and rejects ring or all-gather sequence parallelism at startup.
## Enable the backend ## Enable the backend
@@ -126,8 +119,7 @@ the backend guarantees sparse execution:
- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA - `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
executed; the warning includes the reason. executed; the warning includes the reason.
The Wan route requires CUDA tensors in FP16 or BF16, 256-token blocks, The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks,
standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask, standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
and no active sequence-parallel context. The MiniMax-H3 route uses 64-token and no active sequence-parallel context. NPU and XPU paths do not execute the
`(4, 4, 4)` blocks and supports pure Ulysses. NPU and XPU paths do not execute FastVideo VSA CUDA kernel.
the FastVideo VSA CUDA kernel.
@@ -15,9 +15,9 @@ implementations to reduce numerical drift and sampling divergence.
pip install kernels==0.14.1 pip install kernels==0.14.1
``` ```
If `kernels` is unavailable, an explicit `FLASH_ATTN_HUB` or If `kernels` is unavailable, vLLM-Omni warns and falls back to the
`FLASH_ATTN_3_HUB` selection raises. Automatic / `FLASH_ATTN` selection in the corresponding local FlashAttention implementation. On CUDA devices older than
Diffusers adapter may still walk hub then local FlashAttention backends. Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
## Usage ## Usage
@@ -1,274 +1,103 @@
# TRTLLM Attention # TRTLLM Attention
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels on datacenter `TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
Blackwell GPUs. Selected on its own it computes dense BF16 attention, and its BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
dense performance is on par with
[FlashAttention-4](dense_backends.md#flashattention-4-on-blackwell). It also
provides two opt-in, lossy acceleration modes that can be enabled independently
or together:
| Mode | Config key | What it changes | For common selection and per-role configuration, see the
| --- | --- | --- | [attention backend overview](../attention_backends.md).
| [Skip-Softmax](#skip-softmax) | `skip_softmax` | Skips the Softmax and `PV` work of KV tiles whose scores are too low to matter |
| [SAGE quantization](#sage-quantization) | `quant` | Runs `QK^T` in INT8 or FP8 and `PV` in FP8 instead of BF16 |
Both keys are set through `--diffusion-attention-config`, as JSON or as
vLLM-style dotted flags; the
[attention backend overview](../attention_backends.md#configuration) covers
both syntaxes and per-role resolution. The examples on this page use JSON.
Attention computes scores `S = QK^T`, probabilities `P = softmax(S)`, and the
output `O = PV`. SAGE lowers the precision of both matrix multiplications.
Skip-Softmax keeps `QK^T` dense and removes the Softmax and `PV` work for
tiles that `QK^T` shows to be unimportant. The two modes therefore compose:
SAGE makes every tile cheaper, and Skip-Softmax reduces the number of tiles
that reach the second half of the kernel.
## Requirements ## Requirements
`TRTLLM_ATTN` requires all of the following: Selecting `TRTLLM_ATTN` requires:
- an `sm100a` or `sm103a` GPU (B200, B300, GB200, GB300); workstation Blackwell - datacenter Blackwell, sm_100 or sm_103;
(`sm120`/`sm121`) is not supported;
- `head_dim=128`; - `head_dim=128`;
- a FlashInfer build that exposes the trtllm-gen kernels (0.6.16rc1 or newer - FlashInfer; and
for SAGE); - a model-declared packed or mask-free compatible path.
- an attention path that is mask-free or provides packed-padding metadata.
Structural suffix padding is expressed through that metadata rather than an
`attn_mask` tensor.
An explicit selection that violates these requirements raises at startup It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
instead of silently falling back to another backend. explicit selection raises instead of silently changing the backend.
### Sequence parallelism ## Dense mode
`TRTLLM_ATTN` runs under no sequence parallelism or under pure Ulysses. Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
Ulysses redistributes the sequence and attention heads around the attention
call, but the local computation still goes through the configured backend, so
both optional modes work unchanged. Ring and AllGather-KV do not:
- Ring runs its own distributed attention and bypasses the backend. Combining
Ring with a `skip_softmax` key raises; a `quant` key would be silently
ignored, so do not combine them either.
- AllGather-KV changes the Q/KV distribution and is rejected when
`TRTLLM_ATTN` is selected.
Configure eight-way Ulysses alone as:
```bash ```bash
--usp 8 --ring 1 --allgather-degree 1 vllm-omni serve <model> --diffusion-attention-backend TRTLLM_ATTN
```
A degree of `1` disables that sequence-parallel mode. It does not limit the
server to one GPU or affect tensor, pipeline, or VAE parallelism.
## Quick start
On datacenter Blackwell the platform selects `TRTLLM_ATTN` by default when the
model declares a compatible path. To select it explicitly, for dense BF16:
```bash
vllm serve <model> --omni \
--diffusion-attention-backend TRTLLM_ATTN
```
To enable both optimizations with a conservative Skip-Softmax setting and FP8
SAGE:
```bash
vllm serve <model> --omni \
--diffusion-attention-config '{
"default": {
"backend": "TRTLLM_ATTN",
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16},
"skip_softmax": {"threshold": 0.05, "disabled_until_timestep": 0.97}
}
}'
```
The startup log reports the selection; look for one of:
```text
Defaulting to diffusion attention backend TRTLLM_ATTN (datacenter Blackwell ..., head_dim 128)
Resolved diffusion attention backend 'TRTLLM_ATTN' for role='self' via attention_config.default
``` ```
## Skip-Softmax ## Skip-Softmax
Skip-Softmax (BLASST) skips the Softmax and `PV` work of KV tiles whose Skip-Softmax trades some fidelity for speed by skipping low-contribution
scores fall far below the running row maximum. `QK^T` always runs, and how attention work. See the [feature design](../../../design/feature/skip_softmax.md)
many tiles qualify depends on the input. The for the algorithm and implementation contract.
[feature design](../../../design/feature/skip_softmax.md) gives the skip test.
### Configuration keys | Key | Valid values | Meaning |
| Key | Range | Meaning |
| --- | --- | --- | | --- | --- | --- |
| `threshold` | `>= 0`; useful values in `(0, 1)` | Skip threshold, independent of sequence length. Calibration-free. | | `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve |
| `target_sparsity` | `[0, 1]` | Requested operating point on the checkpoint's calibrated curve. Requires calibration metadata. | | `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` |
| `disabled_until_timestep` | `[0, 1]`; default `0` | Keeps attention dense while the normalized timestep `t > D`; `0` disables the gate. | | `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` |
`threshold` and `target_sparsity` are two ways to set the same kernel
threshold; setting both is a configuration error. Exactly one of them enables
Skip-Softmax.
### Direct threshold
Set `threshold` when the checkpoint carries no calibration. `threshold=0`
skips nothing; larger values skip more tiles and lower output fidelity. Values
around `0.05` are a reasonable first try for video DiTs; tune against dense
output on the same prompt and seed. The value is independent of sequence
length; see the
[feature design](../../../design/feature/skip_softmax.md#from-configuration-to-the-kernel-threshold)
for how it reaches the kernel.
```bash ```bash
vllm serve <model> --omni \ vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--diffusion-attention-config \ --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
'{"default":{"backend":"TRTLLM_ATTN","skip_softmax":{ "skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}'
"threshold":0.05,"disabled_until_timestep":0.97}}}'
``` ```
### Calibrated target sparsity ```python
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
A fixed `threshold` does not produce a fixed fraction of skipped tiles because config = AttentionConfig(
the score distribution changes with the model, prompt, and shape. default=AttentionSpec(
[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer) backend="TRTLLM_ATTN",
can calibrate a per-model curve from sparsity to threshold and store it in the skip_softmax=SkipSoftmaxSpec(
checkpoint's transformer `config.json` under `sparse_attention_config`. target_sparsity=0.65,
`target_sparsity` selects a point on that curve. The achieved sparsity still disabled_until_timestep=0.86,
varies per prompt, layer, and denoising step; the calibration makes the ),
requested value a meaningful target, not a guarantee. ),
)
What vLLM-Omni takes from the checkpoint:
- The curve coefficients. Only the `a * exp(b * target_sparsity)` form ModelOpt
writes is supported; another formula is rejected at startup.
- The `ignore` list: attention layers matching those patterns stay dense.
- For multi-expert Diffusers checkpoints, `transformer_2/config.json` is read
separately; if it is missing, `transformer_2` stays dense with a warning.
Checkpoint-level `target_sparsity` and `disabled_until_timestep` defaults are
not consumed; set them in the vLLM-Omni configuration. The
[feature design](../../../design/feature/skip_softmax.md#skip-softmax-calibration-config)
documents the checkpoint format and how the coefficients reach the kernel.
To calibrate a checkpoint yourself, follow the
[ModelOpt Skip-Softmax example](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity),
which exports a copy of the checkpoint with the coefficients embedded. The
fitted curve depends on the attention statistics of the calibration shape, so
calibrate at the resolution and frame count you will serve.
Requesting `target_sparsity` for a checkpoint without calibration is a startup
error that names the `threshold` alternative. The
[ModelOpt Wan2.2 FP8 checkpoint](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8/blob/main/transformer/config.json)
is a calibrated example:
```bash
vllm serve nvidia/Wan2.2-T2V-A14B-Diffusers-FP8 --omni \
--diffusion-attention-config \
'{"default":{"backend":"TRTLLM_ATTN","skip_softmax":{
"target_sparsity":0.75,"disabled_until_timestep":0.86}}}'
``` ```
### Timestep gating Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase
sparsity only after comparing quality against dense output at the same seed.
The early, high-noise denoising steps fix the global layout of the output, and The benefit grows with sequence length because only attention is accelerated.
their errors propagate through every later step. `disabled_until_timestep`
keeps those steps dense and enables Skip-Softmax once the normalized timestep
`t` reaches the configured cutoff. The default `0` is a sentinel that turns the
gate off: Skip-Softmax runs on every step and no timestep needs to be
published. `1.0` also leaves no step dense, but goes through the gate and
therefore requires the pipeline to publish `t`.
`t` is the scheduler's own timestep normalized to `[0, 1]`, published by the
pipeline for each denoising step; for rectified-flow models it is the current
sigma. It runs from near `1.0` down to `0.0`, but not linearly in the step
index: flow-shifted schedules spend many of their steps at high `t`. The number
of dense steps a given cutoff produces therefore depends on the model's
schedule and step count, and follows from the actual timestep sequence:
```text
dense_steps = count(t[i] > disabled_until_timestep)
```
Derive the cutoff from the schedule of the model you are serving rather than
reusing another model's value; the
[feature design](../../../design/feature/skip_softmax.md#mapping-the-cutoff-to-denoising-steps)
works through an example.
A pipeline that does not publish `t` stays dense whenever
`disabled_until_timestep > 0` is set, and logs a warning once. Pipelines
publish it through `DenoiseProgressMixin.record_denoise_step`.
## SAGE quantization ## SAGE quantization
SAGE quantization runs both attention matrix multiplications in low precision: TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
Q and K are quantized to INT8 or FP8 E4M3 for `QK^T`, and P and V use FP8 E4M3 matrix multiplications use low precision.
for `PV`. P is quantized inside the FMHA kernel; V is quantized per channel
before the kernel call.
vLLM-Omni exposes the Q/K dtype and the Q/K scale granularity. The P and V
formats are fixed by the kernel, so the `quant` key has no V dtype for this
backend. This mode is distinct from the standalone
[SageAttention backends](sage.md), which use their own kernels.
FP8 Q/K kernels exist on `sm100a` and `sm103a`; INT8 Q/K kernels exist on | Key | Valid values | Meaning |
`sm100a` only.
| Key | Values | Meaning |
| --- | --- | --- | | --- | --- | --- |
| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype. Setting it enables SAGE. | | `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense |
| `q_block_size` | `1`, `4`, `16` | Consecutive query tokens sharing one Q scale; default `1` | | `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` |
| `k_block_size` | `1`, `4`, `16` | Consecutive key tokens sharing one K scale; default `16` | | `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` |
Start with the defaults, `q_block_size=1` and `k_block_size=16`, and try V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
smaller K blocks only if output quality does not meet expectations; smaller Only block sizes with compiled kernels are accepted.
blocks give finer scales at some cost in speed. When a KV sequence in a call is
shorter than `k_block_size`, that call falls back to dense attention and a
warning is logged once.
```bash ```bash
vllm serve <model> --omni \ vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
--diffusion-attention-config \ --diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
'{"default":{"backend":"TRTLLM_ATTN","quant":{ "quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
``` ```
The `quant` key is shared with `FLASHINFER_ATTN`, but each backend validates ```python
its own fields: `float16`/`bfloat16` Q/K dtypes and `dtype_vo` are from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
`FLASHINFER_ATTN` options and are rejected here.
## Composing both modes config = AttentionConfig(
default=AttentionSpec(
`skip_softmax` and `quant` may appear in the same `AttentionSpec`. Their backend="TRTLLM_ATTN",
quality effects compound, so establish a dense baseline, enable one mode at a quant=AttnQuantSpec(
time, and then evaluate the combination on the same prompts and seeds. dtype_qk="fp8_e4m3",
q_block_size=1,
Modes configured in `default` apply to every attention role that has no k_block_size=16,
`per_role` entry. A `per_role` spec replaces the whole spec for that role and ),
does not inherit `quant` or `skip_softmax` from `default`, so ),
`{"backend":"TRTLLM_ATTN"}` is the way to keep a short or sensitive attention )
site dense while the long DiT sequence uses both modes:
```bash
vllm serve MiniMaxAI/MiniMax-H3 --omni \
--diffusion-attention-config '{
"default": {
"backend": "TRTLLM_ATTN",
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16},
"skip_softmax": {"threshold": 0.05, "disabled_until_timestep": 0.97}
},
"per_role": {
"minimax_h3.token_refiner": {"backend": "TRTLLM_ATTN"}
}
}'
``` ```
Role names are declared by each model; the This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are
[attention backend overview](../attention_backends.md#configuration) covers available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only.
the resolution order and the equivalent Python API. The shared `AttnQuantSpec` is backend typed: values intended for
`FLASHINFER_ATTN` are rejected here and vice versa.
End-to-end speedup depends on the share of step time spent in attention, the
sequence length, the chosen Q/K precision, and the tile sparsity the input
actually yields. Benchmark the exact workload rather than extrapolating from
another model.
@@ -94,7 +94,7 @@ In `DiffusionParallelConfig`
!!! info !!! info
Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** and **Boogu-Image-Edit double guidance** also support `cfg_parallel_size=3` for full three-way branch execution. Boogu double guidance additionally supports size 2 via round-robin assignment of its three branches. Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** is an exception: it supports `cfg_parallel_size=3`, which adds a third branch on rank 2 for full three-way CFG parallelism.
--- ---
@@ -120,18 +120,10 @@ Additional requirements:
| Parameter | Type | Default | Description | | Parameter | Type | Default | Description |
|-----------|------|---------|-------------| |-----------|------|---------|-------------|
| `vae_use_tiling` | bool | False | Must be set to `True` when using VAE patch parallelism. | | `vae_use_tiling` | bool | False | Must be set to `True` when using VAE patch parallelism. |
| `vae_stack_tiling` | `"auto"` \| `"true"` \| `"false"` | `"false"` | For VAEs with native support (currently MiniMax H3), batch local decoder tiles into fewer model calls. `"auto"` enables it when the request has at least two tiles per VAE-parallel rank. |
!!! note "Automatic VAE Tiling" !!! note "Automatic VAE Tiling"
When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), the system automatically sets `vae_use_tiling=True` if not already enabled. When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), the system automatically sets `vae_use_tiling=True` if not already enabled.
!!! tip "MiniMax H3 stacked tiling"
MiniMax H3 supports lossless native stacked tiling. Enable it with
`--vae-stack-tiling auto` (or `true`) to batch each rank's spatial tiles
into one decoder call per temporal window. `auto` keeps the sequential
path for shapes that do not provide multiple tiles per rank, while
`false` remains the default to avoid increasing the VAE activation peak.
--- ---
## Spatially-Sharded Decode (Wan) ## Spatially-Sharded Decode (Wan)
+2 -3
View File
@@ -117,7 +117,6 @@ The following tables show which models support each feature:
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:| |--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **Boogu-Image** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
@@ -139,7 +138,7 @@ The following tables show which models support each feature:
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **SenseNova-U1 / U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **SenseNova-U1** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
@@ -163,7 +162,7 @@ The following tables show which models support each feature:
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ | | **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
| **SANA-Video-2B T2V I2V** | ❌ | ❌ | ❌ | ✅ | ✅ (TP=2 only) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
@@ -2,6 +2,7 @@
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/text_to_speech>. Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/text_to_speech>.
vLLM-Omni exposes TTS models through the OpenAI-compatible vLLM-Omni exposes TTS models through the OpenAI-compatible
[`POST /v1/audio/speech`](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md) endpoint, [`POST /v1/audio/speech`](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md) endpoint,
launched with `vllm serve <model> --omni`. Each TTS model has its own launched with `vllm serve <model> --omni`. Each TTS model has its own
@@ -15,7 +16,7 @@ For the full list of supported architectures across all modalities, see
## Supported Models ## Supported Models
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
| --- | --- | --- | --- | --- | --- | |---|---|---|---|---|---|
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — | | higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
@@ -103,7 +104,6 @@ For full request-shape documentation (all parameters, response formats, error co
2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets. 2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets.
### Prerequisites ### Prerequisites
```bash ```bash
huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512 huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
``` ```
@@ -111,7 +111,6 @@ huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class). If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class).
### Launch ### Launch
```bash ```bash
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
# or: # or:
@@ -121,16 +120,13 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path. Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
### CLI client ### CLI client
The client defaults to the official upstream zero-shot prompt, so it runs without extra flags: The client defaults to the official upstream zero-shot prompt, so it runs without extra flags:
```bash ```bash
cd examples/online_serving/text_to_speech/cosyvoice3 cd examples/online_serving/text_to_speech/cosyvoice3
python speech_client.py --text "收到好友从远方寄来的生日礼物。" python speech_client.py --text "收到好友从远方寄来的生日礼物。"
``` ```
Pass your own reference clip and transcript for a different voice: Pass your own reference clip and transcript for a different voice:
```bash ```bash
python speech_client.py --text "Hello, this is a cloned voice." \ python speech_client.py --text "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \ --ref-audio /path/to/reference.wav \
@@ -138,7 +134,6 @@ python speech_client.py --text "Hello, this is a cloned voice." \
``` ```
Stream PCM instead of WAV: Stream PCM instead of WAV:
```bash ```bash
python speech_client.py --text "Hello world" --stream --output output.pcm python speech_client.py --text "Hello world" --stream --output output.pcm
``` ```
@@ -146,7 +141,6 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`. The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`.
### Notes ### Notes
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform. - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override. - Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
- For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3). - For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
@@ -158,23 +152,19 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec. 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
### Prerequisites ### Prerequisites
```bash ```bash
pip install fish-speech pip install fish-speech
``` ```
### Launch ### Launch
```bash ```bash
vllm serve fishaudio/s2-pro --omni --port 8091 vllm serve fishaudio/s2-pro --omni --port 8091
# or: # or:
./fish_speech/run_server.sh ./fish_speech/run_server.sh
``` ```
The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`). The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`).
### Voice cloning ### Voice cloning
```bash ```bash
curl -X POST http://localhost:8091/v1/audio/speech \ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \ -H "Content-Type: application/json" \
@@ -187,7 +177,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
``` ```
### CLI client ### CLI client
```bash ```bash
cd examples/online_serving/text_to_speech/fish_speech cd examples/online_serving/text_to_speech/fish_speech
python speech_client.py --text "Hello, how are you?" python speech_client.py --text "Hello, how are you?"
@@ -195,14 +184,12 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
``` ```
### Gradio demo ### Gradio demo
```bash ```bash
./fish_speech/run_gradio_demo.sh # launches server + Gradio ./fish_speech/run_gradio_demo.sh # launches server + Gradio
python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server already running python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server already running
``` ```
### Notes ### Notes
- Output: 44.1 kHz mono. - Output: 44.1 kHz mono.
- Streaming PCM player command must use `-r 44100`. - Streaming PCM player command must use `-r 44100`.
@@ -221,15 +208,12 @@ pip install -U "transformers>=5.3.0"
``` ```
### Launch ### Launch
```bash ```bash
GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
``` ```
Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`. Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
### Sending requests ### Sending requests
```bash ```bash
# Plain TTS # Plain TTS
python higgs_audio_v2/batch_speech_client.py \ python higgs_audio_v2/batch_speech_client.py \
@@ -247,17 +231,14 @@ python higgs_audio_v2/batch_speech_client.py \
``` ```
### Notes ### Notes
- Output: 24 kHz mono. - Output: 24 kHz mono.
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality. - `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`. - Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
## GLM-TTS ## GLM-TTS
2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`. 2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
### Launch ### Launch
```bash ```bash
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
# or: # or:
@@ -265,7 +246,6 @@ bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-T
``` ```
### Sending requests ### Sending requests
```bash ```bash
# Voice cloning (required) # Voice cloning (required)
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
@@ -282,13 +262,11 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
``` ```
### Gradio demo ### Gradio demo
```bash ```bash
bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
``` ```
### Notes ### Notes
- Output: 24 kHz mono WAV via HiFT vocoder. - Output: 24 kHz mono WAV via HiFT vocoder.
- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds. - `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer. - Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
@@ -300,15 +278,12 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline. Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
### Prerequisites ### Prerequisites
```bash ```bash
huggingface-cli download k2-fsa/OmniVoice huggingface-cli download k2-fsa/OmniVoice
``` ```
Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`. Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`.
### Launch ### Launch
```bash ```bash
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
# or: # or:
@@ -316,7 +291,6 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
``` ```
### CLI client ### CLI client
```bash ```bash
cd examples/online_serving/text_to_speech/omnivoice cd examples/online_serving/text_to_speech/omnivoice
python speech_client.py --text "Hello, how are you?" python speech_client.py --text "Hello, how are you?"
@@ -326,7 +300,6 @@ python speech_client.py --text "Bonjour, comment allez-vous?" --language French
The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`. The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`.
### Notes ### Notes
- Voice cloning and voice design require offline inference; see the [offline OmniVoice section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#omnivoice). - Voice cloning and voice design require offline inference; see the [offline OmniVoice section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#omnivoice).
--- ---
@@ -336,7 +309,7 @@ The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--l
Three model variants exposed via separate checkpoints: Three model variants exposed via separate checkpoints:
| Variant | HF repo | Use | | Variant | HF repo | Use |
| --- | --- | --- | |---|---|---|
| CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions | | CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions |
| VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description | | VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description |
| Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio | | Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio |
@@ -344,7 +317,6 @@ Three model variants exposed via separate checkpoints:
Each variant ships smaller `0.6B` companions where available. Each variant ships smaller `0.6B` companions where available.
### Launch ### Launch
```bash ```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
# or: # or:
@@ -354,30 +326,27 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
``` ```
### Executor backend ### Executor backend
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override. Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor. To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
### Tuning stage 1 `max_num_seqs` per task type ### Tuning stage 1 `max_num_seqs` per task type
The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost. The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost.
CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types: CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types:
```bash ```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
--stage-overrides '{"1": {"max_num_seqs": 1}}' --stage-overrides '{"1": {"max_num_seqs": 1}}'
``` ```
### Sending requests ### Sending requests
```bash ```bash
# CustomVoice with a predefined speaker # CustomVoice with a predefined speaker
python qwen3_tts/openai_speech_client.py \ python qwen3_tts/openai_speech_client.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--text "今天天气真好" \ --text "今天天气真好" \
--speaker ryan \ --voice ryan \
--instructions "用开心的语气说" --instructions "用开心的语气说"
# VoiceDesign with a style description # VoiceDesign with a style description
@@ -397,9 +366,7 @@ python qwen3_tts/openai_speech_client.py \
``` ```
### Voices endpoint ### Voices endpoint
List available voices, or upload a custom one for Base cloning: List available voices, or upload a custom one for Base cloning:
```bash ```bash
# List # List
curl http://localhost:8091/v1/audio/voices curl http://localhost:8091/v1/audio/voices
@@ -412,18 +379,10 @@ curl -X POST http://localhost:8091/v1/audio/voices \
-F "ref_text=The exact transcript of the audio sample." \ -F "ref_text=The exact transcript of the audio sample." \
-F "speaker_description=warm narrator" -F "speaker_description=warm narrator"
``` ```
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
For Qwen3-TTS, uploaded voices are Base voice-cloning inputs and require a Base
checkpoint. When a request names an uploaded voice, the server infers
`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
CustomVoice speakers and require a CustomVoice checkpoint.
### Precomputed custom voices ### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
For reused Base voice-cloning speakers, precompute the reference artifacts once
and load them at server startup. Precomputed voices use the same Base task and
checkpoint-matching rules as uploaded voices:
```bash ```bash
python qwen3_tts/precompute_custom_voice.py \ python qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \ --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
@@ -433,15 +392,11 @@ python qwen3_tts/precompute_custom_voice.py \
--mode icl \ --mode icl \
--output-dir /path/to/custom_voices --output-dir /path/to/custom_voices
``` ```
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config: `--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
```yaml ```yaml
custom_voice_dir: /path/to/custom_voices custom_voice_dir: /path/to/custom_voices
``` ```
Then start the server with that config and call the Speech API with only the voice name: Then start the server with that config and call the Speech API with only the voice name:
```bash ```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
@@ -452,7 +407,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
``` ```
### Streaming PCM ### Streaming PCM
```bash ```bash
curl -X POST http://localhost:8091/v1/audio/speech \ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \ -H "Content-Type: application/json" \
@@ -465,20 +419,16 @@ curl -X POST http://localhost:8091/v1/audio/speech \
"response_format": "pcm" "response_format": "pcm"
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
``` ```
Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
### Streaming WebSocket ### Streaming WebSocket
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
```bash ```bash
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine." python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1 python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
``` ```
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection: `input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
```bash ```bash
python qwen3_tts/streaming_speech_client.py \ python qwen3_tts/streaming_speech_client.py \
--text "First utterance." \ --text "First utterance." \
@@ -486,7 +436,6 @@ python qwen3_tts/streaming_speech_client.py \
``` ```
### Gradio demos ### Gradio demos
```bash ```bash
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default) ./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign ./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign
@@ -498,15 +447,12 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
``` ```
### Speaker embedding interpolation ### Speaker embedding interpolation
`qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage. `qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage.
### Batch client ### Batch client
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement. `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes ### Notes
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above. - Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request. - With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`. - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
@@ -521,15 +467,12 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
Single-stage native AR TTS at 48 kHz. Single-stage native AR TTS at 48 kHz.
### Launch ### Launch
```bash ```bash
vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000 vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000
``` ```
Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks. Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks.
### Sending requests ### Sending requests
```bash ```bash
# Zero-shot synthesis # Zero-shot synthesis
python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2." python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2."
@@ -539,13 +482,10 @@ python voxcpm2/openai_speech_client.py \
--text "This should sound like the reference speaker." \ --text "This should sound like the reference speaker." \
--ref-audio /path/to/reference.wav --ref-audio /path/to/reference.wav
``` ```
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs. The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
### Precomputed custom voices ### Precomputed custom voices
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`: For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
```bash ```bash
python voxcpm2/precompute_custom_voice.py \ python voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \ --model openbmb/VoxCPM2 \
@@ -555,21 +495,16 @@ python voxcpm2/precompute_custom_voice.py \
--prompt-text "Original transcript of the reference audio" \ --prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices --output-dir /path/to/custom_voices
``` ```
Add the output directory to the deploy config: Add the output directory to the deploy config:
```yaml ```yaml
custom_voice_dir: /path/to/custom_voices custom_voice_dir: /path/to/custom_voices
``` ```
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`. After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
### Gradio demo (gapless streaming via AudioWorklet) ### Gradio demo (gapless streaming via AudioWorklet)
```bash ```bash
python voxcpm2/gradio_demo.py python voxcpm2/gradio_demo.py
``` ```
Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`. Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`.
--- ---
@@ -579,31 +514,24 @@ Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free p
Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific. Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific.
### Prerequisites ### Prerequisites
Latest `mistral_common` with `SpeechRequest` support: Latest `mistral_common` with `SpeechRequest` support:
```bash ```bash
pip install -e /path/to/mistral-common # or upgrade from PyPI when available pip install -e /path/to/mistral-common # or upgrade from PyPI when available
``` ```
### Launch ### Launch
```bash ```bash
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091 vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
``` ```
Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`. Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`.
### Gradio demo ### Gradio demo
```bash ```bash
python voxtral_tts/gradio_demo.py python voxtral_tts/gradio_demo.py
``` ```
The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients). The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients).
### Notes ### Notes
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
- Voice cloning is gated upstream and may require a recent `mistral_common`. - Voice cloning is gated upstream and may require a recent `mistral_common`.
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now. - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
@@ -46,7 +46,7 @@ More failure modes (for example, network jitter and network interruption) will b
## Fault Injection Scenario Matrix ## Fault Injection Scenario Matrix
| Scenario | Fault Type | System Behavior | Current Status | | Scenario | Fault Type | System Behavior | Current Status |
| ------ | ---------------------------- | ---------- | ---------- | |------|----------------------------|----------|----------|
| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | | | No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | |
| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | | | No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | |
| No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725) <br>[#43060](https://github.com/vllm-project/vllm/issues/43060) | | No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725) <br>[#43060](https://github.com/vllm-project/vllm/issues/43060) |
@@ -1,259 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Run Mage-VL offline end-to-end inference through the upstream Mage scripts.
Mage-VL currently ships its custom Mage-ViT and codec processors as
Transformers remote code in the Microsoft Mage repository. This example keeps
the vLLM-Omni entrypoint stable while delegating the checkpoint-specific forward
path to that upstream implementation.
"""
from __future__ import annotations
import argparse
import json
import os
import subprocess
import sys
import time
from importlib.util import find_spec
from pathlib import Path
DEFAULT_MODEL = "microsoft/Mage-VL"
DEFAULT_IMAGE_QUESTION = "Describe this image in detail."
DEFAULT_VIDEO_QUESTION = "Describe this video."
def _default_mage_repo() -> Path | None:
value = os.getenv("MAGE_REPO_DIR")
if value:
return Path(value)
sibling = Path(__file__).resolve().parents[3].parent / "Mage"
if sibling.exists():
return sibling
return None
def _resolve_mage_repo(value: str | None) -> Path:
repo = Path(value) if value else _default_mage_repo()
if repo is None:
raise ValueError("Set --mage-repo or MAGE_REPO_DIR to a checkout of https://github.com/microsoft/Mage.")
repo = repo.expanduser().resolve()
if not (repo / "mage_vl" / "inference_base.py").is_file():
raise FileNotFoundError(f"Missing Mage-VL scripts under {repo}. Expected mage_vl/inference_base.py.")
return repo
def _default_media_path(mage_repo: Path, task: str) -> Path:
if task == "image":
return mage_repo / "mage_vl" / "assets" / "examples" / "dog.jpg"
return mage_repo / "mage_vl" / "assets" / "examples" / "soccer-broadcast.mp4"
def _resolve_media_path(path: str | None, mage_repo: Path, task: str) -> Path:
resolved = Path(path).expanduser().resolve() if path else _default_media_path(mage_repo, task)
if not resolved.is_file():
raise FileNotFoundError(f"Input media file not found: {resolved}")
return resolved
def build_command(args: argparse.Namespace, mage_repo: Path, media_path: Path) -> list[str]:
if args.task == "streaming":
return [
sys.executable,
str(mage_repo / "mage_vl" / "inference_streaming.py"),
"--video",
str(media_path),
"--checkpoint",
args.model,
"--video_backend",
args.video_backend,
"--num_frames",
str(args.num_frames),
"--cur_fps",
str(args.cur_fps),
"--segment_sec",
str(args.segment_sec),
"--max_new_tokens",
str(args.max_new_tokens),
"--max_segments",
str(args.max_segments),
"--gate_threshold",
str(args.gate_threshold),
"--device",
args.device,
"--attn_impl",
args.attn_impl,
]
cmd = [
sys.executable,
str(mage_repo / "mage_vl" / "inference_base.py"),
"--mode",
"offline",
"--question",
args.question or (DEFAULT_IMAGE_QUESTION if args.task == "image" else DEFAULT_VIDEO_QUESTION),
"--model",
args.model,
"--max-new-tokens",
str(args.max_new_tokens),
]
if args.task == "image":
cmd.extend(["--image", str(media_path)])
return cmd
cmd.extend(
[
"--video",
str(media_path),
"--video-backend",
args.video_backend,
"--num-frames",
str(args.num_frames),
"--max-pixels",
str(args.max_pixels),
]
)
if args.video_backend == "codec":
cmd.extend(["--codec-engine", args.codec_engine])
return cmd
def check_runtime_dependencies(task: str) -> None:
required = ["torch"]
if task == "streaming":
required.append("mamba_ssm")
missing = [name for name in required if find_spec(name) is None]
if not missing:
return
install_hint = (
"Install Mage-VL runtime dependencies in this environment first:\n"
" pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n"
" grep -v '^mamba-ssm' /path/to/Mage/mage_vl/requirements.txt > /tmp/mage_vl_requirements_no_mamba.txt\n"
" pip install -r /tmp/mage_vl_requirements_no_mamba.txt\n"
" pip install 'causal-conv1d>=1.4.0' --no-build-isolation\n"
" pip install 'mamba-ssm>=2.2' --no-build-isolation"
)
raise RuntimeError(f"Missing Python packages: {', '.join(missing)}.\n{install_hint}")
def run_command(cmd: list[str], cwd: Path, timeout: int) -> subprocess.CompletedProcess[str]:
return subprocess.run(
cmd,
cwd=str(cwd),
input="y\n",
text=True,
capture_output=True,
timeout=timeout,
check=False,
)
def write_outputs(
*,
args: argparse.Namespace,
mage_repo: Path,
media_path: Path,
cmd: list[str],
result: subprocess.CompletedProcess[str] | None,
elapsed_s: float | None,
) -> Path:
output_dir = Path(args.output_dir).expanduser().resolve()
output_dir.mkdir(parents=True, exist_ok=True)
summary = {
"task": args.task,
"model": args.model,
"mage_repo": str(mage_repo),
"media_path": str(media_path),
"command": cmd,
"dry_run": args.dry_run,
"elapsed_s": elapsed_s,
"returncode": None if result is None else result.returncode,
}
if result is not None:
stdout = result.stdout.strip()
stderr = result.stderr.strip()
summary["stdout"] = stdout
summary["stderr"] = stderr
(output_dir / "stdout.txt").write_text(result.stdout, encoding="utf-8")
(output_dir / "stderr.txt").write_text(result.stderr, encoding="utf-8")
summary_path = output_dir / "summary.json"
summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
return summary_path
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Mage-VL offline end-to-end inference example.")
parser.add_argument(
"--mage-repo",
default=None,
help="Path to a Microsoft Mage checkout. Defaults to MAGE_REPO_DIR.",
)
parser.add_argument("--model", default=DEFAULT_MODEL, help="HF model id or local checkpoint path.")
parser.add_argument("--task", choices=("image", "video", "streaming"), default="image")
parser.add_argument("--image-path", default=None, help="Image input. Defaults to the Mage sample image.")
parser.add_argument("--video-path", default=None, help="Video input. Defaults to the Mage sample video.")
parser.add_argument("--question", default=None, help="Question for image/video offline QA.")
parser.add_argument("--video-backend", choices=("frames", "codec"), default="frames")
parser.add_argument("--codec-engine", choices=("traditional", "neural"), default="traditional")
parser.add_argument("--num-frames", type=int, default=32)
parser.add_argument("--max-pixels", type=int, default=150000)
parser.add_argument("--max-new-tokens", type=int, default=256)
parser.add_argument("--cur-fps", type=float, default=2.0)
parser.add_argument("--segment-sec", type=float, default=8.0)
parser.add_argument("--max-segments", type=int, default=0)
parser.add_argument("--gate-threshold", type=float, default=0.5)
parser.add_argument("--device", default="cuda")
parser.add_argument("--attn-impl", choices=("flash_attention_2", "sdpa", "eager"), default="flash_attention_2")
parser.add_argument("--timeout", type=int, default=1800)
parser.add_argument("--output-dir", default="outputs/mage_vl")
parser.add_argument("--dry-run", action="store_true", help="Build the command and summary without executing it.")
args = parser.parse_args()
if args.num_frames <= 0:
parser.error("--num-frames must be positive")
if args.segment_sec <= 0:
parser.error("--segment-sec must be positive")
if args.task == "image" and args.video_path:
parser.error("--video-path cannot be used with --task image")
if args.task in {"video", "streaming"} and args.image_path:
parser.error("--image-path cannot be used with video tasks")
return args
def main() -> None:
args = parse_args()
mage_repo = _resolve_mage_repo(args.mage_repo)
media_arg = args.image_path if args.task == "image" else args.video_path
media_path = _resolve_media_path(media_arg, mage_repo, args.task)
cmd = build_command(args, mage_repo, media_path)
result = None
elapsed_s = None
if not args.dry_run:
check_runtime_dependencies(args.task)
start = time.perf_counter()
result = run_command(cmd, cwd=mage_repo, timeout=args.timeout)
elapsed_s = time.perf_counter() - start
summary_path = write_outputs(
args=args,
mage_repo=mage_repo,
media_path=media_path,
cmd=cmd,
result=result,
elapsed_s=elapsed_s,
)
print(f"summary saved to {summary_path}")
if result is not None:
if result.stdout:
print(result.stdout, end="")
if result.returncode != 0:
if result.stderr:
print(result.stderr, file=sys.stderr, end="")
raise SystemExit(result.returncode)
if __name__ == "__main__":
main()
@@ -1,79 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations
import argparse
import asyncio
import base64
import json
import time
from pathlib import Path
import websockets
async def run(args: argparse.Namespace) -> None:
headers = {"Authorization": f"Bearer {args.auth_token}"} if args.auth_token else None
started = time.perf_counter()
async with websockets.connect(
f"{args.url}?session_id={args.session_id}",
additional_headers=headers,
max_size=args.max_message_bytes,
open_timeout=args.timeout,
) as websocket:
created = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout))
print(json.dumps(created, ensure_ascii=False))
await websocket.send(
json.dumps(
{
"type": "input.append",
"modality": "text",
"data": args.prompt,
}
)
)
encoded = base64.b64encode(Path(args.video).read_bytes()).decode("ascii")
await websocket.send(
json.dumps(
{
"type": "input.append",
"modality": "video",
"data": {
"kind": "video/mp4",
"segment_id": args.segment_id,
"pts_ms": args.pts_ms,
"duration_ms": args.duration_ms,
"data": {"video_base64": encoded},
},
}
)
)
while True:
event = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout))
print(json.dumps(event, ensure_ascii=False))
if event.get("type") in {"response.done", "error"}:
break
await websocket.send(json.dumps({"type": "close"}))
print(f"elapsed_s={time.perf_counter() - started:.3f}")
def main() -> None:
parser = argparse.ArgumentParser(description="Send one Mage-VL video window over the duplex WebSocket.")
parser.add_argument("--url", default="ws://127.0.0.1:8090/v1/mage-vl/duplex")
parser.add_argument("--video", required=True)
parser.add_argument("--prompt", default="Describe this video segment in detail.")
parser.add_argument("--session-id", default="mage-vl-example")
parser.add_argument("--segment-id", default="segment-0")
parser.add_argument("--pts-ms", type=int, default=0)
parser.add_argument("--duration-ms", type=int, default=4000)
parser.add_argument("--auth-token")
parser.add_argument("--timeout", type=float, default=120.0)
parser.add_argument("--max-message-bytes", type=int, default=16 * 1024 * 1024)
args = parser.parse_args()
asyncio.run(run(args))
if __name__ == "__main__":
main()
@@ -13,7 +13,7 @@ For the full list of supported architectures across all modalities, see
## Supported Models ## Supported Models
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
| --- | --- | --- | --- | --- | --- | |---|---|---|---|---|---|
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — | | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
@@ -104,7 +104,6 @@ For full request-shape documentation (all parameters, response formats, error co
2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`. 2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
### Launch ### Launch
```bash ```bash
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091 vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
# or: # or:
@@ -112,7 +111,6 @@ bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-T
``` ```
### Sending requests ### Sending requests
```bash ```bash
# Voice cloning (required) # Voice cloning (required)
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
@@ -129,13 +127,11 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
``` ```
### Gradio demo ### Gradio demo
```bash ```bash
bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
``` ```
### Notes ### Notes
- Output: 24 kHz mono WAV via HiFT vocoder. - Output: 24 kHz mono WAV via HiFT vocoder.
- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds. - `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer. - Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
@@ -155,7 +151,6 @@ pip install 'vllm-omni[indextts2]'
``` ```
### Launch ### Launch
```bash ```bash
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092 vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
# or, to pass the bundled deploy config explicitly: # or, to pass the bundled deploy config explicitly:
@@ -169,7 +164,6 @@ bash examples/online_serving/text_to_speech/indextts2/run_server.sh
``` ```
### Sending requests ### Sending requests
```bash ```bash
# Voice cloning (ref_audio required) # Voice cloning (ref_audio required)
python examples/online_serving/text_to_speech/indextts2/speech_client.py \ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
@@ -200,7 +194,6 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
``` ```
### Notes ### Notes
- Output: 22.05 kHz mono WAV. - Output: 22.05 kHz mono WAV.
- Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; neither version provides a built-in text-only preset voice. - Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; neither version provides a built-in text-only preset voice.
- IndexTTS-2.5 request controls `lang` and `text_normalization` are carried in `extra_params`; the client does this for `--model-version 2.5`. - IndexTTS-2.5 request controls `lang` and `text_normalization` are carried in `extra_params`; the client does this for `--model-version 2.5`.
@@ -219,7 +212,6 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec. 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
### Prerequisites ### Prerequisites
```bash ```bash
pip install fish-speech pip install fish-speech
``` ```
@@ -245,17 +237,14 @@ export VLLM_OMNI_FISH_KVCACHE_ATTN=0
``` ```
### Launch ### Launch
```bash ```bash
vllm serve fishaudio/s2-pro --omni --port 8091 vllm serve fishaudio/s2-pro --omni --port 8091
# or: # or:
./fish_speech/run_server.sh ./fish_speech/run_server.sh
``` ```
The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`). The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`).
### Voice cloning ### Voice cloning
```bash ```bash
curl -X POST http://localhost:8091/v1/audio/speech \ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \ -H "Content-Type: application/json" \
@@ -268,7 +257,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
``` ```
### CLI client ### CLI client
```bash ```bash
cd examples/online_serving/text_to_speech/fish_speech cd examples/online_serving/text_to_speech/fish_speech
python speech_client.py --text "Hello, how are you?" python speech_client.py --text "Hello, how are you?"
@@ -276,14 +264,12 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
``` ```
### Gradio demo ### Gradio demo
```bash ```bash
./fish_speech/run_gradio_demo.sh # launches server + Gradio ./fish_speech/run_gradio_demo.sh # launches server + Gradio
python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server already running python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server already running
``` ```
### Notes ### Notes
- Output: 44.1 kHz mono. - Output: 44.1 kHz mono.
- Streaming PCM player command must use `-r 44100`. - Streaming PCM player command must use `-r 44100`.
@@ -294,13 +280,10 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`. Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`.
### Launch ### Launch
```bash ```bash
bash examples/online_serving/text_to_speech/ming_tts/run_server.sh bash examples/online_serving/text_to_speech/ming_tts/run_server.sh
``` ```
Equivalent manual command: Equivalent manual command:
```bash ```bash
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \ vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
--deploy-config vllm_omni/deploy/ming_tts.yaml \ --deploy-config vllm_omni/deploy/ming_tts.yaml \
@@ -309,14 +292,12 @@ vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
``` ```
### Sending requests ### Sending requests
```bash ```bash
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
--text "你好,这是 Ming 在线语音合成测试。" --text "你好,这是 Ming 在线语音合成测试。"
``` ```
Structured dialect control: Structured dialect control:
```bash ```bash
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
--text "我觉得社会企业同个人都有责任" \ --text "我觉得社会企业同个人都有责任" \
@@ -325,7 +306,6 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
``` ```
Zero-shot cloning: Zero-shot cloning:
```bash ```bash
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \ --text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
@@ -334,7 +314,6 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
``` ```
### Notes ### Notes
- `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook. - `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook.
- Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples. - Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples.
- Full request details live in [`ming_tts/README.md`](ming_tts/README.md). - Full request details live in [`ming_tts/README.md`](ming_tts/README.md).
@@ -346,14 +325,11 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`. Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`.
### Launch ### Launch
```bash ```bash
# from repo root # from repo root
bash examples/online_serving/text_to_speech/ming_flash_omni_tts/run_server.sh bash examples/online_serving/text_to_speech/ming_flash_omni_tts/run_server.sh
``` ```
Equivalent manual command: Equivalent manual command:
```bash ```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 \ vllm serve Jonathan1909/Ming-flash-omni-2.0 \
--deploy-config vllm_omni/deploy/ming_flash_omni_tts.yaml \ --deploy-config vllm_omni/deploy/ming_flash_omni_tts.yaml \
@@ -362,7 +338,6 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 \
``` ```
### Sending requests ### Sending requests
```bash ```bash
python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.py \ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.py \
--text "我们当迎着阳光辛勤耕作,去摘取,去制作,去品尝,去馈赠。" \ --text "我们当迎着阳光辛勤耕作,去摘取,去制作,去品尝,去馈赠。" \
@@ -370,7 +345,6 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
``` ```
ASMR-style caption via `instructions`: ASMR-style caption via `instructions`:
```bash ```bash
python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.py \ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.py \
--text "我会一直在这里陪着你,直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗?" \ --text "我会一直在这里陪着你,直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗?" \
@@ -379,7 +353,6 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
``` ```
### Notes ### Notes
- Server uses `use_zero_spk_emb=True` and the cookbook decode defaults (`max_decode_steps=200`, `cfg=2.0`, `sigma=0.25`, `temperature=0.0`). For other caption fields (`语速`, `基频`, `IP`, BGM, etc.) or overriding decode args, use the offline example where `additional_information` is set explicitly. - Server uses `use_zero_spk_emb=True` and the cookbook decode defaults (`max_decode_steps=200`, `cfg=2.0`, `sigma=0.25`, `temperature=0.0`). For other caption fields (`语速`, `基频`, `IP`, BGM, etc.) or overriding decode args, use the offline example where `additional_information` is set explicitly.
- This is the online counterpart of [`examples/offline_inference/text_to_speech/ming_flash_omni_tts/`](../../offline_inference/text_to_speech/ming_flash_omni_tts/). - This is the online counterpart of [`examples/offline_inference/text_to_speech/ming_flash_omni_tts/`](../../offline_inference/text_to_speech/ming_flash_omni_tts/).
- For multimodal Ming-flash-omni online serving, see [`examples/online_serving/ming_flash_omni/`](../../ming_flash_omni/). - For multimodal Ming-flash-omni online serving, see [`examples/online_serving/ming_flash_omni/`](../../ming_flash_omni/).
@@ -393,17 +366,14 @@ Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every
> The OpenAI-schema `voice` and `ref_text` fields are accepted but ignored — `voice_clone` does not consume a transcript, and upstream's `continuation` mode (the only path that accepts `prompt_text`) emits near-silent output, so it is not exposed here. Sample reference clips ship in the upstream repo under [`assets/audio/`](https://github.com/OpenMOSS/MOSS-TTS-Nano/tree/main/assets/audio). > The OpenAI-schema `voice` and `ref_text` fields are accepted but ignored — `voice_clone` does not consume a transcript, and upstream's `continuation` mode (the only path that accepts `prompt_text`) emits near-silent output, so it is not exposed here. Sample reference clips ship in the upstream repo under [`assets/audio/`](https://github.com/OpenMOSS/MOSS-TTS-Nano/tree/main/assets/audio).
### Launch ### Launch
```bash ```bash
vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091 vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
# or: # or:
./moss_tts_nano/run_server.sh ./moss_tts_nano/run_server.sh
``` ```
The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--deploy-config`, `--trust-remote-code`, or `--enforce-eager` flags are needed. The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--deploy-config`, `--trust-remote-code`, or `--enforce-eager` flags are needed.
### Sending requests ### Sending requests
```bash ```bash
# One-off fetch of a sample reference clip; cache under XDG_CACHE_HOME. # One-off fetch of a sample reference clip; cache under XDG_CACHE_HOME.
REF_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/moss-tts-nano" REF_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/moss-tts-nano"
@@ -422,7 +392,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
``` ```
### Streaming PCM ### Streaming PCM
```bash ```bash
curl -X POST http://localhost:8091/v1/audio/speech \ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \ -H "Content-Type: application/json" \
@@ -436,7 +405,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
``` ```
### Gradio demo ### Gradio demo
```bash ```bash
# Option 1: launch server + Gradio together # Option 1: launch server + Gradio together
./moss_tts_nano/run_gradio_demo.sh ./moss_tts_nano/run_gradio_demo.sh
@@ -444,11 +412,9 @@ curl -X POST http://localhost:8091/v1/audio/speech \
# Option 2: server already running # Option 2: server already running
python moss_tts_nano/gradio_demo.py --api-base http://localhost:8091 python moss_tts_nano/gradio_demo.py --api-base http://localhost:8091
``` ```
Then open http://localhost:7860 in your browser.
Then open <http://localhost:7860> in your browser.
### Notes ### Notes
- Output is 48 kHz mono PCM (the upstream tokenizer is internally stereo at 48 kHz; the wrapper averages to mono before reaching the engine). - Output is 48 kHz mono PCM (the upstream tokenizer is internally stereo at 48 kHz; the wrapper averages to mono before reaching the engine).
- Standard `/v1/audio/speech` request shape: `input`, `ref_audio` (base64 data URL), `response_format`, `stream`, `max_new_tokens`. The `voice` and `ref_text` fields from the OpenAI schema are accepted but ignored. - Standard `/v1/audio/speech` request shape: `input`, `ref_audio` (base64 data URL), `response_format`, `stream`, `max_new_tokens`. The `voice` and `ref_text` fields from the OpenAI schema are accepted but ignored.
@@ -459,15 +425,12 @@ Then open <http://localhost:7860> in your browser.
Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline. Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
### Prerequisites ### Prerequisites
```bash ```bash
huggingface-cli download k2-fsa/OmniVoice huggingface-cli download k2-fsa/OmniVoice
``` ```
Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`. Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`.
### Launch ### Launch
```bash ```bash
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
# or: # or:
@@ -475,7 +438,6 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
``` ```
### CLI client ### CLI client
```bash ```bash
cd examples/online_serving/text_to_speech/omnivoice cd examples/online_serving/text_to_speech/omnivoice
# Text-only (auto voice) # Text-only (auto voice)
@@ -501,12 +463,13 @@ python speech_client.py --text "Hello, how are you?" --seed 42
The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`. The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`.
## Qwen3-TTS ## Qwen3-TTS
Three model variants exposed via separate checkpoints: Three model variants exposed via separate checkpoints:
| Variant | HF repo | Use | | Variant | HF repo | Use |
| --- | --- | --- | |---|---|---|
| CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions | | CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions |
| VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description | | VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description |
| Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio | | Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio |
@@ -514,7 +477,6 @@ Three model variants exposed via separate checkpoints:
Each variant ships smaller `0.6B` companions where available. Each variant ships smaller `0.6B` companions where available.
### Launch ### Launch
```bash ```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
# or: # or:
@@ -524,13 +486,11 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
``` ```
### Executor backend ### Executor backend
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override. Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor. To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
### Sending requests ### Sending requests
```bash ```bash
# CustomVoice with a predefined speaker # CustomVoice with a predefined speaker
python qwen3_tts/openai_speech_client.py \ python qwen3_tts/openai_speech_client.py \
@@ -556,9 +516,7 @@ python qwen3_tts/openai_speech_client.py \
``` ```
### Voices endpoint ### Voices endpoint
List available voices, or upload a custom one for Base cloning: List available voices, or upload a custom one for Base cloning:
```bash ```bash
# List # List
curl http://localhost:8091/v1/audio/voices curl http://localhost:8091/v1/audio/voices
@@ -571,18 +529,10 @@ curl -X POST http://localhost:8091/v1/audio/voices \
-F "ref_text=The exact transcript of the audio sample." \ -F "ref_text=The exact transcript of the audio sample." \
-F "speaker_description=warm narrator" -F "speaker_description=warm narrator"
``` ```
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
For Qwen3-TTS, uploaded voices are Base voice-cloning inputs and require a Base
checkpoint. When a request names an uploaded voice, the server infers
`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
CustomVoice speakers and require a CustomVoice checkpoint.
### Precomputed custom voices ### Precomputed custom voices
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
For reused Base voice-cloning speakers, precompute the reference artifacts once
and load them at server startup. Precomputed voices use the same Base task and
checkpoint-matching rules as uploaded voices:
```bash ```bash
python qwen3_tts/precompute_custom_voice.py \ python qwen3_tts/precompute_custom_voice.py \
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \ --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
@@ -592,15 +542,11 @@ python qwen3_tts/precompute_custom_voice.py \
--mode icl \ --mode icl \
--output-dir /path/to/custom_voices --output-dir /path/to/custom_voices
``` ```
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config: `--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
```yaml ```yaml
custom_voice_dir: /path/to/custom_voices custom_voice_dir: /path/to/custom_voices
``` ```
Then start the server with that config and call the Speech API with only the voice name: Then start the server with that config and call the Speech API with only the voice name:
```bash ```bash
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
@@ -611,7 +557,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
``` ```
### Streaming PCM ### Streaming PCM
```bash ```bash
curl -X POST http://localhost:8091/v1/audio/speech \ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \ -H "Content-Type: application/json" \
@@ -624,20 +569,16 @@ curl -X POST http://localhost:8091/v1/audio/speech \
"response_format": "pcm" "response_format": "pcm"
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
``` ```
Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
### Streaming WebSocket ### Streaming WebSocket
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
```bash ```bash
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine." python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1 python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
``` ```
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection: `input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
```bash ```bash
python qwen3_tts/streaming_speech_client.py \ python qwen3_tts/streaming_speech_client.py \
--text "First utterance." \ --text "First utterance." \
@@ -645,7 +586,6 @@ python qwen3_tts/streaming_speech_client.py \
``` ```
To receive word-level timestamps, launch the server with a forced aligner: To receive word-level timestamps, launch the server with a forced aligner:
```bash ```bash
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--omni \ --omni \
@@ -653,9 +593,7 @@ vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
--trust-remote-code \ --trust-remote-code \
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
``` ```
Then request PCM JSON sidecar chunks: Then request PCM JSON sidecar chunks:
```bash ```bash
python qwen3_tts/streaming_speech_client.py \ python qwen3_tts/streaming_speech_client.py \
--text "Hello world. How are you?" \ --text "Hello world. How are you?" \
@@ -663,7 +601,6 @@ python qwen3_tts/streaming_speech_client.py \
--response-format pcm \ --response-format pcm \
--word-timestamps --word-timestamps
``` ```
The client writes one PCM file per sentence and a matching The client writes one PCM file per sentence and a matching
`sentence_XXX_timestamps.json` sidecar. `sentence_XXX_timestamps.json` sidecar.
@@ -676,18 +613,15 @@ still returns — use the WebSocket path for long transcripts.
To *see* the alignment instead of reading a JSON sidecar, run the To *see* the alignment instead of reading a JSON sidecar, run the
word-timestamp Gradio demo (server must be launched with `--forced-aligner`): word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
```bash ```bash
python qwen3_tts/word_timestamps_demo.py --api-base http://localhost:8091 python qwen3_tts/word_timestamps_demo.py --api-base http://localhost:8091
``` ```
Each sentence's audio plays in an `<audio>` element while its text is rendered Each sentence's audio plays in an `<audio>` element while its text is rendered
as inline word spans; the current word highlights as `audio.currentTime` as inline word spans; the current word highlights as `audio.currentTime`
crosses each `start_ms`. The **Stop (barge-in)** button cuts playback and crosses each `start_ms`. The **Stop (barge-in)** button cuts playback and
reports the last-spoken word, useful for the voice-agent barge-in case. reports the last-spoken word, useful for the voice-agent barge-in case.
### Gradio demos ### Gradio demos
```bash ```bash
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default) ./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign ./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign
@@ -695,15 +629,12 @@ reports the last-spoken word, useful for the voice-agent barge-in case.
``` ```
### Speaker embedding interpolation ### Speaker embedding interpolation
`qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage. `qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage.
### Batch client ### Batch client
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement. `qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes ### Notes
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above. - Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request. - With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`. - `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
@@ -715,15 +646,12 @@ reports the last-spoken word, useful for the voice-agent barge-in case.
Single-stage native AR TTS at 48 kHz. Single-stage native AR TTS at 48 kHz.
### Launch ### Launch
```bash ```bash
vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000 vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000
``` ```
Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks. Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks.
### Sending requests ### Sending requests
```bash ```bash
# Zero-shot synthesis # Zero-shot synthesis
python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2." python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2."
@@ -733,13 +661,10 @@ python voxcpm2/openai_speech_client.py \
--text "This should sound like the reference speaker." \ --text "This should sound like the reference speaker." \
--ref-audio /path/to/reference.wav --ref-audio /path/to/reference.wav
``` ```
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs. The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
### Precomputed custom voices ### Precomputed custom voices
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`: For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
```bash ```bash
python voxcpm2/precompute_custom_voice.py \ python voxcpm2/precompute_custom_voice.py \
--model openbmb/VoxCPM2 \ --model openbmb/VoxCPM2 \
@@ -749,21 +674,16 @@ python voxcpm2/precompute_custom_voice.py \
--prompt-text "Original transcript of the reference audio" \ --prompt-text "Original transcript of the reference audio" \
--output-dir /path/to/custom_voices --output-dir /path/to/custom_voices
``` ```
Add the output directory to the deploy config: Add the output directory to the deploy config:
```yaml ```yaml
custom_voice_dir: /path/to/custom_voices custom_voice_dir: /path/to/custom_voices
``` ```
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`. After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
### Gradio demo (gapless streaming via AudioWorklet) ### Gradio demo (gapless streaming via AudioWorklet)
```bash ```bash
python voxcpm2/gradio_demo.py python voxcpm2/gradio_demo.py
``` ```
Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`. Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`.
--- ---
@@ -773,31 +693,24 @@ Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free p
Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific. Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific.
### Prerequisites ### Prerequisites
Latest `mistral_common` with `SpeechRequest` support: Latest `mistral_common` with `SpeechRequest` support:
```bash ```bash
pip install -e /path/to/mistral-common # or upgrade from PyPI when available pip install -e /path/to/mistral-common # or upgrade from PyPI when available
``` ```
### Launch ### Launch
```bash ```bash
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091 vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
``` ```
Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`. Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`.
### Gradio demo ### Gradio demo
```bash ```bash
python voxtral_tts/gradio_demo.py python voxtral_tts/gradio_demo.py
``` ```
The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients). The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients).
### Notes ### Notes
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
- Voice cloning is gated upstream and may require a recent `mistral_common`. - Voice cloning is gated upstream and may require a recent `mistral_common`.
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now. - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
@@ -1,6 +1,3 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""OpenAI-compatible client for Qwen3-TTS via /v1/audio/speech endpoint. """OpenAI-compatible client for Qwen3-TTS via /v1/audio/speech endpoint.
This script demonstrates how to use the OpenAI-compatible speech API This script demonstrates how to use the OpenAI-compatible speech API
@@ -77,8 +74,6 @@ def run_tts_generation(args) -> None:
"voice": args.speaker, "voice": args.speaker,
"response_format": args.response_format, "response_format": args.response_format,
} }
if args.sample_rate is not None:
payload["sample_rate"] = args.sample_rate
# Add optional parameters # Add optional parameters
if args.instructions: if args.instructions:
@@ -251,13 +246,6 @@ def parse_args():
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio output format (default: wav)", help="Audio output format (default: wav)",
) )
parser.add_argument(
"--sample-rate",
type=int,
default=None,
choices=[8000, 24000],
help="Output sample rate in Hz (default: model native)",
)
parser.add_argument( parser.add_argument(
"--output", "--output",
"-o", "-o",
+7 -87
View File
@@ -39,10 +39,9 @@ rendering. The Edit and Edit-Turbo checkpoints use the same native
## Hardware Support ## Hardware Support
This recipe documents supported configurations for CUDA GPU serving. The native This recipe documents tested configurations for CUDA GPU serving. The native
pipeline supports single-GPU inference and classifier-free-guidance (CFG) pipeline runs single-GPU; multi-GPU parallelism, CPU offload, and cache
parallelism. CPU offload, cache acceleration, and other multi-GPU dimensions acceleration are not yet supported for this model (see Notes).
remain unsupported (see Notes).
## GPU ## GPU
@@ -109,24 +108,8 @@ curl -s http://localhost:8091/v1/chat/completions \
The model's maximum native resolution is 2K. The model's maximum native resolution is 2K.
- **Known limitations (not yet supported):** CPU offload - **Known limitations (not yet supported):** CPU offload
(`--enable-cpu-offload` / `--enable-layerwise-offload`), Cache-DiT (`--enable-cpu-offload` / `--enable-layerwise-offload`), Cache-DiT
(`--cache-backend cache_dit`), and TP / SP / HSDP multi-GPU parallelism. (`--cache-backend cache_dit`), and multi-GPU parallelism (TP / SP / CFG /
HSDP) are planned follow-ups and are not validated for this model yet.
### 2 x H100 (CFG parallel, Base T2I)
Base T2I has two predictions per guided denoising step, so two CFG ranks map
one-to-one to the positive and negative-text branches:
```bash
vllm serve Boogu/Boogu-Image-0.1-Base \
--omni \
--cfg-parallel-size 2 \
--port 8091
```
Use `guidance_scale > 1.0` to activate the parallel two-branch path. Requests
with `guidance_scale=1.0` remain valid on the same server: every rank evaluates
only the positive branch, no negative embeddings are built, and no guidance is
applied.
## Image editing (Boogu-Image-0.1-Edit) ## Image editing (Boogu-Image-0.1-Edit)
@@ -141,33 +124,6 @@ is sent).
vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091 vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
``` ```
For text-only or image-only guidance, use two CFG ranks:
```bash
vllm serve Boogu/Boogu-Image-0.1-Edit \
--omni \
--cfg-parallel-size 2 \
--port 8091
```
Double guidance has three predictions per denoising step. It supports either
two ranks (round-robin assignment: rank 0 runs branches 0 and 2) or three ranks
(one branch per rank):
```bash
# Practical two-GPU configuration
vllm serve Boogu/Boogu-Image-0.1-Edit \
--omni \
--cfg-parallel-size 2 \
--port 8091
# Maximum branch parallelism
vllm serve Boogu/Boogu-Image-0.1-Edit \
--omni \
--cfg-parallel-size 3 \
--port 8091
```
### Verification ### Verification
Edit an image with `/v1/images/edits` (the model-card example — change a photo Edit an image with `/v1/images/edits` (the model-card example — change a photo
@@ -219,29 +175,11 @@ curl -s http://localhost:8091/v1/chat/completions \
default `1.0` = off). Setting it `> 1.0` enables the double-guidance path default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
(3 model predictions per step), steering more strongly toward the reference (3 model predictions per step), steering more strongly toward the reference
image. image.
- **CFG-size recommendations:**
| Request mode | Scales | Predictions per step | Recommended CFG size |
| --- | --- | ---: | ---: |
| Base T2I text guidance | `guidance_scale > 1`, no reference | 2 | 2 |
| Edit text-only guidance | `guidance_scale > 1`, `guidance_scale_2 = 1` | 2 | 2 |
| Edit image-only guidance | `guidance_scale = 1`, `guidance_scale_2 > 1` | 2 | 2 |
| Edit double guidance | both scales `> 1` | 3 | 2 (practical) or 3 (maximum parallelism) |
| CFG off | both scales `= 1` | 1 | Reuse a size-2 server safely; no speedup expected |
Double guidance preserves Boogu's original combination:
```text
pred = positive_with_reference
+ (text_scale - 1) * (positive_with_reference - negative_with_reference)
+ (image_scale - 1) * (negative_with_reference - uncond)
```
- **Output resolution:** the output size follows the reference image (upstream - **Output resolution:** the output size follows the reference image (upstream
`align_res`, on by default for a single reference), so `height`/`width` are `align_res`, on by default for a single reference), so `height`/`width` are
derived from the input and requested sizes are not applied for edits. derived from the input and requested sizes are not applied for edits.
- **Same limitations** as the Base checkpoint apply except that CFG parallelism - **Same limitations** as the Base checkpoint apply (no CPU offload, Cache-DiT,
is supported as described above; CPU offload, Cache-DiT, TP, SP, and HSDP are or multi-GPU parallelism yet).
still unsupported.
- **Offline editing:** the shared offline example - **Offline editing:** the shared offline example
[`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_to_image.md) [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_to_image.md)
supports Boogu-Image-Edit directly supports Boogu-Image-Edit directly
@@ -313,21 +251,3 @@ curl -s http://localhost:8091/v1/chat/completions \
the regular Edit checkpoint applies. the regular Edit checkpoint applies.
- **Known limitations:** the same single-GPU limitations as Base and Edit apply; - **Known limitations:** the same single-GPU limitations as Base and Edit apply;
CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated. CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
## Performance validation
The checked-in single-device A40 measurements remain the current Boogu
baseline: Base T2I is about 16.37 s request latency at concurrency 1 and Edit
text-only guidance is about 29.64 s for 512x512, 28-step requests, with roughly
36.8 GB peak memory. CFG-parallel A/B numbers must be collected on the same
hardware and software revision before reporting a speedup; do not compare a
multi-GPU result against these A40 numbers as if it were a controlled A/B.
Run the dedicated perf rows in
`tests/dfx/perf/tests/test_boogu_image_vllm_omni.json` and
`tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json` on 2x/3x H100, then
record latency, throughput, and per-GPU peak memory beside the single-device
run. The expected shape is close to the general CFG-parallel range (roughly
1.5x-1.8x for balanced two-branch workloads); the three-branch Edit result will
depend on whether two-rank round-robin or full three-rank execution is used and
must be reported from measurement rather than assumed.
-114
View File
@@ -1,114 +0,0 @@
# Mage-VL
> Codec-native image/video understanding plus proactive full-duplex streaming.
## Summary
- Vendor: Microsoft
- Model: [`microsoft/Mage-VL`](https://huggingface.co/microsoft/Mage-VL)
- Task: Image understanding, video understanding, codec-native video understanding,
and event-gated streaming commentary
- Mode: Full-duplex WebSocket server and session adapter in
`vllm_omni/experimental/fullduplex/mage_vl`
- Maintainer: Community
## Status
This experimental integration provides the Mage-VL adapter contract: bounded causal
visual windows, codec-window input, rolling gate evaluation, explicit query handling,
per-session state isolation, and stale-output cancellation through the shared
full-duplex runtime.
The production transport exposes `WebSocket /v1/mage-vl/duplex`, with bearer-token
authentication, bounded concurrent sessions and message sizes, idle timeouts,
disconnect cancellation, per-session state cleanup, and serialized access to the
shared Transformers checkpoint. Native vLLM model execution remains follow-up work.
The default `frames` backend uses decoded frames and has no codec preprocessing
requirement. The optional `codec` backend requires `ffmpeg` and `ffprobe`, plus the
`cv-preinfer` executable provided by `codec-video-prep>=0.2.5`, on `PATH`. Server
startup fails with an actionable error when `cv-preinfer` is unavailable; use
`--video-backend frames` when codec preprocessing is not installed.
## Serving
```bash
python -m vllm_omni.experimental.fullduplex.mage_vl.serving.server \
--model microsoft/Mage-VL --host 0.0.0.0 --port 8090 \
--auth-token "$MAGE_VL_API_KEY" --max-sessions 32 \
--video-backend frames --num-frames 8 --target-fps 1 \
--gate-threshold 0.5 --attn-impl sdpa
python examples/online_serving/mage_vl/duplex_client.py \
--url ws://127.0.0.1:8090/v1/mage-vl/duplex \
--auth-token "$MAGE_VL_API_KEY" \
--video /path/to/segment.mp4 \
--prompt "Describe this video segment in detail."
```
Each input window is a JSON `input.append` event whose `data.data.video_base64`
contains an MP4 segment. A text `input.append` before the video requests an explicit
answer; without text, the StreamMind gate decides whether to respond proactively.
The server emits `session.created`, `response.created`, one or more
`response.delta`, and `response.done` events. Clients can send `response.cancel` for
barge-in and `close` for graceful shutdown.
The adapter can still be wired to another backend through `--adapter module:factory`,
including:
- a Transformers `trust_remote_code=True` backend for offline/local testing;
- the model author's SGLang-compatible serving path;
- a future native vLLM `mage_vl` backend.
## Adapter Shape
```python
from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig
from vllm_omni.experimental.fullduplex.mage_vl import MageVLDuplexAdapter, MageVLDuplexRuntime
adapter = MageVLDuplexAdapter(gate=run_mage_gate, generate=run_mage_decoder)
session = DuplexSession(
"stream-1",
DuplexSessionConfig(
input_modalities=("codec_window", "video", "image", "text"),
output_modalities=("text",),
proactive=True,
),
)
runtime = MageVLDuplexRuntime(session, adapter)
await runtime.run(input_events, emit_event)
```
Input events use the shared full-duplex protocol. For decoded video/image, pass the frame
payload directly. For codec-native streaming, pass `modality="codec_window"` and include
fields such as `segment_id`, `pts_ms`, `duration_ms`, `codec`, and `metadata`.
```python
{
"type": "input.append",
"modality": "codec_window",
"data": {
"kind": "h264",
"segment_id": "seg-42",
"pts_ms": 42000,
"duration_ms": 1000,
"codec": {
"motion_vectors": ...,
"residual_energy": ...,
},
"metadata": {"gop": "P"},
},
}
```
The gate callable receives the current rolling window and returns either a
`MageVLGateDecision`, a mapping with `should_respond` / `text` / `event_id`, a bool, or
a text string. Repeated `event_id` values are de-duplicated at the session adapter.
## Verification
```bash
pytest tests/e2e/features/fullduplex/test_mage_vl_adapter.py \
tests/examples/offline_inference/test_mage_vl.py
```
+21 -54
View File
@@ -278,7 +278,7 @@ steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
### Attention Backends ### Attention Backends
On supported datacenter Blackwell systems, MiniMax H3 defaults to dense BF16 On datacenter Blackwell GPUs, MiniMax H3 defaults to dense BF16
`TRTLLM_ATTN`; no attention backend flag is required. To select it explicitly, `TRTLLM_ATTN`; no attention backend flag is required. To select it explicitly,
use: use:
@@ -298,33 +298,17 @@ FA4 remains available by explicitly selecting the `FLASH_ATTN` backend:
--diffusion-attention-backend FLASH_ATTN --diffusion-attention-backend FLASH_ATTN
``` ```
With the optional `[fa4]` dependency installed, `FLASH_ATTN` prefers FA4 on On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
Blackwell. Confirm the server log contains `Using CuTe FlashAttention-4 on `Using CuTe FlashAttention-4 on Blackwell` before recording FA4 measurements.
Blackwell` before recording FA4 measurements.
`TRTLLM_ATTN` additionally offers two **lossy** optimizations for the long main `TRTLLM_ATTN` additionally supports two **lossy** optimizations for the long main
DiT attention sequence: SAGE attention quantization and Skip-Softmax sparse DiT attention sequence: SAGE attention quantization and Skip-Softmax Sparse
attention. Both work under the pure Ulysses parallelism of the profile above Attention. SAGE quantizes Q/K to the configured dtype and V to FP8. This example uses
(`--usp 4 --ring 1`). The example below enables both: `fp8_e4m3` for Q/K; B200 also supports `int8` Q/K. The TRTLLM SAGE path fixes V
to FP8, so vLLM-Omni only exposes the Q/K dtype. The token refiner is a short
- SAGE with `fp8_e4m3` Q/K; P and V are always FP8 in this kernel. B200 attention path, so the `per_role` override leaves SAGE and Skip-Softmax disabled
additionally supports `int8` Q/K, which preserves accuracy better than FP8. for it. The example enables the calibration-free Skip-Softmax path with
- Skip-Softmax with a direct `threshold=0.05` (the calibrated `threshold=0.05`, after the normalized timestep reaches `0.97`:
`target_sparsity` control needs ModelOpt metadata that the official H3
checkpoint does not include). Together with the cutoff below this is a
**conservative** setting: a low threshold skips only clearly negligible tiles,
and a cutoff close to `1.0` leaves a substantial dense prefix. Raise
`threshold` or lower `disabled_until_timestep` for more speedup once quality
is verified.
- `disabled_until_timestep=0.97` keeps the early high-noise steps dense. The
gate compares against the video sigma, which H3's default flow shift of 12
keeps high for much of the run: at 50 steps, `0.99`, `0.97`, and `0.95`
leave the first 6, 14, and 19 of 49 denoiser forwards dense. See the
[Skip-Softmax design](https://github.com/vllm-project/vllm-omni/blob/main/docs/design/feature/skip_softmax.md#timestep-gating)
for how the cutoff maps to steps.
- A `per_role` entry that keeps the token refiner, a short attention path,
dense. A per-role spec does not inherit `quant` or `skip_softmax` from
`default`.
```bash ```bash
--diffusion-attention-config '{ --diffusion-attention-config '{
@@ -348,13 +332,10 @@ attention. Both work under the pure Ulysses parallelism of the profile above
}' }'
``` ```
Both optimizations trade fidelity for speed and their effects compound. For configuration details, see
Compare against dense output on the same prompt and seed before adopting them. [TRTLLM_ATTN Backend and Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-backend-and-skip-softmax)
For the full key reference, see
[Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends/trtllm.md#skip-softmax)
and and
[SAGE quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends/trtllm.md#sage-quantization) [TRTLLM_ATTN SAGE Quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-sage-quantization).
in the TRTLLM attention guide.
### Text encoder tensor parallelism ### Text encoder tensor parallelism
@@ -981,27 +962,13 @@ mostly base H3 weights on a four-step schedule. Offload is refused for the same
reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and
`--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast. `--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast.
The VSA variants are supported through FastVideo's external kernel. Install a > [!NOTE]
`fastvideo-kernel` build that provides the `fastvideo_kernel` Python module, > Only the **Dense / Data-Free** variant is supported today. The three VSA
then add the following flags to the same command: > variants carry `.set_weight` compression gates for modules that do not exist
> in vLLM-Omni's dense H3 attention, and the release states dense attention is
```bash > not a supported substitute for them, so they are refused at startup rather
--diffusion-attention-backend FASTVIDEO_VSA \ > than silently run wrong. Video Sparse Attention for H3's packed
--fastvideo-vsa-topk 64 > `[text | cond | audio | video]` sequence is follow-up work.
```
FastH3 VSA applies its learned `.set_weight` compression gates to the complete
packed `[text | cond | audio | video]` document using the official H3 geometry:
text/condition/audio prefix tiles never cross segment boundaries, and target
video rows use `(4, 4, 4)` 3-D tiles (64 tokens). Prefix queries remain dense;
video queries select all prefix tiles plus the configured top-k video tiles.
Pure Ulysses sequence parallelism is supported: the learned gate follows the
same sequence-to-head all-to-all as Q/K/V before VSA runs. Ring and all-gather
SP remain unsupported because they do not present a complete packed sequence
to each block-sparse kernel rank.
The Dense / Data-Free variant does not require `fastvideo-kernel` and should be
served with a dense attention backend.
Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at
1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded: 1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded:
+2 -101
View File
@@ -119,30 +119,6 @@ MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
bash examples/online_serving/text_to_video/run_curl_sana_video.sh bash examples/online_serving/text_to_video/run_curl_sana_video.sh
``` ```
##### Parallel native serving
The native pipelines support tensor parallelism (up to 2 GPUs) and CFG
parallelism. Tensor parallelism on 2 GPUs:
```bash
vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
--omni \
--model-class-name SanaVideoPipeline \
--tensor-parallel-size 2 \
--dtype bfloat16 \
--port 8091
```
For CFG parallelism on 2 GPUs, replace `--tensor-parallel-size 2` with
`--cfg-parallel-size 2`; it splits the guided and unguided branches across the
two GPUs and only helps when `guidance_scale` is above 1. Passing both flags
combines the two on 4 GPUs.
Whether tensor parallelism lowers latency depends on the interconnect. Each
transformer block adds an all-reduce, so it speeds up generation only on fast
GPU links such as NVLink and can be slower than a single GPU on PCIe-only
systems. Measure on your hardware before enabling it.
To run the black-box compatibility backend for T2V, replace the server script To run the black-box compatibility backend for T2V, replace the server script
with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request
works; `num_frames` is adapted to Diffusers' `frames` argument. The script works; `num_frames` is adapted to Diffusers' `frames` argument. The script
@@ -202,75 +178,6 @@ Use the native `SanaVideoPipeline` and `SanaImageToVideoPipeline` for the
primary SANA execution paths. The Diffusers adapter is retained as a primary SANA execution paths. The Diffusers adapter is retained as a
validated compatibility/reference backend. validated compatibility/reference backend.
#### Native Cache-DiT and CPU offload
The native T2V and I2V pipelines support Cache-DiT, model-level CPU offload,
and layerwise CPU offload through the common diffusion flags. Layerwise mode
keeps non-block transformer modules, the text encoder, and the VAE on the
runtime device while prefetching DiT blocks in order.
Add one of the following flag sets to either native offline command above:
```bash
# Cache-DiT
--cache-backend cache_dit
# Model-level CPU offload
--enable-cpu-offload
# Layerwise CPU offload
--enable-layerwise-offload
# Cache-DiT plus model-level offload
--cache-backend cache_dit --enable-cpu-offload
# Cache-DiT plus layerwise offload
--cache-backend cache_dit --enable-layerwise-offload
# Distributed layerwise offload (weights sharded across the DP group)
--enable-distributed-layerwise-offload
```
If both CPU offload flags are supplied, the common offloader keeps its existing
layerwise precedence. Cache-DiT refreshes each request from the explicit
`num_inference_steps`; when that field is omitted, both native SANA pipelines
default to 50 inference steps.
Cache-DiT and CPU offload require TP1, CFG1, and SP1. Combining them with
tensor, CFG, or sequence parallelism raises before checkpoint components are
loaded, as do other cache backends such as TeaCache. Cache-DiT cannot be
combined with distributed layerwise offload: per-rank cache skips would
desynchronize the weight AllGather.
Measured on one A800-SXM4-80GB (driver 580.126.09, CUDA 13.0, PyTorch
2.11.0+cu130, Diffusers 0.38.0) with `DIFFUSION_ATTENTION_BACKEND=CUDNN_ATTN`,
81 frames, 50 steps, seed 42, each configuration in its own process:
| Configuration | 480p T2V latency | 480p T2V generation peak |
|---|---:|---:|
| Baseline | 120.19 s | 24.07 GiB |
| Cache-DiT | 77.11 s (1.56x) | 24.07 GiB |
| Model CPU offload | 121.15 s | 16.81 GiB (-7.26 GiB) |
| Layerwise offload | 118.63 s | 20.75 GiB (-3.32 GiB) |
Cache-DiT gives 1.54x on 720p I2V (37.0 s to 24.1 s). Latency is the median of
three measured runs after one warmup. Offload trades startup peak too: 14.28
GiB baseline against 6.14 GiB for model-level and 6.80 GiB for layerwise.
Cache-DiT is approximate: it skips block computations, so its output differs
from an uncached run. Similarity against a reference run is not used as a gate
here because this pipeline does not reproduce its own trajectory on this
hardware -- two runs of identical code differ by SSIM 0.9712 / rel_l2 8.7e-2 at
50 steps. Cache-DiT output was checked frame by frame against the uncached run
on 480p T2V and 720p I2V instead, and matches it in subject, composition,
sharpness, and lighting.
The numbers above use the offline example's default cache configuration.
Lowering `residual_diff_threshold` in `cache_config` caches fewer steps,
trading speed for a trajectory closer to the uncached run; see the
[Cache-DiT guide](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md)
for the tuning knobs.
#### Notes #### Notes
- Key flags: select I2V explicitly with `--model-class-name - Key flags: select I2V explicitly with `--model-class-name
@@ -287,14 +194,8 @@ for the tuning knobs.
The denoising loop intentionally retains the checkpoint-compatible The denoising loop intentionally retains the checkpoint-compatible
Diffusers `DPMSolverMultistepScheduler`. Diffusers `DPMSolverMultistepScheduler`.
- Known limitations: - Known limitations:
- Tensor parallelism (up to 2 GPUs) and CFG parallelism are supported for - Sequence/tensor/CFG parallelism, Cache-DiT, TeaCache, and step execution
the native pipeline. Sequence parallelism, TeaCache, and step execution are not validated for the native pipeline.
are not validated.
- Cache-DiT and CPU offload are limited to TP1/CFG1/SP1. Cache-DiT with
distributed layerwise offload is not supported by the native pipeline.
- Cache-DiT speedup and offload memory numbers above are single-GPU A800
measurements; other hardware will differ. Distributed layerwise offload is
only exercised single-rank here.
- The Diffusers backend is a compatibility path and does not provide native - The Diffusers backend is a compatibility path and does not provide native
vLLM-Omni parallelism or continuous batching. vLLM-Omni parallelism or continuous batching.
- Native describes pipeline and Transformer ownership, not a zero-Diffusers - Native describes pipeline and Transformer ownership, not a zero-Diffusers
-1
View File
@@ -166,7 +166,6 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
- Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config. - Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`. - Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`.
- Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail. - Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
- Stored voices: Uploaded and precomputed voices use the Base task and require a Base checkpoint. Built-in preset speakers remain CustomVoice voices.
- Base codec termination: Base requests without an explicit `max_new_tokens` - Base codec termination: Base requests without an explicit `max_new_tokens`
use a text-scaled safety ceiling (at least 192 codec frames and no more than use a text-scaled safety ceiling (at least 192 codec frames and no more than
the configured model limit). If the Talker reaches that ceiling without the configured model limit). If the Talker reaches that ceiling without
-1
View File
@@ -78,7 +78,6 @@ recipes/
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB | | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X | | [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB / 1x AMD MI300X |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X | | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X |
| [`SenseNova/SenseNova-U1.5.md`](./SenseNova/SenseNova-U1.5.md) | Unified image generation and understanding, with an 8-step distilled LoRA | 1x A800 80GB |
| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X | | [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X |
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB | | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 | | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
-159
View File
@@ -1,159 +0,0 @@
# SenseNova-U1.5
> Unified image generation and understanding, with an official 8-step distilled LoRA
## Summary
- Vendor: SenseNova
- Model: `sensenova/SenseNova-U1.5-8B-MoT`
- LoRA: `sensenova/SenseNova-U1.5-8B-MoT-LoRAs` (`SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors`)
- Task: text2img, img2img, img2text (visual understanding), text2text (chat)
- Mode: Offline inference, Online serving (OpenAI-compatible API)
- Maintainer: Community
## When to use this recipe
U1.5 runs on the same `SenseNovaU1Pipeline` as U1 — the checkpoint keeps
`model_type: neo_chat`, so it resolves without `--model-class-name`. Relative to U1 the
config flips two fields, `use_pixel_head` to `true` (the flow-matching head becomes a
`ConvDecoder`) and `noise_scale_max_value` from 8.0 to 16.0; both are already read by
`SenseNovaU1Config`. The checkpoint is 13 shards / 50.2 GB on disk, 30.3 GB of that fp32, and
loads to roughly 34 GB in bf16.
Use this recipe for U1.5 specifically, including its distilled few-step LoRA. For U1 see
[SenseNova-U1](SenseNova-U1.md).
## Hardware Support
### GPU
#### 1x A800 80GB
- 1x NVIDIA A800 80GB PCIe with an Intel Xeon Gold 6336Y, driver 595.84, Linux
- Python 3.12.13, vLLM 0.28.0, torch 2.13.0+cu130, CUDA 13.0, diffusers 0.40.0,
transformers 5.14.1, flashinfer 0.6.16.post3, BF16, TP=1
- Peak GPU memory: 34.3 GB at 1024x1024, 36.4 GB at 1536x2720
##### Text-to-Image
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model sensenova/SenseNova-U1.5-8B-MoT \
--prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \
--width 1024 --height 1024 \
--seed 42 --num-inference-steps 50 --cfg-scale 4.0 \
--extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
--output sensenova_u15_t2i.png
```
Think mode (`"think": true`) is recommended for higher image quality.
##### Image-to-Image Editing
```bash
python examples/offline_inference/image_to_image/image_edit.py \
--model sensenova/SenseNova-U1.5-8B-MoT \
--prompt "Turn this into an oil painting" \
--image input.png --resolution 1024 \
--seed 42 --num-inference-steps 25 --cfg-scale 4.0 \
--extra-args '{"think": true, "img_cfg_scale": 1.0, "cfg_norm": "none", "timestep_shift": 3.0}' \
--output sensenova_u15_edit.png
```
##### 8-step distilled LoRA
The distilled LoRA is fused into the generation tower at load time:
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model sensenova/SenseNova-U1.5-8B-MoT \
--lora-path SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors --lora-backend distill \
--prompt "Close portrait of an elderly woman by a farmhouse window, warm natural light." \
--width 1024 --height 1024 \
--seed 42 --num-inference-steps 8 --cfg-scale 1.0 \
--extra-body '{"think": false, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \
--output sensenova_u15_lora8.png
```
**Use `--cfg-scale 1.0` with this LoRA.** It is distilled with DMD and runs without
classifier-free guidance; the default `4.0` applies guidance twice and produces a
blown-out, posterised image.
##### Online serving
```bash
vllm serve sensenova/SenseNova-U1.5-8B-MoT --omni --port 8091
python examples/online_serving/sensenova_u1/openai_chat_client.py \
-s http://127.0.0.1:8091 -m img2text -i input.png -p "Describe this image."
```
`-s` takes the base URL; the client appends `/v1` itself.
#### Measured latency (1x A800 80GB, 25 steps, median of 3 after a warmup)
| Resolution | Step latency | Total |
| --- | --- | --- |
| 1024x1024 | 209.5 ms | 5.24 s |
| 1536x1536 | 458.5 ms | 11.46 s |
At 1536x2720 with 50 steps, end-to-end is 43.0 s.
#### Measured latency (1x NVIDIA H200 139GB, single run, no warmup)
Reported by @hsliuustc0106 against PR head `29c090d`, on one reserved CUDA device.
- Python 3.12.13, vLLM 0.28.0, torch 2.13.0+cu130, CUDA 13.0, diffusers 0.40.0,
transformers 5.14.1, flashinfer 0.6.16.post3, BF16, TP=1
- `sensenova/SenseNova-U1.5-8B-MoT` with `SenseNova-U1.5-8B-MoT-LoRA-8step.safetensors`,
seed 42, 1024x1024
| Case | Stage latency | Peak GPU memory |
| --- | --- | --- |
| 50 steps, cfg 4.0, think off | 9017 ms | 34384 MiB |
| 8 steps, cfg 1.0, no LoRA control | 624.00 ms, 77.85 ms/step | 34376 MiB |
| 8 steps, cfg 1.0, distilled LoRA | 623.85 ms, 77.98 ms/step | 34364 MiB |
| 50 steps, cfg 4.0, think on, `VLLM_OMNI_SENSENOVA_PAGED_DECODE=1` | 81622 ms | 34594 MiB |
Each row runs the Text-to-Image command above with `--num-inference-steps` and
`--cfg-scale` as listed and `think` set through `--extra-body`; the LoRA row adds
`--lora-path` and `--lora-backend distill`.
Single runs without a preceding warmup, so the think-on row carries the first-request
compile rather than a steady state. The distilled LoRA fused into 168 parameters, and its
image differs from the same-seed no-LoRA control in 1,048,574 of 1,048,576 pixels, MAE 58.62.
```bash
pytest -m "cpu and not cuda" tests/diffusion/models/sensenova_u1/ tests/diffusion/lora/test_loader.py -q
pytest -m "cpu and not cuda" tests/config/test_environment_variables.py tests/diffusion/test_diffusion_worker.py -q
pytest -m cuda tests/diffusion/models/sensenova_u1/test_sensenova_u1_attention_gqa.py \
tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py -q --tb=short
```
71, 13 and 9 passed. Serving reached `/health`; `img2text` returned a 1921-character
description and `text2text` the expected answer.
#### Verification
```bash
pytest -q tests/diffusion/models/sensenova_u1/
```
## Notes
- Both `use_pixel_head` and `noise_scale_max_value` come from `config.json`; no flag is needed.
- The checkpoint carries no `configuration_neo_chat.py`, so the loader logs a few
"does not appear to have a file named configuration_neo_chat.py" errors during startup and
then proceeds on the in-tree config. Generation is unaffected.
- The LoRA targets the generation tower only (`*_mot_gen`); understanding-tower weights are
untouched.
- Autoregressive decode (think, text-to-text and image-to-text) runs on a paged K/V cache under
a captured CUDA graph. It falls back to the ordinary cache when the device or the bundled
`flash_attn_varlen_func` cannot support it; set `VLLM_OMNI_SENSENOVA_PAGED_DECODE=0` to force
that fallback.
- The first request after startup costs about 0.7 s more than the steady state whether the paged
path is on or off. Measured on one A800 with the inductor, triton and vLLM compile caches all
cleared, median of three runs: 718 ms above steady with the path on, 679 ms with it off.
- Each request captures its own graphs, and a think request captures twice because the sequence
grows past the 512 bucket, so the capture cost is paid per request rather than once at
startup.
@@ -1 +0,0 @@
首先,我分析所有输入图像:图像1是一个圆形的logo,设计现代且抽象。它由不同色调的蓝色(深蓝、中蓝、浅蓝)和白色构成,这些色块以流畅的曲线相互交织,形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形,背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴,<recaption>这幅图像以产品摄影的精致风格,呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。冰箱贴居中放置在`image_2`的灰色织物背景上,其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感:表面覆盖着一层光滑如镜的珐琅釉面,反射出柔和而清晰的高光;图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒,增强了立体感。柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡,并在其下方投下淡淡的、轮廓模糊的阴影,使其与织物背景无缝融合,营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓,以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构,精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果,以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>
@@ -1 +0,0 @@
首先,我分析所有输入图像:图像1是一个圆形logo,背景为黑色,主体由深蓝、中蓝、浅蓝和白色组成的抽象曲线构成,整体呈现出一种动态的、类似旋涡的视觉效果。图像2展示了四个并排摆放在灰色织物表面上的冰箱贴。这些冰箱贴的造型分别是长颈鹿、小鸟、蜜蜂和海豚,其关键特征是材质:它们具有光滑、高反光的珐琅质感,边缘有明显的银色金属包边,呈现出立体感。用户的指令是“基于图一的logo,参考图二中冰箱贴的材质,制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计(logo)转化为一个具有特定三维物理属性(冰箱贴材质)的实体物品。核心在于保留logo的设计精髓,同时精确复刻参考图像中冰箱贴的材质和工艺感。为了构建最终的答案图像,我会将图一的圆形logo作为新冰箱贴的主体图案。这个圆形冰箱贴将拥有与图二冰箱贴一致的厚度和立体感。其表面会覆盖一层高光泽的透明釉质,完美再现珐琅的质感。logo内部的蓝色和白色区域将保持其原有的色彩,但会因为表面的光泽而出现明亮的高光反射。整个圆形的边缘会包裹一圈纤细而有光泽的银色金属边框,这不仅增加了产品的精致感,也与图二中的金属包边工艺相呼应。最后,将这个新制作的圆形冰箱贴放置在图二所示的灰色织物背景上,确保其光照、阴影和反射效果与周围环境和其他冰箱贴完全一致,使其看起来就像是和它们一起被拍摄的,从而创造出一个逼真且风格统一的视觉效果。</think><recaption>这幅图像以产品摄影的精准风格,展示了一枚根据`image_1`的logo设计定制的圆形冰箱贴。最终图像使用`image_2`的分辨率。该冰箱贴居中放置于`image_2`的灰色织物背景之上,其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝及白色构成的动态旋涡状抽象图案。整个冰箱贴被赋予了与`image_2`中冰箱贴完全一致的高级珐琅质感,表面覆盖着一层晶莹剔透的高光釉面,在柔和的顶光照射下,呈现出清晰而温润的高光反射。一圈纤薄而精致的银色金属边框紧密包裹着圆形主体,其金属光泽与釉面的质感形成精妙对比。冰箱贴在织物背景上投下了一层极其自然的柔和阴影,证明了其与环境光影的无缝融合,整体视觉效果真实可信,宛如一件精心制作的实体工艺品。<relation_1>最终图像完整保留了`image_1`中logo的全部核心设计元素。这包括其完美的圆形轮廓,以及内部由深蓝、中蓝、浅蓝和白色四种颜色构成的抽象旋涡图案的精确形态、色彩和布局。</relation_1><relation_2>最终图像的分辨率、背景和材质风格完全来自`image_2`。这包括其作为背景的灰色织物表面的独特纹理,以及应用于新冰箱贴上的、与`image_2`中冰箱贴完全相同的珐琅材质(高光釉面质感)和银色金属边框的物理属性。画面的整体光照环境,包括柔和的光源方向和质感,也与`image_2`保持一致。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep><guidance><timestep_r>[<img>]{3600}<eoi></answer>
Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.7 MiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 1.6 MiB

-42
View File
@@ -1,42 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from pathlib import Path
from shlex import split
import pytest
import yaml
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
AMD_MERGE_PIPELINE = Path(".buildkite/amd/test-amd-merge.yml")
def _find_step(label: str) -> dict:
pipeline = yaml.safe_load(AMD_MERGE_PIPELINE.read_text(encoding="utf-8"))
def walk(steps: list[dict]) -> dict | None:
for step in steps:
if step.get("label") == label:
return step
if nested := walk(step.get("steps", [])):
return nested
return None
step = walk(pipeline.get("steps", []))
assert step is not None, f"missing AMD pipeline step: {label}"
return step
def test_qwen3_tts_base_preserves_advanced_model_arguments() -> None:
step = _find_step("Qwen3-TTS Base E2E Test")
commands = step["commands"]
assert all("bash -c" not in command for command in commands)
pytest_command = next(command for command in commands if "pytest" in command)
argv = split(pytest_command)
marker_index = argv.index("-m")
run_level_index = argv.index("--run-level")
assert argv[marker_index + 1] == "advanced_model and cuda"
assert argv[run_level_index + 1] == "advanced_model"
-1
View File
@@ -1,5 +1,4 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from __future__ import annotations from __future__ import annotations
import sys import sys
+3 -176
View File
@@ -13,7 +13,6 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "com
from skip_ci import resolve_ci_decision # noqa: E402 from skip_ci import resolve_ci_decision # noqa: E402
from upload_pipeline import ( # noqa: E402 from upload_pipeline import ( # noqa: E402
_expand_mirror_hardwares, _expand_mirror_hardwares,
_get_mirror_hw_selector,
_load_bootstrap_steps, _load_bootstrap_steps,
_render_bootstrap_pipeline, _render_bootstrap_pipeline,
_render_test_pipeline, _render_test_pipeline,
@@ -117,8 +116,7 @@ def test_yaml_gated_l2_still_enables_image_via_ready_base() -> None:
assert "if: true" not in rendered assert "if: true" not in rendered
def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest.MonkeyPatch) -> None: def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins() -> None:
monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "")
doc = { doc = {
"steps": [ "steps": [
{ {
@@ -131,29 +129,8 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest
rendered = _render_test_pipeline(doc, changed_files=None) rendered = _render_test_pipeline(doc, changed_files=None)
step = rendered["steps"][0] step = rendered["steps"][0]
assert "mirror_hardwares" not in step assert "mirror_hardwares" not in step
assert step["agents"]["queue"] == "l4-k8s" assert step["agents"]["queue"] == "gpu_1_queue"
assert step["retry"] == { assert step["plugins"][0]["docker#v5.2.0"]["image"].endswith("$BUILDKITE_COMMIT")
"automatic": [
{"exit_status": -1, "limit": 1},
{"exit_status": 128, "limit": 1},
{"signal_reason": "agent_stop", "limit": 1},
{"signal_reason": "agent_refused", "limit": 1},
],
}
container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0]
assert container["image"].endswith("$BUILDKITE_COMMIT")
assert container["resources"]["limits"]["nvidia.com/gpu"] == 1
def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
step = _expand_mirror_hardwares(
{
"label": "opt-out",
"mirror_hardwares": "l4_1",
"retry": {"automatic": [{"exit_status": 255, "limit": 2}]},
},
)
assert step["retry"] == {"automatic": [{"exit_status": 255, "limit": 2}]}
def test_mirror_hardwares_conflicts_with_explicit_agents() -> None: def test_mirror_hardwares_conflicts_with_explicit_agents() -> None:
@@ -184,156 +161,6 @@ 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]: def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
rendered = _render_test_pipeline(doc, changed_files=changed_files) rendered = _render_test_pipeline(doc, changed_files=changed_files)
labels: set[str] = set() labels: set[str] = set()
-63
View File
@@ -2399,21 +2399,6 @@ class TestBaseConfigInheritance:
# CI overrides max_tokens # CI overrides max_tokens
assert s0["max_tokens"] == 150 assert s0["max_tokens"] == 150
def test_qwen3_omni_colocate_async_bounds_only_rocm_kv_cache(self):
ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe_colocate_async.yaml"))
pipeline = resolve_pipeline_config("qwen3_omni_moe_thinker_only")
assert isinstance(pipeline, PipelineConfig)
cuda = _apply_platform_overrides(load_deploy_config(ci_path), platform="cuda")
cuda_stage = merge_pipeline_deploy(pipeline, cuda)[0]
assert cuda_stage.yaml_engine_args["gpu_memory_utilization"] == 0.9
assert "kv_cache_memory_bytes" not in cuda_stage.yaml_engine_args
rocm = _apply_platform_overrides(load_deploy_config(ci_path), platform="rocm")
rocm_stage = merge_pipeline_deploy(pipeline, rocm)[0]
assert "gpu_memory_utilization" not in rocm_stage.yaml_engine_args
assert rocm_stage.yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3
def test_pure_inheritance_overlay(self, tmp_path): def test_pure_inheritance_overlay(self, tmp_path):
"""An overlay with only ``base_config`` inherits everything.""" """An overlay with only ``base_config`` inherits everything."""
base = Path(get_deploy_config_path("qwen3_omni_moe.yaml")) base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
@@ -2875,54 +2860,6 @@ stages:
assert omega_config.engine_args.max_model_len == 8192 assert omega_config.engine_args.max_model_len == 8192
assert omega_config.engine_args.enforce_eager is True assert omega_config.engine_args.enforce_eager is True
def test_cli_attention_shorthand_replaces_yaml_structured_default(self):
stage = StageConfig(
stage_id=0,
model_stage="dit",
stage_type=StageType.DIFFUSION,
yaml_engine_args={
"diffusion_attention_config": {
"default": {"backend": "FLASH_ATTN"},
"per_role": {"cross": {"backend": "TORCH_SDPA"}},
},
},
runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"},
)
engine_args = stage.to_omegaconf().engine_args
assert engine_args.diffusion_attention_backend == "SAGE_ATTN"
assert "default" not in engine_args.diffusion_attention_config
assert engine_args.diffusion_attention_config.per_role.cross.backend == "TORCH_SDPA"
def test_cli_attention_shorthand_keeps_yaml_per_role_only_config(self):
stage = StageConfig(
stage_id=0,
model_stage="dit",
stage_type=StageType.DIFFUSION,
yaml_engine_args={"diffusion_attention_config": {"per_role": {"cross": {"backend": "TORCH_SDPA"}}}},
runtime_overrides={"diffusion_attention_backend": "SAGE_ATTN"},
)
engine_args = stage.to_omegaconf().engine_args
assert engine_args.diffusion_attention_backend == "SAGE_ATTN"
assert engine_args.diffusion_attention_config == {"per_role": {"cross": {"backend": "TORCH_SDPA"}}}
def test_cli_attention_config_replaces_yaml_shorthand(self):
stage = StageConfig(
stage_id=0,
model_stage="dit",
stage_type=StageType.DIFFUSION,
yaml_engine_args={"diffusion_attention_backend": "TORCH_SDPA"},
runtime_overrides={"diffusion_attention_config": {"default": {"backend": "SAGE_ATTN"}}},
)
engine_args = stage.to_omegaconf().engine_args
assert "diffusion_attention_backend" not in engine_args
assert engine_args.diffusion_attention_config.default.backend == "SAGE_ATTN"
class TestSentinelDefaultPrecedence: class TestSentinelDefaultPrecedence:
"""Caller-typed (non-None) values win over YAML; None values fall through """Caller-typed (non-None) values win over YAML; None values fall through
+2 -2
View File
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
EnvironmentVariableCategory.PUBLIC_OMNI: 24, EnvironmentVariableCategory.PUBLIC_OMNI: 24,
EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.INHERITED_VLLM: 20,
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
EnvironmentVariableCategory.MODEL_SPECIFIC: 58, EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
EnvironmentVariableCategory.INTERNAL: 2, EnvironmentVariableCategory.INTERNAL: 2,
} }
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
) )
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
ModelEnvironmentVariableDisposition.PROMOTE: 35, ModelEnvironmentVariableDisposition.PROMOTE: 33,
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.EXTERNAL: 0,
ModelEnvironmentVariableDisposition.INTERNALIZE: 11, ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
@@ -50,6 +50,9 @@ from tests.dfx.conftest import (
) )
from tests.helpers.runtime import get_open_port 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 # Inline field processing
@@ -834,7 +837,7 @@ def run_benchmark(
"Quantization": _to_quantization_value(server_type, serve_args_dict), "Quantization": _to_quantization_value(server_type, serve_args_dict),
"offload": _to_offload_string(server_type, serve_args_dict), "offload": _to_offload_string(server_type, serve_args_dict),
"compile": _to_compile_value(server_type, serve_args_dict), "compile": _to_compile_value(server_type, serve_args_dict),
"Attn_backend": str(serve_args_dict.get("diffusion-attention-backend") or ""), "Attn_backend": os.environ.get("DIFFUSION_ATTENTION_BACKEND", ""),
"num_inference_steps": params.get("num-inference-steps", ""), "num_inference_steps": params.get("num-inference-steps", ""),
"completed": completed, "completed": completed,
"failed": failed, "failed": failed,
@@ -1,7 +1,7 @@
[ [
{ {
"test_name": "test_boogu_image_edit_single_device", "test_name": "test_boogu_image_edit_single_device",
"description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. This row measures text-only edit guidance (2 predictions/step); dedicated rows below use extra-body to measure double guidance. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.", "description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. Text-only edit path (2 predictions/step); a double-guidance (guidance_scale_2) case is a follow-up since the random dataset cannot inject extra_body params. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.",
"server_type": "vllm-omni", "server_type": "vllm-omni",
"server_params": { "server_params": {
"model": "Boogu/Boogu-Image-0.1-Edit", "model": "Boogu/Boogu-Image-0.1-Edit",
@@ -33,184 +33,5 @@
} }
} }
] ]
},
{
"test_name": "test_boogu_image_edit_double_guidance_single_device",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 1
}
},
"full_model",
"diffusion"
],
"description": "Single-device sequential three-branch double-guidance control row for direct A/B comparison with cfg_parallel_size=2 and 3 on identical H100 hardware.",
"server_type": "vllm-omni",
"server_params": {
"model": "Boogu/Boogu-Image-0.1-Edit",
"serve_args": {
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "512x512_steps28_i2i_double_cfg1",
"dataset": "random",
"task": "i2i",
"width": 512,
"height": 512,
"num-inference-steps": 28,
"num-prompts": 10,
"max-concurrency": 1,
"warmup-requests": 4,
"warmup-concurrency": 1,
"warmup-num-inference-steps": 28,
"enable-negative-prompt": true,
"extra-body": {
"guidance_scale": 5.0,
"guidance_scale_2": 2.0,
"seed": 42
}
}
]
},
{
"test_name": "test_boogu_image_edit_text_cfg_parallel_2",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 2
}
},
"full_model",
"diffusion"
],
"description": "Two-device CFG-parallel Edit text-only guidance measurement row. Compare with the single-device row on identical H100 hardware and software.",
"server_type": "vllm-omni",
"server_params": {
"model": "Boogu/Boogu-Image-0.1-Edit",
"serve_args": {
"cfg-parallel-size": 2,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "512x512_steps28_i2i_text_cfg2",
"dataset": "random",
"task": "i2i",
"width": 512,
"height": 512,
"num-inference-steps": 28,
"num-prompts": 10,
"max-concurrency": 1,
"warmup-requests": 4,
"warmup-concurrency": 1,
"warmup-num-inference-steps": 28,
"enable-negative-prompt": true,
"extra-body": {
"guidance_scale": 5.0,
"guidance_scale_2": 1.0,
"seed": 42
}
}
]
},
{
"test_name": "test_boogu_image_edit_double_cfg_parallel_2",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 2
}
},
"full_model",
"diffusion"
],
"description": "Two-device round-robin execution of the three Edit double-guidance branches. Populate A/B latency, throughput, and per-GPU memory after GPU validation.",
"server_type": "vllm-omni",
"server_params": {
"model": "Boogu/Boogu-Image-0.1-Edit",
"serve_args": {
"cfg-parallel-size": 2,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "512x512_steps28_i2i_double_cfg2",
"dataset": "random",
"task": "i2i",
"width": 512,
"height": 512,
"num-inference-steps": 28,
"num-prompts": 10,
"max-concurrency": 1,
"warmup-requests": 4,
"warmup-concurrency": 1,
"warmup-num-inference-steps": 28,
"enable-negative-prompt": true,
"extra-body": {
"guidance_scale": 5.0,
"guidance_scale_2": 2.0,
"seed": 42
}
}
]
},
{
"test_name": "test_boogu_image_edit_double_cfg_parallel_3",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 3
}
},
"full_model",
"diffusion"
],
"description": "Full three-device execution of the three Edit double-guidance branches. Compare directly with sequential and cfg_parallel_size=2 on identical H100 hardware.",
"server_type": "vllm-omni",
"server_params": {
"model": "Boogu/Boogu-Image-0.1-Edit",
"serve_args": {
"cfg-parallel-size": 3,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "512x512_steps28_i2i_double_cfg3",
"dataset": "random",
"task": "i2i",
"width": 512,
"height": 512,
"num-inference-steps": 28,
"num-prompts": 10,
"max-concurrency": 1,
"warmup-requests": 4,
"warmup-concurrency": 1,
"warmup-num-inference-steps": 28,
"enable-negative-prompt": true,
"extra-body": {
"guidance_scale": 5.0,
"guidance_scale_2": 2.0,
"seed": 42
}
}
]
} }
] ]
@@ -33,49 +33,5 @@
} }
} }
] ]
},
{
"test_name": "test_boogu_image_cfg_parallel_2",
"mark": [
{
"hardware_marks": {
"res": {
"cuda": "H100"
},
"num_cards": 2
}
},
"full_model",
"diffusion"
],
"description": "Two-device CFG-parallel Base T2I measurement row. Run beside test_boogu_image_single_device on the same H100 revision and populate latency, throughput, speedup, and per-GPU peak-memory baselines from the resulting report.",
"server_type": "vllm-omni",
"server_params": {
"model": "Boogu/Boogu-Image-0.1-Base",
"serve_args": {
"cfg-parallel-size": 2,
"enable-diffusion-pipeline-profiler": true
}
},
"benchmark_params": [
{
"name": "512x512_steps28_cfg2",
"dataset": "random",
"task": "t2i",
"width": 512,
"height": 512,
"num-inference-steps": 28,
"num-prompts": 10,
"max-concurrency": 1,
"warmup-requests": 4,
"warmup-concurrency": 1,
"warmup-num-inference-steps": 28,
"enable-negative-prompt": true,
"extra-body": {
"guidance_scale": 4.0,
"seed": 42
}
}
]
} }
] ]
@@ -30,6 +30,7 @@
"vae-patch-parallel-size": 4, "vae-patch-parallel-size": 4,
"vae-parallel-mode": "tile", "vae-parallel-mode": "tile",
"vae-use-tiling": true, "vae-use-tiling": true,
"diffusion-attention-backend": "FLASH_ATTN",
"stage-init-timeout": 1800, "stage-init-timeout": 1800,
"init-timeout": 1800, "init-timeout": 1800,
"enable-diffusion-pipeline-profiler": true "enable-diffusion-pipeline-profiler": true
@@ -101,6 +102,7 @@
"vae-patch-parallel-size": 4, "vae-patch-parallel-size": 4,
"vae-parallel-mode": "tile", "vae-parallel-mode": "tile",
"vae-use-tiling": true, "vae-use-tiling": true,
"diffusion-attention-backend": "FLASH_ATTN",
"stage-init-timeout": 1800, "stage-init-timeout": 1800,
"init-timeout": 1800, "init-timeout": 1800,
"enable-diffusion-pipeline-profiler": true "enable-diffusion-pipeline-profiler": true
@@ -173,6 +175,7 @@
"vae-patch-parallel-size": 4, "vae-patch-parallel-size": 4,
"vae-parallel-mode": "tile", "vae-parallel-mode": "tile",
"vae-use-tiling": true, "vae-use-tiling": true,
"diffusion-attention-backend": "FLASH_ATTN",
"stage-init-timeout": 1800, "stage-init-timeout": 1800,
"init-timeout": 1800, "init-timeout": 1800,
"enable-diffusion-pipeline-profiler": true "enable-diffusion-pipeline-profiler": true
@@ -244,6 +247,7 @@
"vae-patch-parallel-size": 4, "vae-patch-parallel-size": 4,
"vae-parallel-mode": "tile", "vae-parallel-mode": "tile",
"vae-use-tiling": true, "vae-use-tiling": true,
"diffusion-attention-backend": "FLASH_ATTN",
"stage-init-timeout": 1800, "stage-init-timeout": 1800,
"init-timeout": 1800, "init-timeout": 1800,
"enable-diffusion-pipeline-profiler": true, "enable-diffusion-pipeline-profiler": true,
@@ -197,7 +197,6 @@ def test_create_unique_server_pytest_params_applies_marks(tmp_path):
assert len(by_id["test_with_mark"].values) == 1 assert len(by_id["test_with_mark"].values) == 1
assert isinstance(by_id["test_with_mark"].values[0], tuple) 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 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) assert not any(m.name == "H100" for m in by_id["test_without_mark"].marks)
@@ -18,7 +18,6 @@ import torch
import vllm_omni.diffusion.attention.layer as layer_mod import vllm_omni.diffusion.attention.layer as layer_mod
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata 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.attention.layer import Attention
from vllm_omni.diffusion.config import ( from vllm_omni.diffusion.config import (
get_current_diffusion_config, get_current_diffusion_config,
@@ -74,12 +73,6 @@ class TestAttentionSpec:
assert bk["dtype_vo"] == "fp8_e4m3" assert bk["dtype_vo"] == "fp8_e4m3"
assert bk["flashinfer_backend"] == "trtllm-gen" 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): def test_quant_and_skip_softmax_coexist(self):
bk = AttentionSpec( bk = AttentionSpec(
backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5} backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5}
@@ -157,8 +150,8 @@ class TestAttentionSpec:
def test_block_sparse_precision_enum_input_normalized(self): def test_block_sparse_precision_enum_input_normalized(self):
"""RainFusionPrecision enum members must normalize to their string value. """RainFusionPrecision enum members must normalize to their string value.
Regression: ``str(str, Enum)`` later produced ``RainFusionPrecision.MIX`` Regression: the str,Enum members passed the value check but str() later
instead of ``mix``. ``BlockSparseSpec`` stores ``member.value``. produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it).
""" """
from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig
from vllm_omni.diffusion.data import RainFusionPrecision from vllm_omni.diffusion.data import RainFusionPrecision
@@ -516,49 +509,6 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
dense_flash.assert_not_called() dense_flash.assert_not_called()
sdpa.assert_called_once() 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): def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
attention = Attention.__new__(Attention) attention = Attention.__new__(Attention)
attention._scheduler_paged_kv = False attention._scheduler_paged_kv = False
@@ -629,106 +579,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
assert len(calls) == 1 assert len(calls) == 1
assert attention.attn_backend is _DenseBackend assert attention.attn_backend is _DenseBackend
# Auto path records the resolved backend name for Ring/SP; it is not explicit. assert attention.backend_pref is None
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): def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
class _FakeAttentionImpl: class _FakeAttentionImpl:
@@ -763,10 +614,9 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
return _FakeBackend, AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5}) return _FakeBackend, AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})
class _FakeRingParallelAttention: class _FakeRingParallelAttention:
def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False): def __init__(self, sp_group, attn_backend_pref=None):
self.sp_group = sp_group self.sp_group = sp_group
self.attn_backend_pref = attn_backend_pref 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, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl)) monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
@@ -812,134 +662,11 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
assert attn.use_ring is True assert attn.use_ring is True
assert attn.ring_runner is not None assert attn.ring_runner is not None
assert attn.ring_runner.attn_backend_pref == "TRTLLM_ATTN" 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: class TestDiffusionKvCacheQuantization:
@staticmethod @staticmethod
def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True): def _install_attention_init_stubs(monkeypatch):
class _FakeAttentionImpl: class _FakeAttentionImpl:
def __init__(self, **kwargs): def __init__(self, **kwargs):
self.kwargs = kwargs self.kwargs = kwargs
@@ -947,10 +674,6 @@ class TestDiffusionKvCacheQuantization:
def forward(self, query, key, value, attn_metadata=None): def forward(self, query, key, value, attn_metadata=None):
return query return query
@staticmethod
def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
return supports_kv_cache_dtype
class _FakeBackend: class _FakeBackend:
@staticmethod @staticmethod
def get_name() -> str: def get_name() -> str:
@@ -960,11 +683,14 @@ class TestDiffusionKvCacheQuantization:
def get_impl_cls(): def get_impl_cls():
return _FakeAttentionImpl return _FakeAttentionImpl
@staticmethod
def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
return True
class _FakeRingParallelAttention: class _FakeRingParallelAttention:
def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False): def __init__(self, sp_group, attn_backend_pref=None):
self.sp_group = sp_group self.sp_group = sp_group
self.attn_backend_pref = attn_backend_pref self.attn_backend_pref = attn_backend_pref
self.attn_backend_explicit = attn_backend_explicit
monkeypatch.setattr( monkeypatch.setattr(
layer_mod, layer_mod,
@@ -1038,22 +764,3 @@ class TestDiffusionKvCacheQuantization:
causal=False, causal=False,
softmax_scale=1.0, 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,
)
@@ -1,265 +0,0 @@
# 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
+35 -69
View File
@@ -1,84 +1,62 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from contextlib import contextmanager, nullcontext from contextlib import nullcontext
from typing import Any
import pytest import pytest
import torch import torch
from torch.nn.attention import SDPBackend 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.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend, CuDNNAttentionImpl from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch): @pytest.mark.parametrize(
"""Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V.""" ("query_sequence_length", "key_sequence_length", "expected_backend"),
[
(1, 1, SDPBackend.MATH),
(8, 1, SDPBackend.MATH),
(1, 8, SDPBackend.MATH),
(8, 8, SDPBackend.CUDNN_ATTENTION),
],
)
def test_cudnn_routes_single_token_attention_to_math(
monkeypatch, query_sequence_length, key_sequence_length, expected_backend
):
selected_backends = [] selected_backends = []
@contextmanager
def fake_sdpa_kernel(backends): def fake_sdpa_kernel(backends):
selected_backends.append(tuple(backends)) selected_backends.append(backends)
yield return nullcontext()
def fake_sdpa(query, key, value, **kwargs): monkeypatch.setattr(
return query "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
fake_sdpa_kernel,
monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel) )
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) monkeypatch.setattr(
torch.nn.functional,
impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5) "scaled_dot_product_attention",
query = torch.randn(1, 2, 2, 8) lambda query, _key, _value, **_kwargs: query,
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 == [(SDPBackend.MATH,)]
def test_explicit_cudnn_rejects_kv_seq_len_one():
impl = CuDNNAttentionImpl( impl = CuDNNAttentionImpl(
num_heads=2, num_heads=2,
head_size=8, head_size=4,
softmax_scale=0.5, softmax_scale=0.5,
backend_explicit=True,
) )
query = torch.randn(1, 2, 2, 8) query = torch.randn(1, query_sequence_length, 2, 4)
singleton_kv = torch.randn(1, 1, 2, 8) key = torch.randn(1, key_sequence_length, 2, 4)
value = torch.randn_like(key)
with pytest.raises(ValueError, match="explicitly selected.*sequence length 1"): output = impl.forward_cuda(query, key, value)
impl.forward_cuda(query, singleton_kv, singleton_kv)
assert output.shape == query.shape
def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch): assert selected_backends == [[expected_backend]]
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): def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
observed: dict[str, Any] = {} observed = {}
def fake_sdpa(query, key, value, **kwargs): def fake_sdpa(query, key, value, **kwargs):
observed.update(query=query, key=key, value=value, kwargs=kwargs) observed.update(query=query, key=key, value=value, kwargs=kwargs)
@@ -127,15 +105,3 @@ def test_cudnn_rejects_invalid_valid_kv_length():
query, query,
AttentionMetadata(extra={"valid_kv_length": 9}), AttentionMetadata(extra={"valid_kv_length": 9}),
) )
@pytest.mark.parametrize("head_size", [8, 64, 128, 256])
def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size):
assert CuDNNAttentionBackend.supports_head_size(head_size)
assert head_size in CuDNNAttentionBackend.get_supported_head_sizes()
@pytest.mark.parametrize("head_size", [0, 7, 12, 320])
def test_cudnn_backend_rejects_incompatible_head_sizes(head_size):
assert not CuDNNAttentionBackend.supports_head_size(head_size)
assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes()
+1 -203
View File
@@ -1,7 +1,6 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import math
import sys import sys
import types import types
from typing import Any from typing import Any
@@ -9,17 +8,9 @@ from typing import Any
import pytest import pytest
import torch import torch
from vllm_omni.diffusion.attention.backends.abstract import ( from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
AttentionMetadata,
PackedPaddingMetadata,
VideoTokenLayout,
VideoTokenSpan,
)
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import ( from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
FastVideoVSABackend,
FastVideoVSAImpl, FastVideoVSAImpl,
_build_h3_block_map,
_get_h3_tile_metadata,
) )
from vllm_omni.diffusion.attention.backends.registry import ( from vllm_omni.diffusion.attention.backends.registry import (
DiffusionAttentionBackendEnum, DiffusionAttentionBackendEnum,
@@ -32,12 +23,6 @@ def test_fastvideo_vsa_backend_is_registered():
assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend") assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
def test_fastvideo_vsa_reports_missing_optional_kernel(monkeypatch):
monkeypatch.setattr("importlib.util.find_spec", lambda name: None)
with pytest.raises(ImportError, match="fastvideo-kernel"):
FastVideoVSABackend.validate_available()
def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch): def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
calls: dict[str, Any] = {} calls: dict[str, Any] = {}
fake_module = types.ModuleType("fastvideo_kernel") fake_module = types.ModuleType("fastvideo_kernel")
@@ -216,190 +201,3 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
# CPU/float32 is rejected later, but k=N itself must not trigger fallback. # CPU/float32 is rejected later, but k=N itself must not trigger fallback.
assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported" assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2" assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
extra: dict[str, Any] = {
"vsa_h3_prefix_segments": prefix_segments,
}
if gate is not None:
extra["gate_compress"] = gate
return AttentionMetadata(
packed_padding=packed_padding,
video_layout=VideoTokenLayout(
used_len=sum(prefix_segments) + math.prod(video_shape),
video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
),
extra=extra,
)
def _fake_block_sparse_kernel(monkeypatch, calls):
"""Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
calls["block_map"] = block_map.clone()
calls["variable_block_sizes"] = variable_block_sizes.tolist()
calls["q_shape"] = tuple(q.shape)
blocks = block_map.shape[-1]
assert blocks * 64 == q.shape[2]
# Expand the per-block map to a per-token mask so the stand-in honours
# the sparsity contract the real kernel implements, and drop the rows
# each partial block pads with, which the real kernel skips through
# variable_block_sizes.
token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
token_mask = token_mask & within.reshape(1, 1, 1, -1)
out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
return out.to(q.dtype), None
module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
setattr(module, "block_sparse_attn", block_sparse_attn)
monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
def _h3_impl(**backend_kwargs):
kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
kwargs.update(backend_kwargs)
return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
calls: dict[str, Any] = {}
_fake_block_sparse_kernel(monkeypatch, calls)
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5, 70), (2, 4, 4)
seq_len = sum(prefix_segments) + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
assert output.shape == query.shape
# topk covers every video block, so the block map is dense and the tiled
# kernel must reproduce plain attention over the same rows.
reference = torch.nn.functional.scaled_dot_product_attention(
*(x.float().transpose(1, 2) for x in (query, query, query))
).transpose(1, 2)
torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
# 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
assert calls["variable_block_sizes"] == [5, 64, 6, 32]
def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
calls: dict[str, Any] = {}
_fake_block_sparse_kernel(monkeypatch, calls)
impl = _h3_impl(topk=1)
prefix_segments, video_shape = (5,), (8, 4, 4)
seq_len = 5 + 8 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
# 1 prefix block + 2 video blocks. The odd count is padded to an even one
# for the paired-CTA kernel contract, and the Triton route slices that
# transport-only partner off again before launching.
block_map = calls["block_map"]
assert block_map.shape[-1] == 3
assert calls["variable_block_sizes"] == [5, 64, 64]
assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
assert block_map[..., :1].all(), "prefix keys are exempt from selection"
assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
# The same odd geometry must still untile back onto the original rows.
dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
reference = torch.nn.functional.scaled_dot_product_attention(
*(x.float().transpose(1, 2) for x in (query, query, query))
).transpose(1, 2)
torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
_fake_block_sparse_kernel(monkeypatch, {})
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5,), (2, 4, 4)
seq_len = 5 + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
zero_gate = torch.zeros_like(query)
with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
live_gate = torch.full_like(query, 0.5)
with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
assert not torch.allclose(with_live_gate.float(), without_gate.float())
def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
_fake_block_sparse_kernel(monkeypatch, {})
impl = _h3_impl(topk=64)
prefix_segments, video_shape = (5,), (2, 4, 4)
valid = 5 + 2 * 4 * 4
torch.manual_seed(0)
query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
padding = PackedPaddingMetadata(
q_length=valid,
kv_length=valid,
cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
)
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
assert output.shape == query.shape
assert torch.count_nonzero(output[:, valid:]) == 0
def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
# The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
# building the padding mask. Every fallback out of the VSA path must honour
# that contract itself; SDPA reads attn_mask and nothing else.
impl = FastVideoVSAImpl(
num_heads=2,
head_size=8,
softmax_scale=8**-0.5,
backend_kwargs={"topk": 1, "min_seq_len": 4096},
)
valid = 40
torch.manual_seed(0)
query = torch.randn(1, 64, 2, 8)
padding = PackedPaddingMetadata(
q_length=valid,
kv_length=valid,
cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
)
baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
perturbed_input = query.clone()
perturbed_input[:, valid:] += 100.0
perturbed = impl.forward_cuda(
perturbed_input, perturbed_input, perturbed_input, AttentionMetadata(packed_padding=padding, extra={})
)
assert baseline.shape == query.shape
torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
assert torch.count_nonzero(baseline[:, valid:]) == 0
def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
(5, 70, 9), (5, 6, 6), torch.device("cpu")
)
assert sizes[:4].tolist() == [5, 64, 6, 9]
assert prefix_blocks == 4
assert video_blocks == 8
assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
assert block_map[:, :, :2].all()
assert block_map[..., :2].all()
assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()
@@ -1,112 +0,0 @@
# 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
+17 -66
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys import sys
import types import types
@@ -9,13 +9,18 @@ import torch
from tests.helpers.mark import hardware_test from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata 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.diffusion.attention.backends.sdpa import SDPAImpl
from vllm_omni.platforms import current_omni_platform from vllm_omni.platforms import current_omni_platform
@pytest.mark.core_model @pytest.mark.core_model
@pytest.mark.cpu @pytest.mark.cpu
def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch): def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
"""Test that when kernels package is not available, the platform fallback logic
routes to local FLASH_ATTN.
"""
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
# Temporarily hide/remove kernels module if present # Temporarily hide/remove kernels module if present
@@ -33,73 +38,19 @@ def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.Mo
) )
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True}) monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
with pytest.raises(ImportError, match="explicitly selected"): # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
with pytest.raises(ImportError, match="explicitly selected"): assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
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") kernels_module = types.ModuleType("kernels")
monkeypatch.setitem(sys.modules, "kernels", kernels_module) monkeypatch.setitem(sys.modules, "kernels", kernels_module)
with pytest.raises(ValueError, match="require.*Hopper GPU"): backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
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 @pytest.mark.core_model
@@ -1,199 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import importlib
import sys
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.diffusion.attention.backends.ring import ring_globals, ring_kernels
from vllm_omni.diffusion.attention.parallel import factory, ring
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_source_build_fa3_is_hopper_only(monkeypatch):
fake_interface = SimpleNamespace(
_flash_attn_forward=lambda *args, **kwargs: None,
flash_attn_func=lambda *args, **kwargs: None,
)
try:
with monkeypatch.context() as context:
context.setitem(sys.modules, "flash_attn_interface", fake_interface)
reloaded = importlib.reload(ring_globals)
assert reloaded.HAS_FA3
assert reloaded.FA3_SUPPORTED_CUDA_MAJORS == frozenset({9})
finally:
importlib.reload(ring_globals)
@pytest.mark.parametrize(
("capability", "expected"),
[
((8, 0), True),
((8, 9), True),
((9, 0), True),
((10, 0), False),
((10, 3), False),
((12, 0), False),
],
)
def test_can_use_fa3_checks_device_arch(monkeypatch, capability, expected):
monkeypatch.setattr(ring, "HAS_FA3", True)
monkeypatch.setattr(ring, "FA3_SUPPORTED_CUDA_MAJORS", frozenset({8, 9}))
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: capability)
assert ring._can_use_fa3(torch.device("cuda")) is expected
def test_can_use_fa3_preserves_unknown_extension_contract(monkeypatch):
monkeypatch.setattr(ring, "HAS_FA3", True)
monkeypatch.setattr(ring, "FA3_SUPPORTED_CUDA_MAJORS", None)
assert ring._can_use_fa3(torch.device("cuda"))
@pytest.mark.parametrize(
("capability", "expected"),
[
((8, 0), True),
((9, 0), True),
((10, 0), False),
((10, 3), False),
((12, 0), False),
],
)
def test_can_use_fa2_checks_device_arch(monkeypatch, capability, expected):
monkeypatch.setattr(ring, "HAS_FLASH_ATTN", True)
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: capability)
assert ring._can_use_fa2(torch.device("cuda")) is expected
def test_ring_falls_back_to_sdpa_when_fa3_extension_has_no_device_kernel(monkeypatch):
monkeypatch.setattr(ring, "HAS_FA3", True)
monkeypatch.setattr(ring, "HAS_FA4", False)
monkeypatch.setattr(ring, "HAS_FLASH_ATTN", False)
monkeypatch.setattr(ring, "HAS_AITER", False)
monkeypatch.setattr(ring, "_can_use_fa3", lambda _device: False)
expected = torch.randn(1, 2, 1, 8)
def fake_ring_sdpa(*args, **kwargs):
return expected
module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
monkeypatch.setitem(sys.modules, module_name, SimpleNamespace(ring_pytorch_attn_func=fake_ring_sdpa))
strategy = ring.RingParallelAttention(SimpleNamespace(ring_group=object()))
query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
actual = strategy.run_attention(query, query, query, attn_metadata=None)
assert actual is expected
def test_explicit_ring_backend_does_not_fallback_to_sdpa(monkeypatch):
monkeypatch.setattr(ring, "HAS_FA3", True)
monkeypatch.setattr(ring, "HAS_FA4", False)
monkeypatch.setattr(ring, "HAS_FLASH_ATTN", False)
monkeypatch.setattr(ring, "HAS_AITER", False)
monkeypatch.setattr(ring, "_can_use_fa3", lambda _device: False)
strategy = ring.RingParallelAttention(
SimpleNamespace(ring_group=object()),
attn_backend_pref="FLASH_ATTN",
attn_backend_explicit=True,
)
query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
with pytest.raises(RuntimeError, match="explicitly selected"):
strategy.run_attention(query, query, query, attn_metadata=None)
@pytest.mark.parametrize(
"backend_pref",
[
"CUDNN_ATTN",
"FLASH_ATTN_HUB",
"FLASH_ATTN_3_HUB",
],
)
def test_explicit_non_ring_backend_is_rejected(monkeypatch, backend_pref):
# Local FA is available so a missing rejection would silently run FA4/FA3/FA2
# instead of the requested Hub/cuDNN backend.
monkeypatch.setattr(ring, "HAS_FA4", True)
monkeypatch.setattr(ring, "_can_use_fa4", lambda _device: True)
module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
monkeypatch.setitem(
sys.modules,
module_name,
SimpleNamespace(ring_pytorch_attn_func=lambda *args, **kwargs: pytest.fail("unexpected SDPA fallback")),
)
monkeypatch.setitem(
sys.modules,
"vllm_omni.diffusion.attention.backends.ring_flash_attn",
SimpleNamespace(ring_flash_attn_func=lambda *args, **kwargs: pytest.fail("unexpected local FA ring")),
)
strategy = ring.RingParallelAttention(
SimpleNamespace(ring_group=object()),
attn_backend_pref=backend_pref,
attn_backend_explicit=True,
)
query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
with pytest.raises(ValueError, match="ring sequence parallelism has no implementation"):
strategy.run_attention(query, query, query, attn_metadata=None)
def test_explicit_flash_ring_rejects_float32_instead_of_falling_back(monkeypatch):
module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
monkeypatch.setitem(
sys.modules,
module_name,
SimpleNamespace(ring_pytorch_attn_func=lambda *args, **kwargs: pytest.fail("unexpected SDPA fallback")),
)
strategy = ring.RingParallelAttention(
SimpleNamespace(ring_group=object()),
attn_backend_pref="FLASH_ATTN",
attn_backend_explicit=True,
)
query = torch.randn(1, 2, 1, 8, dtype=torch.float32)
with pytest.raises(ValueError, match="does not support float32"):
strategy.run_attention(query, query, query, attn_metadata=None)
def test_pytorch_ring_flash_op_rejects_float32():
query = torch.randn(1, 2, 1, 8, dtype=torch.float32)
with pytest.raises(ValueError, match="does not support float32"):
ring_kernels.pytorch_attn_forward(query, query, query, op_type="flash")
def test_configured_sp_does_not_fallback_when_group_is_unavailable(monkeypatch):
parallel_config = SimpleNamespace(ulysses_degree=2, ring_degree=1, allgather_degree=1)
forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(parallel_config=parallel_config))
monkeypatch.setattr(factory, "is_forward_context_available", lambda: True)
monkeypatch.setattr(factory, "get_forward_context", lambda: forward_context)
monkeypatch.setattr(
factory,
"get_sp_group",
lambda: (_ for _ in ()).throw(RuntimeError("SP group is not initialized")),
)
with pytest.raises(RuntimeError, match="SP is configured"):
factory.build_parallel_attention_strategy(scatter_idx=2, gather_idx=1, use_sync=False)
def test_configured_sp_rejects_single_rank_group(monkeypatch):
parallel_config = SimpleNamespace(ulysses_degree=2, ring_degree=1, allgather_degree=1)
forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(parallel_config=parallel_config))
monkeypatch.setattr(factory, "is_forward_context_available", lambda: True)
monkeypatch.setattr(factory, "get_forward_context", lambda: forward_context)
monkeypatch.setattr(factory, "get_sp_group", lambda: SimpleNamespace())
monkeypatch.setattr(factory, "get_sequence_parallel_world_size", lambda: 1)
with pytest.raises(RuntimeError, match="world size is not greater than one"):
factory.build_parallel_attention_strategy(scatter_idx=2, gather_idx=1, use_sync=False)
@@ -72,25 +72,3 @@ def test_ring_prefers_fa4_on_blackwell(monkeypatch):
assert strategy.run_attention(query, "key", "value", None) == "output" assert strategy.run_attention(query, "key", "value", None) == "output"
assert captured["attn_type"] is AttnType.FA4 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)
@@ -1,92 +0,0 @@
# 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))))
@@ -1,36 +0,0 @@
# 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)
+23 -96
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import importlib import importlib
import sys import sys
@@ -9,8 +9,6 @@ import pytest
import torch import torch
from vllm.platforms import current_platform 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] pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -19,7 +17,7 @@ SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"
def load_sage_attn3_module(monkeypatch: pytest.MonkeyPatch, kernel_impl): def load_sage_attn3_module(monkeypatch: pytest.MonkeyPatch, kernel_impl):
fake_module = types.ModuleType("sageattn3") fake_module = types.ModuleType("sageattn3")
setattr(fake_module, "sageattn3_blackwell", kernel_impl) fake_module.sageattn3_blackwell = kernel_impl
monkeypatch.setitem(sys.modules, "sageattn3", fake_module) monkeypatch.setitem(sys.modules, "sageattn3", fake_module)
sys.modules.pop(SAGE_ATTN3_MODULE, None) sys.modules.pop(SAGE_ATTN3_MODULE, None)
return importlib.import_module(SAGE_ATTN3_MODULE) return importlib.import_module(SAGE_ATTN3_MODULE)
@@ -53,11 +51,21 @@ def test_sage_attn3_forward_uses_blackwell_layout(monkeypatch: pytest.MonkeyPatc
assert torch.allclose(output, expected) assert torch.allclose(output, expected)
def test_sage_attn3_rejects_gqa_instead_of_falling_back(monkeypatch: pytest.MonkeyPatch): def test_sage_attn3_falls_back_to_sdpa_for_gqa(monkeypatch: pytest.MonkeyPatch):
def fake_kernel(*args, **kwargs): def fake_kernel(*args, **kwargs):
raise AssertionError("sageattn3_blackwell should not be used for GQA") raise AssertionError("sageattn3_blackwell should not be used for GQA")
backend_module = load_sage_attn3_module(monkeypatch, fake_kernel) 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( impl = backend_module.SageAttention3Impl(
num_heads=4, num_heads=4,
head_size=64, head_size=64,
@@ -69,38 +77,13 @@ def test_sage_attn3_rejects_gqa_instead_of_falling_back(monkeypatch: pytest.Monk
key = torch.randn(2, 8, 2, 64) key = torch.randn(2, 8, 2, 64)
value = torch.randn(2, 8, 2, 64) value = torch.randn(2, 8, 2, 64)
with pytest.raises(NotImplementedError, match="does not support GQA/MQA"): output = impl.forward_cuda(query, key, value)
impl.forward_cuda(query, key, value)
assert sdpa_calls["query_shape"] == (2, 4, 8, 64)
def test_sage_attn3_rejects_mask_instead_of_ignoring_it(monkeypatch: pytest.MonkeyPatch): assert sdpa_calls["key_shape"] == (2, 2, 8, 64)
def fake_kernel(*args, **kwargs): assert sdpa_calls["enable_gqa"] is True
raise AssertionError("sageattn3_blackwell should not run with an unsupported mask") expected = (query.permute(0, 2, 1, 3) + 1).permute(0, 2, 1, 3)
assert torch.allclose(output, expected)
backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
impl = backend_module.SageAttention3Impl(
num_heads=4,
head_size=64,
softmax_scale=1.0 / 8.0,
causal=False,
)
query = torch.randn(2, 8, 4, 64)
metadata = AttentionMetadata(attn_mask=torch.ones(2, 8, dtype=torch.bool))
with pytest.raises(ValueError, match="does not support attn_mask"):
impl.forward_cuda(query, query, query, metadata)
def test_sage_attn3_rejects_custom_softmax_scale(monkeypatch: pytest.MonkeyPatch):
backend_module = load_sage_attn3_module(monkeypatch, lambda *args, **kwargs: None)
with pytest.raises(ValueError, match="does not expose a custom softmax scale"):
backend_module.SageAttention3Impl(
num_heads=4,
head_size=64,
softmax_scale=1.0,
causal=False,
)
@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform") @pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
@@ -132,39 +115,10 @@ 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") @pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
@pytest.mark.parametrize("head_size", [32, 320]) def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch: pytest.MonkeyPatch):
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.platforms.interface import DeviceCapability
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
from vllm_omni.diffusion.envs import PACKAGES_CHECKER from vllm_omni.diffusion.envs import PACKAGES_CHECKER
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
@@ -175,33 +129,6 @@ def test_cuda_platform_rejects_explicit_sage_attn3_on_unsupported_gpu(monkeypatc
) )
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False}) monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
with pytest.raises(ValueError, match="explicitly selected"): backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
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)
+4 -40
View File
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for diffusion attention backend selection.""" """Unit tests for diffusion attention backend selection."""
from types import SimpleNamespace from types import SimpleNamespace
@@ -100,12 +100,9 @@ def test_platform_default_used_without_resolved_spec(monkeypatch, attention_conf
def test_backend_class_resolution_is_cached(monkeypatch): def test_backend_class_resolution_is_cached(monkeypatch):
platform_calls = [] platform_calls = []
load_calls = [] load_calls = []
fake_platform = SimpleNamespace(
def fake_get_cls(**kwargs): get_diffusion_attn_backend_cls=lambda **kwargs: (platform_calls.append(kwargs) or "fake.module.Backend")
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) monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform)
@@ -129,39 +126,6 @@ def test_backend_class_resolution_is_cached(monkeypatch):
assert load_calls == ["fake.module.Backend"] 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(): def test_load_backend_cls_reports_missing_module():
with pytest.raises(ImportError, match="Failed to import module missing_attention_backend"): with pytest.raises(ImportError, match="Failed to import module missing_attention_backend"):
selector._load_backend_cls("missing_attention_backend.Backend") selector._load_backend_cls("missing_attention_backend.Backend")
+38 -448
View File
@@ -29,343 +29,31 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.mark.core_model @pytest.mark.core_model
@pytest.mark.cpu @pytest.mark.cpu
def test_uaa_gqa_head_padding_preserves_the_query_to_kv_ratio( def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
monkeypatch: pytest.MonkeyPatch, monkeypatch.setattr(
) -> None: "vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
"""Q must be padded to a multiple of the *KV* padding, not of world_size. lambda default: "advanced_uaa",
Rounding each up independently is the obvious implementation and is wrong:
28Q/7KV at SP2 would become 28/8, i.e. 14/4 per rank, whose 3.5 ratio is not
a valid GQA shape. Deriving Q from the padded KV count gives 32/8 -> 16/4.
"""
from vllm_omni.diffusion.attention.parallel import ulysses
class FakeGroup:
ulysses_group = object()
ulysses_world_size = 2
ulysses_rank = 0
ring_world_size = 1
observed = []
def fake_all_to_all(pg, tensor, **kwargs):
observed.append(kwargs["padded_head_cnt"])
return tensor, tensor.shape[2]
monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
strategy = ulysses.UlyssesParallelAttention(
FakeGroup(),
scatter_idx=2,
gather_idx=1,
use_sync=False,
) )
with set_forward_context():
strategy.pre_attention(
torch.zeros(1, 3, 28, 4),
torch.zeros(1, 3, 7, 4),
torch.zeros(1, 3, 7, 4),
None,
)
q_padded, k_padded, v_padded = observed
assert (q_padded, k_padded, v_padded) == (32, 8, 8)
# The per-rank shape must still be a valid GQA shape.
assert (q_padded // 2) % (k_padded // 2) == 0
@pytest.mark.core_model
@pytest.mark.cpu
def test_uaa_rejects_head_counts_that_are_not_a_gqa_shape(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.attention.parallel import ulysses
class FakeGroup:
ulysses_group = object()
ulysses_world_size = 2
ulysses_rank = 0
ring_world_size = 1
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
strategy = ulysses.UlyssesParallelAttention(
FakeGroup(),
scatter_idx=2,
gather_idx=1,
use_sync=False,
)
with set_forward_context(), pytest.raises(ValueError, match="multiple of KV heads"):
strategy.pre_attention(
torch.zeros(1, 3, 10, 4),
torch.zeros(1, 3, 4, 4),
torch.zeros(1, 3, 4, 4),
None,
)
@pytest.mark.core_model
@pytest.mark.cpu
@pytest.mark.parametrize(
"q_heads, kv_heads, world_size, expected_warn",
[
# 28Q/7KV @ U=2 -> pad to 32Q, 32/28 = 1.14x, below threshold, no warn.
(28, 7, 2, False),
# 32Q/1KV @ U=8 -> pad KV to 8 -> pad Q to 256, 256/32 = 8x, warn.
(32, 1, 8, True),
# 16Q/2KV @ U=4 -> pad KV to 4 -> pad Q to 32, 32/16 = 2x, warn.
(16, 2, 4, True),
],
)
def test_uaa_padding_ratio_warning_fires_when_blowup_is_large(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
q_heads: int,
kv_heads: int,
world_size: int,
expected_warn: bool,
) -> None:
"""advanced_uaa must warn once when GQA padding materially inflates Q heads.
MQA-style shapes at high ulysses_degree pad Q by U x, which grows attention
FLOPs and temporary VRAM by the same factor. We do not fall back (there is
no cheap alternative -- Ulysses needs KV divisible by U for the head-dim
split), but we surface a one-shot warning so the user can pick a
ulysses_degree that divides KV_heads if the overhead is unacceptable.
"""
from vllm_omni.diffusion.attention.parallel import ulysses
# Reset the warn-once dedup set so parametrized cases do not shadow each other.
ulysses._uaa_pad_ratio_warned.clear()
class FakeGroup:
ulysses_group = object()
ulysses_world_size = world_size
ulysses_rank = 0
ring_world_size = 1
def fake_all_to_all(pg, tensor, **kwargs):
return tensor, tensor.shape[2]
monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *a, **kw: [3] * world_size)
strategy = ulysses.UlyssesParallelAttention(
FakeGroup(),
scatter_idx=2,
gather_idx=1,
use_sync=False,
)
caplog.set_level("WARNING", logger="vllm_omni.diffusion.attention.parallel.ulysses")
with set_forward_context():
strategy.pre_attention(
torch.zeros(1, 3, q_heads, 4),
torch.zeros(1, 3, kv_heads, 4),
torch.zeros(1, 3, kv_heads, 4),
None,
)
warn_msgs = [r.message for r in caplog.records if "GQA padding inflates" in r.message]
if expected_warn:
assert len(warn_msgs) == 1, warn_msgs
# Second call with the same shape must not re-warn (warn-once dedup).
with set_forward_context():
strategy.pre_attention(
torch.zeros(1, 3, q_heads, 4),
torch.zeros(1, 3, kv_heads, 4),
torch.zeros(1, 3, kv_heads, 4),
None,
)
warn_msgs = [r.message for r in caplog.records if "GQA padding inflates" in r.message]
assert len(warn_msgs) == 1, "warn-once dedup should suppress the second call"
else:
assert warn_msgs == [], f"unexpected warning for {q_heads}Q/{kv_heads}KV @ U={world_size}: {warn_msgs}"
@pytest.mark.core_model
@pytest.mark.cpu
def test_uaa_joint_gqa_head_padding_preserves_the_query_to_kv_ratio(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Joint Q/K/V must also preserve the GQA ratio through Ulysses padding.
MMDiT-style models (SD3, FLUX, Qwen-Image, Hunyuan-Video, ...) pass a
separate text stream via ``AttentionMetadata.joint_*``. Before the fix the
joint block padded K/V by Q's amount and integer-divided K/V by world_size,
which for 28Q/7KV at SP=2 produces mismatched joint vs main head counts and
the downstream ring-attention concat blows up. Post-fix the joint block
mirrors the main-path rule: pad KV to a world_size multiple, derive Q as
``padded_kv * (Q // KV)``.
"""
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.parallel import ulysses
class FakeGroup:
ulysses_group = object()
ulysses_world_size = 2
ulysses_rank = 0
ring_world_size = 1
def fake_all_to_all(pg, tensor, **kwargs):
# Simulate the head-dim sharding so the post-a2a joint concat has
# matching per-rank head counts on Q and joint_Q.
heads_per_rank = kwargs["padded_head_cnt"] // 2
return tensor[:, :, :heads_per_rank, :].contiguous(), tensor.shape[2]
monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
strategy = ulysses.UlyssesParallelAttention(
FakeGroup(),
scatter_idx=2,
gather_idx=1,
use_sync=False,
)
attn_metadata = AttentionMetadata(
joint_query=torch.zeros(1, 5, 28, 4),
joint_key=torch.zeros(1, 5, 7, 4),
joint_value=torch.zeros(1, 5, 7, 4),
)
with set_forward_context():
strategy.pre_attention(
torch.zeros(1, 3, 28, 4),
torch.zeros(1, 3, 7, 4),
torch.zeros(1, 3, 7, 4),
attn_metadata,
)
# 28Q/7KV at SP=2 -> pad to 32/8 -> per-rank slice is 16Q/4KV.
assert attn_metadata.joint_key.shape[-2] == 4
assert attn_metadata.joint_value.shape[-2] == 4
# Per-rank joint shape must remain a valid GQA shape (Q multiple of KV).
per_rank_joint_q = 32 // 2
per_rank_joint_kv = attn_metadata.joint_key.shape[-2]
assert per_rank_joint_q % per_rank_joint_kv == 0
@pytest.mark.core_model
@pytest.mark.cpu
def test_uaa_rejects_joint_head_counts_that_are_not_a_gqa_shape(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.parallel import ulysses
class FakeGroup:
ulysses_group = object()
ulysses_world_size = 2
ulysses_rank = 0
ring_world_size = 1
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
strategy = ulysses.UlyssesParallelAttention(
FakeGroup(),
scatter_idx=2,
gather_idx=1,
use_sync=False,
)
attn_metadata = AttentionMetadata(
joint_query=torch.zeros(1, 5, 10, 4),
joint_key=torch.zeros(1, 5, 4, 4),
joint_value=torch.zeros(1, 5, 4, 4),
)
with set_forward_context(), pytest.raises(ValueError, match="multiple of joint KV heads"):
strategy.pre_attention(
torch.zeros(1, 3, 8, 4),
torch.zeros(1, 3, 4, 4),
torch.zeros(1, 3, 4, 4),
attn_metadata,
)
@pytest.mark.core_model
@pytest.mark.cpu
def test_advanced_uaa_hybrid_rejects_non_gqa_shapes(monkeypatch: pytest.MonkeyPatch) -> None:
"""Hybrid Ulysses+Ring still rejects shapes that are not a GQA shape.
Valid GQA shapes are safe under advanced_uaa head padding (the derived-Q
rule keeps real Q heads paired with real K/V heads after the split, see
pre_attention), but a query count that is not a multiple of the K/V count
cannot be padded while preserving the ratio, so it must be rejected.
"""
from vllm_omni.diffusion.attention.parallel import ulysses
sp_group = SimpleNamespace( sp_group = SimpleNamespace(
ulysses_group=None, ulysses_group=None,
ulysses_world_size=2, ulysses_world_size=2,
ulysses_rank=0, ulysses_rank=0,
ring_world_size=2, ring_world_size=2,
ring_group=object(),
) )
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
strategy = UlyssesParallelAttention( strategy = UlyssesParallelAttention(
sp_group=sp_group, sp_group=sp_group,
scatter_idx=2, scatter_idx=2,
gather_idx=1, gather_idx=1,
use_sync=False, use_sync=False,
) )
query = torch.zeros(1, 2, 10, 4) query = torch.zeros(1, 2, 8, 4)
key = torch.zeros(1, 2, 4, 4) key = torch.zeros(1, 2, 1, 4)
value = torch.zeros_like(key) value = torch.zeros_like(key)
with pytest.raises(ValueError, match="multiple of KV heads"): with pytest.raises(ValueError, match="does not support K/V head padding"):
strategy.pre_attention(query, key, value, None) strategy.pre_attention(query, key, value, None)
def test_strict_ulysses_reshards_vsa_gate_with_qkv(monkeypatch) -> None:
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
monkeypatch.setattr(
"vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
lambda default: "strict",
)
calls: list[torch.Tensor] = []
def fake_all_to_all(pg, tensor, scatter_idx, gather_idx, use_sync):
del pg, scatter_idx, gather_idx, use_sync
calls.append(tensor)
return tensor.reshape(tensor.shape[0], tensor.shape[1] * 2, tensor.shape[2] // 2, tensor.shape[3])
monkeypatch.setattr(
"vllm_omni.diffusion.attention.parallel.ulysses.SeqAllToAll4D.apply",
fake_all_to_all,
)
strategy = UlyssesParallelAttention(
sp_group=SimpleNamespace(
ulysses_group=object(),
ulysses_world_size=2,
ulysses_rank=0,
ring_world_size=1,
),
scatter_idx=2,
gather_idx=1,
use_sync=False,
)
query = torch.arange(24).reshape(1, 3, 4, 2)
key = query + 10
value = query + 20
gate = query + 100
metadata = AttentionMetadata(extra={"gate_compress": gate})
query_out, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
assert all(actual is expected for actual, expected in zip(calls, (query, key, value, gate), strict=True))
assert query_out.shape == (1, 6, 2, 2)
assert metadata_out is metadata
assert metadata.extra["gate_compress"].shape == query_out.shape
torch.testing.assert_close(metadata.extra["gate_compress"], gate.reshape_as(query_out))
def _run_attention_case( def _run_attention_case(
local_rank: int, local_rank: int,
world_size: int, world_size: int,
@@ -379,9 +67,6 @@ def _run_attention_case(
ring_degree: int = 1, ring_degree: int = 1,
split_sizes: list[int] | None = None, split_sizes: list[int] | None = None,
sdp_kernel_mode: str = "math", sdp_kernel_mode: str = "math",
num_kv_heads: int | None = None,
has_joint: bool = False,
joint_strategy: str = "front",
) -> None: ) -> None:
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device) current_omni_platform.set_device(device)
@@ -417,32 +102,14 @@ def _run_attention_case(
head_size=head_size, head_size=head_size,
causal=False, causal=False,
softmax_scale=1.0 / (head_size**0.5), softmax_scale=1.0 / (head_size**0.5),
num_kv_heads=num_kv_heads,
).to(device=device, dtype=torch.float32) ).to(device=device, dtype=torch.float32)
with np.load(input_file, allow_pickle=False) as payload: with np.load(input_file, allow_pickle=False) as payload:
q_full = torch.from_numpy(payload["q"]).to(device=device) q_full = torch.from_numpy(payload["q"]).to(device=device)
k_full = torch.from_numpy(payload["k"]).to(device=device) k_full = torch.from_numpy(payload["k"]).to(device=device)
v_full = torch.from_numpy(payload["v"]).to(device=device) v_full = torch.from_numpy(payload["v"]).to(device=device)
if has_joint:
joint_q = torch.from_numpy(payload["joint_q"]).to(device=device)
joint_k = torch.from_numpy(payload["joint_k"]).to(device=device)
joint_v = torch.from_numpy(payload["joint_v"]).to(device=device)
attn_metadata = None
if world_size == 1: if world_size == 1:
if has_joint:
# Baseline: SP is inactive so pre_attention won't concat joint for us.
# Build the full sequence directly and skip the joint metadata.
if joint_strategy == "front":
q = torch.cat([joint_q, q_full], dim=1)
k = torch.cat([joint_k, k_full], dim=1)
v = torch.cat([joint_v, v_full], dim=1)
else:
q = torch.cat([q_full, joint_q], dim=1)
k = torch.cat([k_full, joint_k], dim=1)
v = torch.cat([v_full, joint_v], dim=1)
else:
q, k, v = q_full, k_full, v_full q, k, v = q_full, k_full, v_full
else: else:
if split_sizes is None: if split_sizes is None:
@@ -468,16 +135,6 @@ def _run_attention_case(
# shard manually for testing. # shard manually for testing.
get_forward_context()._sp_shard_depth = 1 get_forward_context()._sp_shard_depth = 1
if has_joint:
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
attn_metadata = AttentionMetadata(
joint_query=joint_q,
joint_key=joint_k,
joint_value=joint_v,
joint_strategy=joint_strategy,
)
if sdp_kernel_mode == "math": if sdp_kernel_mode == "math":
sdp_ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True) sdp_ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
elif sdp_kernel_mode == "mem_efficient": elif sdp_kernel_mode == "mem_efficient":
@@ -486,47 +143,25 @@ def _run_attention_case(
raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}") raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}")
with sdp_ctx: with sdp_ctx:
out_local = attn(q, k, v, attn_metadata=attn_metadata).contiguous() out_local = attn(q, k, v).contiguous()
if world_size == 1: if world_size == 1:
out_full = out_local out_full = out_local
else: else:
if has_joint: local_len = torch.tensor([out_local.shape[1]], device=device, dtype=torch.int64)
joint_len = int(joint_q.shape[1])
if joint_strategy == "front":
joint_out_local = out_local[:, :joint_len]
img_out_local = out_local[:, joint_len:]
else:
img_out_local = out_local[:, : out_local.shape[1] - joint_len]
joint_out_local = out_local[:, -joint_len:]
else:
joint_out_local = None
img_out_local = out_local
local_len = torch.tensor([img_out_local.shape[1]], device=device, dtype=torch.int64)
gathered_lens = [torch.empty_like(local_len) for _ in range(world_size)] gathered_lens = [torch.empty_like(local_len) for _ in range(world_size)]
torch.distributed.all_gather(gathered_lens, local_len) torch.distributed.all_gather(gathered_lens, local_len)
lens = [int(t.item()) for t in gathered_lens] lens = [int(t.item()) for t in gathered_lens]
max_len = max(lens) max_len = max(lens)
if img_out_local.shape[1] < max_len: if out_local.shape[1] < max_len:
pad = max_len - img_out_local.shape[1] pad = max_len - out_local.shape[1]
img_out_local = torch.nn.functional.pad(img_out_local, (0, 0, 0, 0, 0, pad)).contiguous() out_local = torch.nn.functional.pad(out_local, (0, 0, 0, 0, 0, pad)).contiguous()
else:
img_out_local = img_out_local.contiguous()
gathered = [torch.empty_like(img_out_local) for _ in range(world_size)] gathered = [torch.empty_like(out_local) for _ in range(world_size)]
torch.distributed.all_gather(gathered, img_out_local) torch.distributed.all_gather(gathered, out_local)
if local_rank == 0: if local_rank == 0:
img_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1) out_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1)
if has_joint:
out_full = (
torch.cat([joint_out_local, img_full], dim=1)
if joint_strategy == "front"
else torch.cat([img_full, joint_out_local], dim=1)
)
else:
out_full = img_full
else: else:
out_full = None out_full = None
@@ -537,43 +172,20 @@ def _run_attention_case(
@pytest.mark.parametrize( @pytest.mark.parametrize(
"sp_world_size,seq_len,joint_len,num_heads,num_kv_heads", "sp_world_size,seq_len,num_heads",
[ [
# MHA (num_kv_heads=None, joint_len=None) (2, 6, 3), # head_cnt not divisible by P=2
(2, 6, None, 3, None), # head_cnt not divisible by P=2 (2, 5, 4), # seq_len not divisible by P=2
(2, 5, None, 4, None), # seq_len not divisible by P=2 (4, 9, 30), # Z-Image-like: head_cnt not divisible by P=4
(4, 9, None, 30, None), # Z-Image-like: head_cnt not divisible by P=4 (4, 10, 8), # seq_len not divisible by P=4
(4, 10, None, 8, None), # seq_len not divisible by P=4
# GQA (joint_len=None)
(2, 6, None, 28, 7), # BOOGU-like GQA: neither Q nor KV divisible by P=2
(2, 5, None, 6, 3), # KV divisible by P, Q is not
(4, 8, None, 12, 3), # KV not divisible by P=4
# Joint GQA (MMDiT-style text stream via AttentionMetadata.joint_*)
(2, 6, 5, 28, 7), # BOOGU-like GQA joint: KV=7 not divisible by P=2
(2, 5, 4, 6, 3), # KV divisible by P, Q is not
], ],
) )
def test_ulysses_uaa_matches_baseline( def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_heads: int) -> None:
sp_world_size: int,
seq_len: int,
joint_len: int | None,
num_heads: int,
num_kv_heads: int | None,
) -> None:
"""End-to-end SP-vs-single-GPU parity under UAA.
Covers MHA, GQA and joint (MMDiT-style) shapes. The joint variant catches
the pre-fix bug where the joint block padded K/V by Q's head count and
integer-divided K/V by world_size, giving per-rank shapes that mismatched
the main-tensor slice (e.g. 14Q/3KV vs. 16Q/4KV for 28Q/7KV at SP=2).
"""
if current_omni_platform.get_device_count() < sp_world_size: if current_omni_platform.get_device_count() < sp_world_size:
pytest.skip(f"Test requires {sp_world_size} GPUs") pytest.skip(f"Test requires {sp_world_size} GPUs")
batch_size = 2 batch_size = 2
head_size = 8 head_size = 8
kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
has_joint = joint_len is not None
base_init_method = get_distributed_init_method() base_init_method = get_distributed_init_method()
sp_init_method = get_distributed_init_method() sp_init_method = get_distributed_init_method()
@@ -588,23 +200,17 @@ def test_ulysses_uaa_matches_baseline(
try: try:
torch.manual_seed(0) torch.manual_seed(0)
q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
k = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32) k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
v = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32) v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
payload = {"q": q.numpy(), "k": k.numpy(), "v": v.numpy()} np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy())
if has_joint:
payload["joint_q"] = torch.randn(batch_size, joint_len, num_heads, head_size, dtype=torch.float32).numpy()
payload["joint_k"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
payload["joint_v"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
np.savez(input_file, **payload)
common_tail = (1, None, "math", num_kv_heads, has_joint)
# Baseline (no SP) # Baseline (no SP)
torch.multiprocessing.spawn( torch.multiprocessing.spawn(
_run_attention_case, _run_attention_case,
args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict", *common_tail), args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"),
nprocs=1, nprocs=1,
) )
# SP (Ulysses-P with UAA) # SP (Ulysses-P with UAA)
torch.multiprocessing.spawn( torch.multiprocessing.spawn(
_run_attention_case, _run_attention_case,
@@ -617,13 +223,15 @@ def test_ulysses_uaa_matches_baseline(
head_size, head_size,
sp_world_size, sp_world_size,
"advanced_uaa", "advanced_uaa",
*common_tail,
), ),
nprocs=sp_world_size, nprocs=sp_world_size,
) )
baseline_t = torch.from_numpy(np.load(baseline_file, allow_pickle=False)) baseline = np.load(baseline_file, allow_pickle=False)
sp_t = torch.from_numpy(np.load(sp_file, allow_pickle=False)) sp = np.load(sp_file, allow_pickle=False)
baseline_t = torch.from_numpy(baseline)
sp_t = torch.from_numpy(sp)
assert baseline_t.shape == sp_t.shape assert baseline_t.shape == sp_t.shape
torch.testing.assert_close(sp_t, baseline_t, atol=1e-5, rtol=1e-5) torch.testing.assert_close(sp_t, baseline_t, atol=1e-5, rtol=1e-5)
finally: finally:
@@ -634,17 +242,7 @@ def test_ulysses_uaa_matches_baseline(
pass pass
@pytest.mark.parametrize( def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
"num_heads,num_kv_heads,joint_len",
[
(3, None, None), # MHA baseline (head_cnt not divisible by ulysses_degree=2)
(28, 7, None), # GQA: exercises Ring pytorch_attn_forward KV expansion (Q!=K heads)
(28, 7, 5), # GQA + MMDiT joint stream: joint K/V ride the padded Ulysses layout and Ring SDPA
],
)
def test_ulysses_uaa_hybrid_ring_matches_baseline(
num_heads: int, num_kv_heads: int | None, joint_len: int | None
) -> None:
sp_world_size = 4 sp_world_size = 4
ulysses_degree = 2 ulysses_degree = 2
ring_degree = 2 ring_degree = 2
@@ -655,8 +253,9 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline(
batch_size = 2 batch_size = 2
head_size = 8 head_size = 8
seq_len = 10 seq_len = 10
kv_heads = num_kv_heads if num_kv_heads is not None else num_heads # This is the positive Hybrid case: K/V heads must not require padding,
has_joint = joint_len is not None # which is unsupported by advanced_uaa when Ring is also enabled.
num_heads = 4
# Ensure ring ranks see equal post-Ulysses seq_len: # Ensure ring ranks see equal post-Ulysses seq_len:
# rank0/1 -> 3+2=5, rank2/3 -> 3+2=5 # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
@@ -675,14 +274,9 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline(
try: try:
torch.manual_seed(0) torch.manual_seed(0)
q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
k = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32) k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
v = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32) v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
payload = {"q": q.numpy(), "k": k.numpy(), "v": v.numpy()} np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy())
if has_joint:
payload["joint_q"] = torch.randn(batch_size, joint_len, num_heads, head_size, dtype=torch.float32).numpy()
payload["joint_k"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
payload["joint_v"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
np.savez(input_file, **payload)
# Baseline (no SP) # Baseline (no SP)
torch.multiprocessing.spawn( torch.multiprocessing.spawn(
@@ -699,8 +293,6 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline(
1, 1,
None, None,
"mem_efficient", "mem_efficient",
num_kv_heads,
has_joint,
), ),
nprocs=1, nprocs=1,
) )
@@ -720,8 +312,6 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline(
ring_degree, ring_degree,
split_sizes, split_sizes,
"mem_efficient", "mem_efficient",
num_kv_heads,
has_joint,
), ),
nprocs=sp_world_size, nprocs=sp_world_size,
) )
+1 -16
View File
@@ -6,9 +6,7 @@ Model specific tests for CacheDiT enablement.
""" """
import ast import ast
from contextlib import contextmanager
from pathlib import Path from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch from unittest.mock import Mock, patch
import pytest import pytest
@@ -19,8 +17,7 @@ from vllm.distributed import parallel_state
import vllm_omni.diffusion.cache.cachedit as cd_backend import vllm_omni.diffusion.cache.cachedit as cd_backend
import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific 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.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import DiffusionCacheConfig
from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer 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.helios.helios_transformer import HeliosTransformer3DModel
from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
@@ -38,17 +35,6 @@ SEPARATE_CFG_TRANSFORMERS = [
SAMPLE_CACHE_CONFIG = DiffusionCacheConfig() 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(): def test_custom_cache_dit_enablers_are_registered_explicitly():
expected_enablers = { expected_enablers = {
"Wan22Pipeline": cd_model_specific.enable_cache_for_wan22, "Wan22Pipeline": cd_model_specific.enable_cache_for_wan22,
@@ -227,7 +213,6 @@ def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
text_in = torch.full((1, seq_len, 16), 3.0) text_in = torch.full((1, seq_len, 16), 3.0)
audio_text_in = torch.full((1, seq_len, 16), 4.0) audio_text_in = torch.full((1, seq_len, 16), 4.0)
with _force_torch_sdpa():
model = LTX2VideoTransformer3DModel( model = LTX2VideoTransformer3DModel(
in_channels=16, in_channels=16,
out_channels=16, out_channels=16,
@@ -13,8 +13,6 @@ import vllm_omni.diffusion.data as diffusion_data
from tests.helpers.mark import hardware_test from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
from vllm_omni.diffusion.data import ( from vllm_omni.diffusion.data import (
AttentionConfig,
AttentionSpec,
DiffusionOutput, DiffusionOutput,
DiffusionParallelConfig, DiffusionParallelConfig,
OmniDiffusionConfig, OmniDiffusionConfig,
@@ -143,96 +141,9 @@ class TestPipelineArgumentsHandling:
assert isinstance(output.output, MockPipelineOutput) assert isinstance(output.output, MockPipelineOutput)
assert output.output.image is stub_image 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( @pytest.mark.parametrize(
"feature_id", "feature_id",
[ ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"],
"cfg_parallel",
"tensor_parallel",
"ulysses",
"ring",
"teacache",
"cache_dit",
"enforce_eager",
"unsupported_quantization",
],
) )
def test_adapter_guard_unsupported_feature(self, feature_id): def test_adapter_guard_unsupported_feature(self, feature_id):
if feature_id == "cfg_parallel": if feature_id == "cfg_parallel":
@@ -240,11 +151,6 @@ class TestPipelineArgumentsHandling:
parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1), parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1),
cache_backend="none", cache_backend="none",
) )
elif feature_id == "tensor_parallel":
od_config = _make_od_config(
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, sequence_parallel_size=1),
cache_backend="none",
)
elif feature_id == "ulysses": elif feature_id == "ulysses":
od_config = _make_od_config( od_config = _make_od_config(
parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ulysses_degree=2), parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ulysses_degree=2),
@@ -17,12 +17,7 @@ import torch
from tests.helpers.mark import hardware_marks, hardware_test from tests.helpers.mark import hardware_marks, hardware_test
from vllm_omni.diffusion.distributed import parallel_state from vllm_omni.diffusion.distributed import parallel_state
from vllm_omni.diffusion.distributed.cfg_parallel import ( from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _dispatch_branches, _wrap
CFGParallelMixin,
_dispatch_branches,
_get_cfg_world_size_or_one,
_wrap,
)
from vllm_omni.diffusion.distributed.parallel_state import ( from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env, destroy_distributed_env,
get_classifier_free_guidance_rank, get_classifier_free_guidance_rank,
@@ -398,68 +393,6 @@ def _build_pipeline_on_device(
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
def test_cfg_world_size_defaults_only_when_group_is_uninitialized(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(parallel_state, "_CFG", None)
assert not parallel_state.is_cfg_group_initialized()
assert _get_cfg_world_size_or_one() == 1
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
def test_cfg_world_size_does_not_swallow_initialized_group_errors(monkeypatch: pytest.MonkeyPatch):
class BrokenCfgGroup:
@property
def world_size(self) -> int:
raise AssertionError("CFG group is broken")
monkeypatch.setattr(parallel_state, "_CFG", BrokenCfgGroup())
assert parallel_state.is_cfg_group_initialized()
with pytest.raises(AssertionError, match="CFG group is broken"):
_get_cfg_world_size_or_one()
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
def test_cfg_defaults_to_sequential_without_initialized_group():
pipeline = _build_pipeline_on_device(
TestCFGPipeline,
hidden_dim=32,
model_seed=42,
dtype=torch.float32,
device=torch.device("cpu"),
)
positive_kwargs, negative_kwargs = _make_two_branch_inputs(
batch_size=1,
channels=4,
height=4,
width=4,
dtype=torch.float32,
device=torch.device("cpu"),
input_seed=123,
)
with torch.no_grad():
positive = pipeline.predict_noise(**positive_kwargs)
negative = pipeline.predict_noise(**negative_kwargs)
actual = pipeline.predict_noise_maybe_with_cfg(
do_true_cfg=True,
true_cfg_scale=5.0,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
expected = negative + 5.0 * (positive - negative)
torch.testing.assert_close(actual, expected)
@pytest.mark.core_model @pytest.mark.core_model
@pytest.mark.diffusion @pytest.mark.diffusion
@pytest.mark.cpu @pytest.mark.cpu
@@ -989,89 +989,6 @@ class TestDimensionValidation:
# Since expected_dims=3 but tensor has 4 dims, should return original # Since expected_dims=3 but tensor has 4 dims, should return original
assert result.shape == tensor_4d.shape 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 @pytest.mark.cpu
class TestStrictModeSplitValidation: class TestStrictModeSplitValidation:
-55
View File
@@ -1,55 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the CustomOp platform-dispatch base class."""
import pytest
import torch
from vllm_omni.diffusion.layers.custom_op import CustomOp
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
class _NativeOnlyOp(CustomOp):
"""An op that implements only the PyTorch-native path."""
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
return x * 2
def test_forward_xpu_falls_back_to_native():
"""XPU has no bespoke kernels, so the base class must reuse forward_native.
Without this fallback every CustomOp subclass has to define an identical
forward_xpu, and any that forgets crashes with NotImplementedError on XPU.
"""
op = _NativeOnlyOp()
x = torch.randn(4, 8)
assert torch.equal(op.forward_xpu(x), op.forward_native(x))
def test_forward_cuda_still_requires_an_implementation():
"""The native fallback must not mask a missing CUDA kernel."""
op = _NativeOnlyOp()
with pytest.raises(NotImplementedError):
op.forward_cuda(torch.randn(4, 8))
def test_mot_rmsnorm_dispatches_on_xpu_without_an_override():
"""MoTRMSNorm relies on the inherited fallback (regression test)."""
from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
assert "forward_xpu" not in MoTRMSNorm.__dict__
norm = MoTRMSNorm(16)
x = torch.randn(8, 16)
text_indices = torch.arange(0, 4)
vae_indices = torch.arange(4, 8)
assert torch.equal(norm.forward_xpu(x), norm.forward_native(x))
assert torch.equal(
norm.forward_xpu(x, text_indices, vae_indices),
norm.forward_native(x, text_indices, vae_indices),
)
+1 -71
View File
@@ -29,7 +29,7 @@ class _DummyQuantMethod:
return y return y
def test_diffusion_base_linear_apply_multi_slice(monkeypatch): def test_diffusion_base_linear_apply_multi_slice():
# Build a fake diffusion LoRA layer with 2 slices and rank=2. # Build a fake diffusion LoRA layer with 2 slices and rank=2.
layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA) layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
layer.tp_size = 1 layer.tp_size = 1
@@ -72,17 +72,7 @@ def test_diffusion_base_linear_apply_multi_slice(monkeypatch):
layer.lora_b_stacked = (b0, b1) layer.lora_b_stacked = (b0, b1)
layer.output_slices = out_slices layer.output_slices = out_slices
addmm_calls = []
original_addmm = torch.addmm
def _record_addmm(input, mat1, mat2, *, out=None):
addmm_calls.append((input, out))
return original_addmm(input, mat1, mat2, out=out)
monkeypatch.setattr(torch, "addmm", _record_addmm)
x = torch.tensor([[1.0, 2.0, 3.0]]) x = torch.tensor([[1.0, 2.0, 3.0]])
with torch.inference_mode():
out = layer.apply(x) out = layer.apply(x)
# Base output is identity: [1,2,3] # Base output is identity: [1,2,3]
@@ -97,66 +87,6 @@ def test_diffusion_base_linear_apply_multi_slice(monkeypatch):
delta1 = torch.tensor([[6.0]]) delta1 = torch.tensor([[6.0]])
expected = torch.cat([base_out[:, :2] + delta0, base_out[:, 2:3] + delta1], dim=-1) expected = torch.cat([base_out[:, :2] + delta0, base_out[:, 2:3] + delta1], dim=-1)
assert torch.allclose(out, expected) assert torch.allclose(out, expected)
assert len(addmm_calls) == 2
assert all(input is out for input, out in addmm_calls)
def test_diffusion_base_linear_apply_preserves_autograd(monkeypatch):
layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
layer.tp_size = 1
layer.lora_config = _DummyLoRAConfig()
layer.base_layer = type("Base", (), {})()
layer.base_layer.quant_method = _DummyQuantMethod(torch.eye(2))
a = torch.tensor([[[[1.0, 2.0]]]], requires_grad=True)
b = torch.tensor([[[[3.0], [4.0]]]], requires_grad=True)
layer.lora_a_stacked = (a,)
layer.lora_b_stacked = (b,)
layer.output_slices = (2,)
layer._diffusion_lora_active_slices = (True,)
def _unexpected_addmm(*args, **kwargs):
raise AssertionError("the out= fast path must not run with autograd enabled")
monkeypatch.setattr(torch, "addmm", _unexpected_addmm)
x = torch.tensor([[2.0, 1.0]], requires_grad=True)
out = layer.apply(x)
assert torch.allclose(out, torch.tensor([[14.0, 17.0]]))
out.sum().backward()
assert torch.allclose(x.grad, torch.tensor([[8.0, 15.0]]))
assert torch.allclose(a.grad, torch.tensor([[[[14.0, 7.0]]]]))
assert torch.allclose(b.grad, torch.tensor([[[[4.0], [4.0]]]]))
def test_diffusion_base_linear_apply_falls_back_for_mixed_output_dtype(monkeypatch):
layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
layer.tp_size = 1
layer.lora_config = _DummyLoRAConfig()
class _Float64OutputQuantMethod(_DummyQuantMethod):
def apply(self, base_layer, x: torch.Tensor, bias: torch.Tensor | None):
return super().apply(base_layer, x, bias).to(torch.float64)
layer.base_layer = type("Base", (), {})()
layer.base_layer.quant_method = _Float64OutputQuantMethod(torch.eye(2))
layer.lora_a_stacked = (torch.ones((1, 1, 1, 2)),)
layer.lora_b_stacked = (torch.tensor([[[[1.0], [2.0]]]]),)
layer.output_slices = (2,)
layer._diffusion_lora_active_slices = (True,)
def _unexpected_addmm(*args, **kwargs):
raise AssertionError("mixed output and LoRA dtypes must use the functional fallback")
monkeypatch.setattr(torch, "addmm", _unexpected_addmm)
with torch.inference_mode():
out = layer.apply(torch.tensor([[1.0, 2.0]]))
assert out.dtype == torch.float64
assert torch.allclose(out, torch.tensor([[4.0, 8.0]], dtype=torch.float64))
def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch): def test_diffusion_lora_buffer_device_persists_across_reallocation(monkeypatch):
+1 -55
View File
@@ -1,6 +1,3 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
from collections import defaultdict from collections import defaultdict
import pytest import pytest
@@ -109,7 +106,7 @@ def make_dummy_lora_state_dict(
""" """
key_a = f"{base_key}.{lora_a_suffix}" key_a = f"{base_key}.{lora_a_suffix}"
key_b = f"{base_key}.{lora_b_suffix}" key_b = f"{base_key}.{lora_b_suffix}"
key_bias = f"{base_key}.{lora_bias_suffix}" key_bias = f"{base_key}.{lora_bias_suffix}" if bias else None
if weight_a is None: if weight_a is None:
weight_a = torch.randn(RANK, HEAD_DIM) weight_a = torch.randn(RANK, HEAD_DIM)
@@ -439,54 +436,3 @@ def test_wan_load_lora_weights_targets_both_transformers(mocker: MockerFixture):
assert_params_equal(pipeline.transformer, original_high) assert_params_equal(pipeline.transformer, original_high)
assert_params_equal(pipeline.transformer_2, original_low) assert_params_equal(pipeline.transformer_2, original_low)
assert "adapter0" not in pipeline.lora_loaded assert "adapter0" not in pipeline.lora_loaded
def test_prepare_lora_delta_disambiguates_prefix_overlapping_stacked_names():
"""A fused name can be a prefix of another fused name.
SenseNova-U1 carries both `.qkv_proj` (understanding) and `.qkv_proj_mot_gen`
(generation). Selecting stacked mappings by substring matched both rules for
the generation weight and concatenated six shards into a delta twice as tall
as the parameter.
"""
param_to_weight_names = {
".qkv_proj_mot_gen": [".q_proj_mot_gen", ".k_proj_mot_gen", ".v_proj_mot_gen"],
".qkv_proj": [".q_proj", ".k_proj", ".v_proj"],
}
state_dict = {}
for part in ("q", "k", "v"):
for tower in ("_mot_gen", ""):
base = f"layers.0.self_attn.{part}_proj{tower}"
state_dict[f"{base}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
state_dict[f"{base}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
for base_key in ("layers.0.self_attn.qkv_proj_mot_gen", "layers.0.self_attn.qkv_proj"):
delta, _ = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
assert delta.shape == (3 * HEAD_DIM, HEAD_DIM), base_key
@pytest.mark.parametrize(
"fused, shards",
[
# qwen_image_transformer.py stacked_params_mapping
(".to_qkv", (".to_q", ".to_k", ".to_v")),
(".add_kv_proj", (".add_q_proj", ".add_k_proj", ".add_v_proj")),
# wan2_2_transformer.py stacked_params_mapping
(".attn1.to_qkv", (".attn1.to_q", ".attn1.to_k", ".attn1.to_v")),
# wan2_2_s2v_transformer.py stacked_params_mapping
(".self_attn.to_qkv", (".self_attn.q", ".self_attn.k", ".self_attn.v")),
],
)
def test_prepare_lora_delta_unchanged_for_existing_fused_layouts(fused, shards):
"""None of the fused names shipped today overlap by prefix, so matching on the
tail must select exactly the shards substring matching selected."""
param_to_weight_names = {fused: list(shards)}
base_key = f"blocks.0{fused}"
state_dict = {}
for shard in shards:
state_dict[f"blocks.0{shard}.lora_A.weight"] = torch.ones(RANK, HEAD_DIM)
state_dict[f"blocks.0{shard}.lora_B.weight"] = torch.ones(HEAD_DIM, RANK)
delta, used_keys = _prepare_lora_delta(state_dict, base_key, param_to_weight_names)
assert delta.shape == (len(shards) * HEAD_DIM, HEAD_DIM)
assert len(used_keys) == 2 * len(shards)
@@ -614,54 +614,6 @@ def test_initialize_model_sets_current_diffusion_config_during_model_constructio
assert get_current_diffusion_config_or_none() is None assert get_current_diffusion_config_or_none() is None
def test_configure_vae_stack_tiling_enables_tiling_and_native_capability():
from vllm_omni.diffusion.registry import _configure_vae_stack_tiling
calls = []
vae = SimpleNamespace(
supports_stack_tiling=True,
set_stack_tiling=lambda mode: calls.append(mode),
)
model = SimpleNamespace(vae=vae)
od_config = SimpleNamespace(
model_class_name="MiniMaxH3Pipeline",
vae_stack_tiling="auto",
vae_use_tiling=False,
)
_configure_vae_stack_tiling(model, od_config)
assert od_config.vae_use_tiling is True
assert calls == ["auto"]
def test_configure_vae_stack_tiling_rejects_unsupported_vae():
from vllm_omni.diffusion.registry import _configure_vae_stack_tiling
od_config = SimpleNamespace(
model_class_name="UnsupportedPipeline",
vae_stack_tiling="true",
vae_use_tiling=False,
)
with pytest.raises(ValueError, match="does not declare stacked-tiling support"):
_configure_vae_stack_tiling(SimpleNamespace(vae=SimpleNamespace()), od_config)
def test_configure_vae_stack_tiling_auto_falls_back_for_unsupported_vae():
from vllm_omni.diffusion.registry import _configure_vae_stack_tiling
od_config = SimpleNamespace(
model_class_name="UnsupportedPipeline",
vae_stack_tiling="auto",
vae_use_tiling=False,
)
_configure_vae_stack_tiling(SimpleNamespace(vae=SimpleNamespace()), od_config)
assert od_config.vae_use_tiling is False
def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch): def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import logging import logging
import os import os
@@ -23,13 +23,9 @@ NORM_EPS = 1e-5
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True)
def _init_distributed(request): def _init_distributed():
"""Initialize the minimal single-rank distributed environment required by """Initialize the minimal single-rank distributed environment required by
the vLLM parallel linear layers (tensor-parallel group must exist).""" the vLLM parallel linear layers (tensor-parallel group must exist)."""
if request.node.name.startswith("test_real_rotary_"):
yield
return
from vllm.distributed.parallel_state import ( from vllm.distributed.parallel_state import (
cleanup_dist_env_and_memory, cleanup_dist_env_and_memory,
init_distributed_environment, init_distributed_environment,
@@ -63,20 +59,6 @@ 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: def _randomize_parameters(module: torch.nn.Module) -> None:
"""Fill parameters with small random values. """Fill parameters with small random values.
@@ -88,11 +70,11 @@ def _randomize_parameters(module: torch.nn.Module) -> None:
param.uniform_(-0.02, 0.02) param.uniform_(-0.02, 0.02)
def _identity_rotary_emb(batch_size: int, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: def _identity_rotary_emb(batch_size: int, seq_len: int) -> torch.Tensor:
"""Real-valued rotary frequencies encoding a zero rotation.""" """Complex rotary frequencies encoding a zero rotation."""
return ( return torch.polar(
torch.ones(batch_size, seq_len, HEAD_DIM), torch.ones(batch_size, seq_len, HEAD_DIM // 2),
torch.zeros(batch_size, seq_len, HEAD_DIM), torch.zeros(batch_size, seq_len, HEAD_DIM // 2),
) )
@@ -137,72 +119,6 @@ def test_boogu_image_transformer_import():
assert BooguImageTransformer2DModel is not None assert BooguImageTransformer2DModel is not None
def test_real_rotary_emb_matches_complex_reference():
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import apply_rotary_emb
torch.manual_seed(0)
x = torch.randn(2, 7, NUM_HEADS, HEAD_DIM)
angles = torch.randn(2, 7, HEAD_DIM // 2)
rotary_emb = (
angles.cos().repeat_interleave(2, dim=-1),
angles.sin().repeat_interleave(2, dim=-1),
)
actual = apply_rotary_emb(x, rotary_emb)
x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], HEAD_DIM // 2, 2))
freqs_cis = torch.polar(torch.ones_like(angles), angles).unsqueeze(2)
expected = torch.view_as_real(x_complex * freqs_cis).flatten(3).type_as(x)
torch.testing.assert_close(actual, expected)
def test_real_rotary_frequency_tables_are_real_and_repeated():
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
BooguImageDoubleStreamRotaryPosEmbed,
)
freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(AXES_DIM_ROPE, AXES_LENS, theta=10000)
assert len(freqs_real) == len(AXES_DIM_ROPE)
for (freqs_cos, freqs_sin), axis_dim, axis_len in zip(freqs_real, AXES_DIM_ROPE, AXES_LENS):
assert freqs_cos.shape == (axis_len, axis_dim)
assert freqs_sin.shape == (axis_len, axis_dim)
assert not freqs_cos.is_complex()
assert not freqs_sin.is_complex()
torch.testing.assert_close(freqs_cos[..., ::2], freqs_cos[..., 1::2])
torch.testing.assert_close(freqs_sin[..., ::2], freqs_sin[..., 1::2])
def test_real_rotary_frequency_tables_use_float32_on_npu(monkeypatch):
from vllm_omni.diffusion.models.boogu_image import boogu_image_transformer
npu_platform = SimpleNamespace(is_npu=lambda: True, supports_float64=lambda: True)
monkeypatch.setattr(boogu_image_transformer, "current_omni_platform", npu_platform)
freqs_real = boogu_image_transformer.BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
AXES_DIM_ROPE, AXES_LENS, theta=10000
)
assert all(freqs.dtype == torch.float32 for pair in freqs_real for freqs in pair)
def test_real_rotary_frequency_tables_use_float32_on_mps(monkeypatch):
from vllm_omni.diffusion.models.boogu_image import boogu_image_transformer
platform = SimpleNamespace(
is_npu=lambda: False,
supports_float64=lambda: True,
)
monkeypatch.setattr(boogu_image_transformer, "current_omni_platform", platform)
monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True)
freqs_real = boogu_image_transformer.BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(
AXES_DIM_ROPE, AXES_LENS, theta=10000
)
assert all(freqs.dtype == torch.float32 for pair in freqs_real for freqs in pair)
def test_single_stream_block_shape(): def test_single_stream_block_shape():
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
BooguImageTransformerBlock, BooguImageTransformerBlock,
@@ -482,10 +398,10 @@ def test_transformer_forward_t2i_shape():
timestep = torch.full((batch_size,), 0.5) timestep = torch.full((batch_size,), 0.5)
instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim) instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool) instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(model.axes_dim_rope, model.axes_lens, theta=10000) freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
with torch.no_grad(): with torch.no_grad():
out = model(latents, timestep, instruction_hidden_states, freqs_real, instruction_attention_mask) out = model(latents, timestep, instruction_hidden_states, freqs_cis, instruction_attention_mask)
assert out.shape == (batch_size, model.out_channels, latent_h, latent_w) assert out.shape == (batch_size, model.out_channels, latent_h, latent_w)
assert torch.isfinite(out).all() assert torch.isfinite(out).all()
@@ -515,7 +431,7 @@ def test_transformer_forward_ti2i_shape():
timestep = torch.full((batch_size,), 0.5) timestep = torch.full((batch_size,), 0.5)
instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim) instruction_hidden_states = torch.randn(batch_size, instruct_len, instruction_feat_dim)
instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool) instruction_attention_mask = torch.ones(batch_size, instruct_len, dtype=torch.bool)
freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(model.axes_dim_rope, model.axes_lens, theta=10000) freqs_cis = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_cis(model.axes_dim_rope, model.axes_lens, theta=10000)
# One sample, one reference image (Boogu editing supports a single ref). # One sample, one reference image (Boogu editing supports a single ref).
ref_image_hidden_states = [[torch.randn(in_channels, ref_h, ref_w)]] ref_image_hidden_states = [[torch.randn(in_channels, ref_h, ref_w)]]
@@ -525,7 +441,7 @@ def test_transformer_forward_ti2i_shape():
latents, latents,
timestep, timestep,
instruction_hidden_states, instruction_hidden_states,
freqs_real, freqs_cis,
instruction_attention_mask, instruction_attention_mask,
ref_image_hidden_states=ref_image_hidden_states, ref_image_hidden_states=ref_image_hidden_states,
) )
@@ -192,6 +192,7 @@ def test_constructor_weights_sources(boogu_pipeline):
(DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"), (DiffusionParallelConfig(tensor_parallel_size=2), "none", "Tensor parallelism"),
(DiffusionParallelConfig(ulysses_degree=2), "none", "Sequence parallelism"), (DiffusionParallelConfig(ulysses_degree=2), "none", "Sequence parallelism"),
(DiffusionParallelConfig(ring_degree=2), "none", "Sequence parallelism"), (DiffusionParallelConfig(ring_degree=2), "none", "Sequence parallelism"),
(DiffusionParallelConfig(cfg_parallel_size=2), "none", "CFG parallelism"),
( (
DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=2), DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=2),
"none", "none",
@@ -226,26 +227,6 @@ def test_constructor_rejects_unsupported_execution_modes(
mock_dependencies["mllm_wrapper"].model.to.assert_not_called() mock_dependencies["mllm_wrapper"].model.to.assert_not_called()
@pytest.mark.parametrize("cfg_parallel_size", [2, 3])
def test_constructor_accepts_cfg_parallel(mock_dependencies, cfg_parallel_size):
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
BooguImagePipeline,
)
od_config = OmniDiffusionConfig(
model="dummy-boogu",
tf_model_config=TransformerConfig(params={}),
dtype=torch.float32,
parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size),
)
pipeline = BooguImagePipeline(od_config=od_config)
assert pipeline.od_config.parallel_config.cfg_parallel_size == cfg_parallel_size
assert hasattr(pipeline, "predict_noise_maybe_with_cfg")
assert hasattr(pipeline, "predict_noise_with_multi_branch_cfg")
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Prompt-encoding tests (deterministic fakes, no constructor) # Prompt-encoding tests (deterministic fakes, no constructor)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -486,8 +467,7 @@ class _FakeTransformer:
"reduce_type": "mean", "reduce_type": "mean",
} }
def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
assert all(not freqs.is_complex() for pair in freqs_real for freqs in pair)
return torch.zeros_like(latents) return torch.zeros_like(latents)
@@ -545,14 +525,6 @@ def test_forward_returns_diffusion_output():
def test_forward_cfg_off_when_guidance_one(): def test_forward_cfg_off_when_guidance_one():
pipeline = _make_forward_pipeline() pipeline = _make_forward_pipeline()
cfg_calls = []
original_predict_noise_maybe_with_cfg = pipeline.predict_noise_maybe_with_cfg
def record_cfg_call(**kwargs):
cfg_calls.append(kwargs)
return original_predict_noise_maybe_with_cfg(**kwargs)
pipeline.predict_noise_maybe_with_cfg = record_cfg_call
req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2, guidance_scale=1.0) req = _make_request_batch("a cat", height=64, width=64, num_inference_steps=2, guidance_scale=1.0)
# guidance_scale=1.0 is falsy-adjacent but explicitly disables CFG; the # guidance_scale=1.0 is falsy-adjacent but explicitly disables CFG; the
# request layer would set guidance_scale_provided, so emulate that here. # request layer would set guidance_scale_provided, so emulate that here.
@@ -562,30 +534,6 @@ def test_forward_cfg_off_when_guidance_one():
# Only the positive prompt is encoded when CFG is off. # Only the positive prompt is encoded when CFG is off.
assert len(pipeline.processor.calls) == 1 assert len(pipeline.processor.calls) == 1
assert len(cfg_calls) == 2
assert all(call["do_true_cfg"] is False for call in cfg_calls)
assert all(call["negative_kwargs"] is None for call in cfg_calls)
def test_double_guidance_combine_matches_legacy_formula():
pipeline = _make_forward_pipeline()
positive_with_reference = torch.randn(2, 4, 8, 8)
negative_with_reference = torch.randn(2, 4, 8, 8)
uncond = torch.randn(2, 4, 8, 8)
text_scale = 5.0
image_scale = 2.0
actual = pipeline.combine_multi_branch_cfg_noise(
[positive_with_reference, negative_with_reference, uncond],
{"text": text_scale, "image": image_scale},
)
legacy = (
positive_with_reference
+ (text_scale - 1.0) * (positive_with_reference - negative_with_reference)
+ (image_scale - 1.0) * (negative_with_reference - uncond)
)
torch.testing.assert_close(actual, legacy, rtol=0, atol=1e-5)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -853,39 +801,6 @@ def test_build_ref_latents_expands_per_output_and_handles_none():
assert ref_latents[2] is None and ref_latents[3] is None assert ref_latents[2] is None and ref_latents[3] is None
def test_vae_attention_context_is_noop_on_cpu():
from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import (
BooguImagePipeline,
)
with BooguImagePipeline._vae_attention_context(torch.device("cpu")):
pass
def test_vae_attention_context_prioritizes_efficient_cuda(monkeypatch):
from vllm_omni.diffusion.models.boogu_image import pipeline_boogu_image
recorded = {}
def fake_sdpa_kernel(backends, *, set_priority):
recorded["backends"] = backends
recorded["set_priority"] = set_priority
return pipeline_boogu_image.nullcontext()
monkeypatch.setattr(pipeline_boogu_image, "sdpa_kernel", fake_sdpa_kernel)
with pipeline_boogu_image.BooguImagePipeline._vae_attention_context(torch.device("cuda")):
pass
assert recorded == {
"backends": [
pipeline_boogu_image.SDPBackend.EFFICIENT_ATTENTION,
pipeline_boogu_image.SDPBackend.MATH,
],
"set_priority": True,
}
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Editing / TI2I: forward CFG branch selection # Editing / TI2I: forward CFG branch selection
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -896,11 +811,9 @@ class _RecordingRefTransformer(_FakeTransformer):
def __init__(self): def __init__(self):
self.calls = [] self.calls = []
self.instruction_embed_calls = []
def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs): def __call__(self, latents, timestep, instruction_embeds, freqs_cis, instruction_attention_mask, **kwargs):
self.calls.append(kwargs.get("ref_image_hidden_states")) self.calls.append(kwargs.get("ref_image_hidden_states"))
self.instruction_embed_calls.append(instruction_embeds)
return torch.zeros_like(latents) return torch.zeros_like(latents)
@@ -962,9 +875,6 @@ def test_forward_ti2i_text_only_two_predictions_with_ref():
# Text-only ti2i: cond+ref and neg+ref -> 2 predictions/step, both carry ref. # Text-only ti2i: cond+ref and neg+ref -> 2 predictions/step, both carry ref.
assert _count_per_step(pipeline.transformer.calls, num_steps) == 2 assert _count_per_step(pipeline.transformer.calls, num_steps) == 2
assert all(ref is not None for ref in pipeline.transformer.calls) assert all(ref is not None for ref in pipeline.transformer.calls)
for i in range(0, len(pipeline.transformer.instruction_embed_calls), 2):
positive, negative = pipeline.transformer.instruction_embed_calls[i : i + 2]
assert positive is not negative
def test_forward_ti2i_double_guidance_three_predictions(): def test_forward_ti2i_double_guidance_three_predictions():
@@ -984,11 +894,6 @@ def test_forward_ti2i_double_guidance_three_predictions():
per_step = [pipeline.transformer.calls[i : i + 3] for i in range(0, len(pipeline.transformer.calls), 3)] per_step = [pipeline.transformer.calls[i : i + 3] for i in range(0, len(pipeline.transformer.calls), 3)]
for step_calls in per_step: for step_calls in per_step:
assert sum(ref is None for ref in step_calls) == 1 assert sum(ref is None for ref in step_calls) == 1
embed_calls = pipeline.transformer.instruction_embed_calls
for i in range(0, len(embed_calls), 3):
positive, negative_with_reference, uncond = embed_calls[i : i + 3]
assert positive is not negative_with_reference
assert negative_with_reference is uncond
def test_forward_ti2i_image_only_two_predictions_drop_ref(): def test_forward_ti2i_image_only_two_predictions_drop_ref():
@@ -1008,10 +913,6 @@ def test_forward_ti2i_image_only_two_predictions_drop_ref():
per_step = [pipeline.transformer.calls[i : i + 2] for i in range(0, len(pipeline.transformer.calls), 2)] per_step = [pipeline.transformer.calls[i : i + 2] for i in range(0, len(pipeline.transformer.calls), 2)]
for step_calls in per_step: for step_calls in per_step:
assert sum(ref is None for ref in step_calls) == 1 assert sum(ref is None for ref in step_calls) == 1
embed_calls = pipeline.transformer.instruction_embed_calls
for i in range(0, len(embed_calls), 2):
positive_with_reference, positive_without_reference = embed_calls[i : i + 2]
assert positive_with_reference is positive_without_reference
def test_forward_ti2i_no_guidance_single_prediction_with_ref(): def test_forward_ti2i_no_guidance_single_prediction_with_ref():
@@ -1150,72 +1150,6 @@ def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
assert cfg.fps == sp.resolved_frame_rate == 12.0 assert cfg.fps == sp.resolved_frame_rate == 12.0
def test_transfer_vae_executor_requires_distributed_vae() -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
executor = object()
pipeline.vae = SimpleNamespace(distributed_executor=executor, is_distributed_enabled=lambda: True)
assert pipeline._transfer_vae_executor() is executor
pipeline.vae = SimpleNamespace(distributed_executor=executor, is_distributed_enabled=lambda: False)
assert pipeline._transfer_vae_executor() is None
def test_sync_transfer_overlap_slices_output_rank_and_broadcasts() -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
class RecordingExecutor:
rank = 0
def __init__(self) -> None:
self.value = None
def broadcast_tensor(self, value):
self.value = value
return value
pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
output = torch.arange(1 * 3 * 5 * 2 * 2).reshape(1, 3, 5, 2, 2)
executor = RecordingExecutor()
overlap = pipeline._sync_transfer_overlap(
output,
overlap_frames=2,
reference_video=output,
vae_executor=executor,
)
assert overlap is not None
torch.testing.assert_close(overlap, output[:, :, -2:])
assert executor.value is overlap
def test_sync_transfer_overlap_allocates_receive_buffer_on_non_output_rank() -> None:
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
class ReceivingExecutor:
rank = 1
def broadcast_tensor(self, value):
assert value.shape == (1, 3, 2, 4, 5)
return torch.ones_like(value)
pipeline = SimpleNamespace(device=torch.device("cpu"), vae=SimpleNamespace(dtype=torch.float32))
reference = torch.zeros(1, 3, 6, 4, 5)
overlap = Cosmos3OmniDiffusersPipeline._sync_transfer_overlap(
pipeline,
torch.empty(0),
overlap_frames=2,
reference_video=reference,
vae_executor=ReceivingExecutor(),
)
assert overlap is not None
assert torch.equal(overlap, torch.ones_like(overlap))
def test_transfer_edge_uses_rgb_canny(monkeypatch: pytest.MonkeyPatch) -> None: def test_transfer_edge_uses_rgb_canny(monkeypatch: pytest.MonkeyPatch) -> None:
from vllm_omni.diffusion.models.cosmos3 import transfer from vllm_omni.diffusion.models.cosmos3 import transfer
@@ -24,7 +24,6 @@ from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
build_hunyuan_diffusion_kv_requests, build_hunyuan_diffusion_kv_requests,
extract_hunyuan_prompt_inputs, extract_hunyuan_prompt_inputs,
hunyuan_num_image_tokens, hunyuan_num_image_tokens,
hunyuan_num_special_tokens,
normalize_hunyuan_cot_text, normalize_hunyuan_cot_text,
prepare_hunyuan_layout, prepare_hunyuan_layout,
) )
@@ -115,7 +114,7 @@ class _FakeImageProcessor:
self.image_sizes: list[tuple[int, int]] = [] self.image_sizes: list[tuple[int, int]] = []
self.vision_encoder_processor = SimpleNamespace(patch_size=1) self.vision_encoder_processor = SimpleNamespace(patch_size=1)
def build_image_info(self, image_size, **kwargs): def build_image_info(self, image_size):
self.image_sizes.append(image_size) self.image_sizes.append(image_size)
return ImageInfo( return ImageInfo(
image_type="gen_image", image_type="gen_image",
@@ -126,7 +125,6 @@ class _FakeImageProcessor:
image_token_length=self.image_token_length, image_token_length=self.image_token_length,
base_size=1024, base_size=1024,
ratio_index=0, ratio_index=0,
**kwargs,
) )
@@ -168,29 +166,6 @@ def _request(
) )
def test_distilled_layout_uses_embedded_cfg_and_meanflow_tokens() -> None:
tokenizer, image_processor = _components([12])
request = _request(guidance_scale=2.5)
prepared_layout = prepare_hunyuan_layout(
request,
tokenizer_wrapper=tokenizer,
image_processor=image_processor,
generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False),
image_base_size=1024,
cfg_distilled=True,
use_meanflow=True,
)
image_info = prepared_layout.generated_image_info
assert image_info.add_guidance_token
assert image_info.add_timestep_r_token
assert tokenizer.calls[0]["cfg_factor"] == 1
assert hunyuan_num_special_tokens(image_info) == 3
assert hunyuan_num_image_tokens(image_info) == 19
assert len(build_hunyuan_diffusion_kv_requests(request, prepared_layout)) == 1
def test_builds_kv_request_lengths_without_model_execution() -> None: def test_builds_kv_request_lengths_without_model_execution() -> None:
tokenizer, image_processor = _components([12]) tokenizer, image_processor = _components([12])
request = _request(guidance_scale=1.0) request = _request(guidance_scale=1.0)
@@ -410,13 +385,9 @@ def test_prepared_model_inputs_match_local_tokenization(
torch.testing.assert_close(prepared_mask, local_mask) torch.testing.assert_close(prepared_mask, local_mask)
assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"] assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"]
image_info = prepared_layout.generated_image_info num_image_tokens = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
generation_token_counts = { local_inputs.update(attention_mask=local_mask, num_image_tokens=num_image_tokens)
"num_image_tokens": hunyuan_num_image_tokens(image_info), prepared_inputs.update(attention_mask=prepared_mask, num_image_tokens=num_image_tokens)
"num_special_tokens": hunyuan_num_special_tokens(image_info),
}
local_inputs.update(attention_mask=local_mask, **generation_token_counts)
prepared_inputs.update(attention_mask=prepared_mask, **generation_token_counts)
local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs) local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs)
prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs) prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs)
@@ -1,294 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for HunyuanImage3 prepare_inputs_for_generation to ensure
distilled model parameters (guidance, timesteps_r) are properly passed."""
import pytest
import torch
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
HunyuanImage3Text2ImagePipeline,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class MockHunyuanImage3Pipeline:
"""Minimal mock that replicates prepare_inputs_for_generation method."""
def __init__(self):
# Mock config for distilled models
self.config = type("Config", (), {})()
self.hf_config = type("HFConfig", (), {"cfg_distilled": True, "use_meanflow": True})()
# Bind the real method from the pipeline class
from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
HunyuanImage3Pipeline as _Real,
)
prepare_inputs_for_generation = _Real.prepare_inputs_for_generation
_update_model_kwargs_for_generation = _Real._update_model_kwargs_for_generation
class TestPrepareInputsForGenerationDistilled:
"""Tests that distilled model parameters are passed through
prepare_inputs_for_generation."""
def setup_method(self):
self.pipeline = MockHunyuanImage3Pipeline()
def _create_mock_kwargs(self, include_distilled_params: bool = True) -> dict:
"""Create mock kwargs for prepare_inputs_for_generation."""
kwargs = {
"position_ids": torch.tensor([[0, 1, 2]]),
"custom_pos_emb": (torch.tensor([[0.0]]), torch.tensor([[1.0]])),
"mode": "gen_image",
"images": torch.randn(1, 4, 64, 64),
"image_mask": torch.ones(1, 64, 64),
"timestep": torch.tensor([0.5]),
"gen_timestep_scatter_index": torch.tensor([[5]]),
"cond_vae_images": None,
"cond_timestep": None,
"cond_vae_image_mask": None,
"cond_vit_images": None,
"cond_vit_image_mask": None,
"vit_kwargs": None,
"cond_timestep_scatter_index": None,
"query_lens": [10],
"seq_lens": [100],
"num_image_tokens": 4096,
"num_special_tokens": 3,
"ar_kv_reuse_len": 0,
"full_attn_spans": None,
"use_cache": True,
}
if include_distilled_params:
# CFG distilled model parameters
kwargs["guidance"] = torch.tensor([2500.0])
kwargs["guidance_scatter_index"] = torch.tensor([[4]])
# MeanFlow distilled model parameters
kwargs["timesteps_r"] = torch.tensor([0.3])
kwargs["timesteps_r_scatter_index"] = torch.tensor([[3]])
else:
kwargs["guidance"] = None
kwargs["guidance_scatter_index"] = None
kwargs["timesteps_r"] = None
kwargs["timesteps_r_scatter_index"] = None
return kwargs
def test_guidance_passed_to_model_inputs(self):
"""Test that guidance is passed through prepare_inputs_for_generation."""
kwargs = self._create_mock_kwargs(include_distilled_params=True)
input_ids = torch.tensor([[1, 2, 3]])
attention_mask = torch.ones(1, 3)
model_inputs = self.pipeline.prepare_inputs_for_generation(
input_ids=input_ids,
past_key_values=None,
attention_mask=attention_mask,
inputs_embeds=None,
tokenizer_output=None,
batch_gen_image_info=None,
generator=None,
**kwargs,
)
assert "guidance" in model_inputs, "guidance should be in model_inputs"
assert model_inputs["guidance"] is not None, "guidance should not be None"
assert torch.equal(model_inputs["guidance"], kwargs["guidance"])
def test_guidance_scatter_index_passed_to_model_inputs(self):
"""Test that guidance_scatter_index is passed through."""
kwargs = self._create_mock_kwargs(include_distilled_params=True)
input_ids = torch.tensor([[1, 2, 3]])
attention_mask = torch.ones(1, 3)
model_inputs = self.pipeline.prepare_inputs_for_generation(
input_ids=input_ids,
past_key_values=None,
attention_mask=attention_mask,
inputs_embeds=None,
tokenizer_output=None,
batch_gen_image_info=None,
generator=None,
**kwargs,
)
assert "guidance_scatter_index" in model_inputs
assert model_inputs["guidance_scatter_index"] is not None
assert torch.equal(model_inputs["guidance_scatter_index"], kwargs["guidance_scatter_index"])
def test_timesteps_r_passed_to_model_inputs(self):
"""Test that timesteps_r is passed through for MeanFlow models."""
kwargs = self._create_mock_kwargs(include_distilled_params=True)
input_ids = torch.tensor([[1, 2, 3]])
attention_mask = torch.ones(1, 3)
model_inputs = self.pipeline.prepare_inputs_for_generation(
input_ids=input_ids,
past_key_values=None,
attention_mask=attention_mask,
inputs_embeds=None,
tokenizer_output=None,
batch_gen_image_info=None,
generator=None,
**kwargs,
)
assert "timesteps_r" in model_inputs, "timesteps_r should be in model_inputs"
assert model_inputs["timesteps_r"] is not None
assert torch.equal(model_inputs["timesteps_r"], kwargs["timesteps_r"])
def test_timesteps_r_scatter_index_passed_to_model_inputs(self):
"""Test that timesteps_r_scatter_index is passed through."""
kwargs = self._create_mock_kwargs(include_distilled_params=True)
input_ids = torch.tensor([[1, 2, 3]])
attention_mask = torch.ones(1, 3)
model_inputs = self.pipeline.prepare_inputs_for_generation(
input_ids=input_ids,
past_key_values=None,
attention_mask=attention_mask,
inputs_embeds=None,
tokenizer_output=None,
batch_gen_image_info=None,
generator=None,
**kwargs,
)
assert "timesteps_r_scatter_index" in model_inputs
assert model_inputs["timesteps_r_scatter_index"] is not None
assert torch.equal(model_inputs["timesteps_r_scatter_index"], kwargs["timesteps_r_scatter_index"])
def test_num_special_tokens_passed_to_model_inputs(self):
"""Test that num_special_tokens is passed through."""
kwargs = self._create_mock_kwargs(include_distilled_params=True)
input_ids = torch.tensor([[1, 2, 3]])
attention_mask = torch.ones(1, 3)
model_inputs = self.pipeline.prepare_inputs_for_generation(
input_ids=input_ids,
past_key_values=None,
attention_mask=attention_mask,
inputs_embeds=None,
tokenizer_output=None,
batch_gen_image_info=None,
generator=None,
**kwargs,
)
assert "num_special_tokens" in model_inputs
assert model_inputs["num_special_tokens"] == kwargs["num_special_tokens"]
def test_decode_update_keeps_all_special_token_positions(self):
"""Test decode step uses timestep, guidance, and timestep_r positions."""
model_kwargs = self._create_mock_kwargs(include_distilled_params=True)
model_kwargs.update(
{
"tokenizer_output": object(),
"attention_mask": torch.ones(1, 1, 7, 7, dtype=torch.bool),
"image_mask": torch.tensor([[False, False, False, False, True, True, True]]),
"gen_timestep_scatter_index": torch.tensor([[1]]),
"guidance_scatter_index": torch.tensor([[2]]),
"timesteps_r_scatter_index": torch.tensor([[3]]),
}
)
updated = self.pipeline._update_model_kwargs_for_generation({}, model_kwargs)
assert torch.equal(updated["position_ids"], torch.tensor([[1, 2, 3, 4, 5, 6]]))
assert updated["num_special_tokens"] == 3
assert updated["guidance_scatter_index"] is model_kwargs["guidance_scatter_index"]
assert updated["timesteps_r_scatter_index"] is model_kwargs["timesteps_r_scatter_index"]
assert updated["attention_mask"].shape == (1, 1, 6, 7)
def test_distilled_params_none_when_not_provided(self):
"""Test that distilled params are None when not provided."""
kwargs = self._create_mock_kwargs(include_distilled_params=False)
input_ids = torch.tensor([[1, 2, 3]])
attention_mask = torch.ones(1, 3)
model_inputs = self.pipeline.prepare_inputs_for_generation(
input_ids=input_ids,
past_key_values=None,
attention_mask=attention_mask,
inputs_embeds=None,
tokenizer_output=None,
batch_gen_image_info=None,
generator=None,
**kwargs,
)
# All distilled params should be present but None
assert "guidance" in model_inputs
assert model_inputs["guidance"] is None
assert "guidance_scatter_index" in model_inputs
assert model_inputs["guidance_scatter_index"] is None
assert "timesteps_r" in model_inputs
assert model_inputs["timesteps_r"] is None
assert "timesteps_r_scatter_index" in model_inputs
assert model_inputs["timesteps_r_scatter_index"] is None
def test_all_distilled_keys_present_in_model_inputs(self):
"""Test that all 4 distilled keys are always present (even if None)."""
kwargs = self._create_mock_kwargs(include_distilled_params=False)
input_ids = torch.tensor([[1, 2, 3]])
attention_mask = torch.ones(1, 3)
model_inputs = self.pipeline.prepare_inputs_for_generation(
input_ids=input_ids,
past_key_values=None,
attention_mask=attention_mask,
inputs_embeds=None,
tokenizer_output=None,
batch_gen_image_info=None,
generator=None,
**kwargs,
)
required_distilled_keys = [
"guidance",
"guidance_scatter_index",
"timesteps_r",
"timesteps_r_scatter_index",
]
for key in required_distilled_keys:
assert key in model_inputs, f"{key} must be in model_inputs dict"
def test_ar_kv_reuse_rebases_distilled_scatter_indices():
pipeline = object.__new__(HunyuanImage3Text2ImagePipeline)
model_kwargs = {
"query_lens": [8],
"attention_mask": torch.ones(1, 1, 8, 12, dtype=torch.bool),
"position_ids": torch.arange(8).unsqueeze(0),
"image_mask": torch.ones(1, 8, dtype=torch.bool),
"gen_timestep_scatter_index": torch.tensor([[5]]),
"guidance_scatter_index": torch.tensor([[6]]),
"timesteps_r_scatter_index": torch.tensor([[7]]),
}
truncated = pipeline._truncate_reused_prefix(
input_ids=torch.arange(8).unsqueeze(0),
model_kwargs=model_kwargs,
positive_reuse_len=3,
)
assert truncated.tolist() == [[3, 4, 5, 6, 7]]
assert model_kwargs["gen_timestep_scatter_index"].tolist() == [[2]]
assert model_kwargs["guidance_scatter_index"].tolist() == [[3]]
assert model_kwargs["timesteps_r_scatter_index"].tolist() == [[4]]
def test_cfg_parallel_splits_meanflow_scatter_indices():
model_kwargs = {
"timesteps_r_scatter_index": torch.tensor([[1], [2], [3], [4]]),
}
HunyuanImage3Text2ImagePipeline._split_model_kwargs_for_cfg_parallel(
model_kwargs,
batch_size=2,
cfg_rank=1,
)
assert model_kwargs["timesteps_r_scatter_index"].tolist() == [[3], [4]]
@@ -38,7 +38,6 @@ def _pipeline():
cache_backend=None, cache_backend=None,
diffusion_kv_cache_skip_step_indices=None, diffusion_kv_cache_skip_step_indices=None,
) )
pipeline.hf_config = SimpleNamespace(cfg_distilled=False, use_meanflow=False)
pipeline._pipeline = SimpleNamespace() pipeline._pipeline = SimpleNamespace()
return pipeline return pipeline
@@ -452,44 +451,3 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
assert states[1].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (2, 1, 2, 6) assert states[1].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (2, 1, 2, 6)
assert states[0].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(2, 4)], [(2, 4)]] assert states[0].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(2, 4)], [(2, 4)]]
assert states[1].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(4, 6)], [(4, 6)]] assert states[1].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(4, 6)], [(4, 6)]]
def test_distilled_step_supplies_guidance_and_meanflow_timestep(monkeypatch):
pipeline = _pipeline()
pipeline.hf_config = SimpleNamespace(cfg_distilled=True, use_meanflow=True)
monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
state = _state("distilled", 1)
state.latents = torch.zeros(1, 1)
state.extra[_STEP_GUIDANCE_SCALE] = 2.5
state.extra[_STEP_MODEL_KWARGS].update(
{
"attention_mask": torch.ones(1, 1, 2, 4, dtype=torch.bool),
"full_attn_spans": [[(2, 4)]],
}
)
state.extra[_STEP_PROMPT_KV] = [
{
"key": torch.zeros(1, 2, 1, 1),
"value": torch.zeros(1, 2, 1, 1),
"lens": torch.tensor([2]),
}
]
state.scheduler = SimpleNamespace(get_timestep_r=lambda _timestep: torch.tensor(0.25))
captured = {}
pipeline._restore_prompt_kv_cache = lambda *_args: None
def fake_prepare_inputs(input_ids, images, timestep, **model_kwargs):
del input_ids, images, timestep
captured.update(model_kwargs)
return {"model_kwargs": model_kwargs}
pipeline.prepare_inputs_for_generation = fake_prepare_inputs
pipeline.forward_call = lambda **_kwargs: {"diffusion_prediction": torch.tensor([[1.0]])}
pipeline._update_model_kwargs_for_generation = lambda _output, model_kwargs: model_kwargs
output = pipeline.denoise_step(InputBatch.make_batch([state]))
torch.testing.assert_close(output, torch.tensor([[1.0]]))
torch.testing.assert_close(captured["guidance"], torch.tensor([2500.0], dtype=torch.bfloat16))
torch.testing.assert_close(captured["timesteps_r"], torch.tensor([0.25]))
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace from types import SimpleNamespace
@@ -182,8 +182,7 @@ def _fake_ltx_parallel_attention():
return "FLASH_ATTN" return "FLASH_ATTN"
@staticmethod @staticmethod
def supports_attention_mask(attention_spec=None): def supports_attention_mask():
del attention_spec
return True return True
def run(name): def run(name):
@@ -795,10 +795,6 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
def test_packed_attention_skips_mask_for_packed_mask_free_backend(): def test_packed_attention_skips_mask_for_packed_mask_free_backend():
from vllm_omni.diffusion.attention.backends.abstract import (
VideoTokenLayout,
VideoTokenSpan,
)
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention, MiniMaxH3Attention,
) )
@@ -825,10 +821,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
torch.nn.Module.__init__(attention) torch.nn.Module.__init__(attention)
attention.attention = FakeAttention() attention.attention = FakeAttention()
q = torch.randn(8, 2, 4) q = torch.randn(8, 2, 4)
video_layout = VideoTokenLayout(
used_len=5,
video_spans=(VideoTokenSpan(start=3, latent_grid=(1, 1, 2), role="target"),),
)
attention._run_packed_attention( attention._run_packed_attention(
q, q,
@@ -837,9 +829,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
max_seqlen=5, max_seqlen=5,
packed_total=8, packed_total=8,
video_layout=video_layout,
vsa_prefix_segments=(1, 2),
gate_compress=torch.ones_like(q),
) )
metadata = attention.attention.metadata metadata = attention.attention.metadata
@@ -847,11 +836,6 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
assert metadata.attn_mask is None assert metadata.attn_mask is None
assert metadata.extra["valid_kv_length"] == 5 assert metadata.extra["valid_kv_length"] == 5
assert metadata.extra["npu_attn_varlen"] is True assert metadata.extra["npu_attn_varlen"] is True
assert metadata.video_layout is video_layout
assert metadata.extra["vsa_h3_prefix_segments"] == (1, 2)
assert "vsa_h3_video_shape" not in metadata.extra
assert "vsa_h3_target_start" not in metadata.extra
assert metadata.extra["gate_compress"].shape == (1, *q.shape)
packed_padding = metadata.packed_padding packed_padding = metadata.packed_padding
assert packed_padding is not None assert packed_padding is not None
assert packed_padding.q_length == 5 assert packed_padding.q_length == 5
@@ -437,83 +437,6 @@ def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path):
_fusion(tmp_path).check_serving_contract(partition="ref2va", **contract) _fusion(tmp_path).check_serving_contract(partition="ref2va", **contract)
def test_a_vsa_variant_accepts_pure_ulysses_and_refuses_ring(tmp_path):
sparse = tmp_path / "vsa" / "adapter_model.safetensors"
_write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
od_config = SimpleNamespace(
diffusion_attention_backend="FASTVIDEO_VSA",
parallel_config=SimpleNamespace(sequence_parallel_size=8, ulysses_degree=8, ring_degree=1),
)
fusion = _load(sparse.parent)
fusion.check_serving_contract(partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0)
od_config.parallel_config.ring_degree = 2
with pytest.raises(ValueError, match="pure Ulysses"):
fusion.check_serving_contract(partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0)
def test_a_vsa_variant_accepts_normalized_attention_config(tmp_path):
sparse = tmp_path / "vsa" / "adapter_model.safetensors"
_write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
od_config = SimpleNamespace(
diffusion_attention_config=SimpleNamespace(default=SimpleNamespace(backend="FASTVIDEO_VSA")),
parallel_config=SimpleNamespace(sequence_parallel_size=1),
)
_load(sparse.parent).check_serving_contract(
partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0
)
def test_a_vsa_variant_reads_the_backend_the_dit_will_actually_resolve(tmp_path):
# The 50 DiT blocks carry attention role "self", so a per_role entry - not
# the default - decides whether the compression gates ever reach a sparse
# kernel.
sparse = tmp_path / "vsa" / "adapter_model.safetensors"
_write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
fusion = _load(sparse.parent)
contract = {"partition": "fl2va", "video_shift": 12.0, "audio_shift": 3.0}
dense_dit = SimpleNamespace(
diffusion_attention_config=SimpleNamespace(
default=SimpleNamespace(backend="FASTVIDEO_VSA"),
per_role={"self": SimpleNamespace(backend="TRTLLM_ATTN")},
),
parallel_config=SimpleNamespace(sequence_parallel_size=1),
)
with pytest.raises(ValueError, match="TRTLLM_ATTN"):
fusion.check_serving_contract(od_config=dense_dit, **contract)
sparse_dit = SimpleNamespace(
diffusion_attention_config=SimpleNamespace(
default=None,
per_role={"self": SimpleNamespace(backend="FASTVIDEO_VSA")},
),
parallel_config=SimpleNamespace(sequence_parallel_size=1),
)
fusion.check_serving_contract(od_config=sparse_dit, **contract)
def test_a_gate_that_never_reached_the_model_is_refused(tmp_path):
# load_weights only logs a skip for a parameter the model does not have, so
# yielding the gate is not evidence it arrived.
sparse = tmp_path / "vsa" / "adapter_model.safetensors"
_write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
fusion = _load(sparse.parent)
checkpoint = [
("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN))),
("blocks.0.attn.out_proj.weight", torch.zeros((_HIDDEN, _INNER))),
("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))),
("blocks.0.mlp.fc2.weight", torch.zeros((_HIDDEN, _FFN))),
]
gate = "blocks.0.attn.to_gate_compress.weight"
streamed = {name for name, _ in fusion.apply(checkpoint)}
assert gate in streamed
with pytest.raises(FastH3AdapterError, match="never reached the model"):
fusion.validate_fully_applied(streamed - {gate})
fusion.validate_fully_applied(streamed)
def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path): def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path):
# Only low-rank factors are placed into H3's fused QKV and gate/up layouts, # Only low-rank factors are placed into H3's fused QKV and gate/up layouts,
# so a .diff aimed at one would otherwise be added transposed. # so a .diff aimed at one would otherwise be added transposed.
@@ -67,33 +67,6 @@ def test_packed_attention_is_a_regional_compile_boundary():
assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False) assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
def test_denoise_branch_keeps_fast_h3_prefix_geometry_model_local():
from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
MiniMaxH3DenoiseBranch,
)
packed = {
"seq_len": torch.tensor(9),
"img_pos": torch.arange(4),
"audio_pos": torch.empty(0, dtype=torch.long),
"text_pos": torch.arange(2),
"update_mask": torch.ones(4, dtype=torch.bool),
"cu_seqlens": torch.tensor([0, 9, 9], dtype=torch.int32),
"img_position_ids": torch.zeros(9, 3, dtype=torch.long),
"video_spans": ({"start": 5, "latent_grid": (1, 2, 2), "role": "target"},),
}
branch = MiniMaxH3DenoiseBranch(
packed=packed,
text_embeddings=torch.zeros(2, 4),
token_tags=torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2]),
device=torch.device("cpu"),
)
layout = branch.static_kwargs["video_token_layout"]
assert not hasattr(layout, "prefix_segments")
assert branch.static_kwargs["packed_seq_params"]["vsa_prefix_segments"] == (2, 3)
def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims(): def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.layers.rope import RotaryEmbedding
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
@@ -30,7 +30,6 @@ class _FakeCheckpointModel:
def __init__(self): def __init__(self):
self.vae_ratio = RATIO self.vae_ratio = RATIO
self.parallel_tiling = True self.parallel_tiling = True
self.stack_tiling = False
self.decoded_with = [] self.decoded_with = []
def split_tiles(self, input_len, is_decoder=False): def split_tiles(self, input_len, is_decoder=False):
@@ -58,7 +57,6 @@ def _vae(parallel_size, state):
vae.model = _FakeCheckpointModel() vae.model = _FakeCheckpointModel()
vae.remote = remote vae.remote = remote
vae.parallel_size = parallel_size vae.parallel_size = parallel_size
vae.stack_tiling_mode = "false"
return vae return vae
@@ -91,33 +89,6 @@ def test_tile_count_uses_the_decoder_grid_not_the_encoder_one(h, w, decoder_tile
assert as_encoder == encoder_tiles != decoder_tiles assert as_encoder == encoder_tiles != decoder_tiles
@pytest.mark.parametrize(
("parallel_size", "h", "w", "expected"),
[
(1, 48, 84, True),
(2, 48, 84, True),
(4, 48, 84, True),
(8, 48, 84, True),
(16, 48, 84, False),
(4, 16, 16, False),
],
)
def test_auto_stack_tiling_requires_multiple_tiles_per_rank(parallel_size, h, w, expected):
vae = _vae(parallel_size, {})
vae.set_stack_tiling("auto")
assert MiniMaxH3VideoVAE._should_stack_decode_tiles(vae, _latent(h, w)) is expected
@pytest.mark.parametrize(("mode", "expected"), [(False, False), (True, True), ("false", False), ("true", True)])
def test_stack_tiling_configuration_accepts_bool_and_string_modes(mode, expected):
vae = _vae(1, {})
vae.set_stack_tiling(mode)
assert vae.model.stack_tiling is expected
def test_rank_local_tiling_restores_the_group_state(): def test_rank_local_tiling_restores_the_group_state():
state = {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"} state = {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"}
vae = _vae(4, state) vae = _vae(4, state)
@@ -185,7 +156,6 @@ def _dispatch_vae(parallel_size):
seen["sp_enabled"] = state["sp_enabled"] seen["sp_enabled"] = state["sp_enabled"]
seen["sp_process_group"] = state["sp_process_group"] seen["sp_process_group"] = state["sp_process_group"]
seen["parallel_tiling"] = vae.model.parallel_tiling seen["parallel_tiling"] = vae.model.parallel_tiling
seen["stack_tiling"] = vae.model.stack_tiling
# the guard never inspects the decoded shape, so keep it small # the guard never inspects the decoded shape, so keep it small
return torch.zeros(1, 3, 2, 4, 4) return torch.zeros(1, 3, 2, 4, 4)
@@ -242,15 +212,3 @@ def test_decode_latent_leaves_the_group_state_alone_when_the_grid_is_large_enoug
assert seen["sp_enabled"] is True assert seen["sp_enabled"] is True
assert seen["sp_process_group"] is state["sp_process_group"] assert seen["sp_process_group"] is state["sp_process_group"]
assert seen["parallel_tiling"] is True assert seen["parallel_tiling"] is True
def test_decode_latent_restores_native_stack_tiling_state():
vae, _, seen = _dispatch_vae(4)
vae.set_stack_tiling("auto")
# The adapter must not leak its per-request decision into the remote VAE.
vae.model.stack_tiling = True
MiniMaxH3VideoVAE.decode_latent(vae, _latent(16, 16))
assert seen["stack_tiling"] is False
assert vae.model.stack_tiling is True
@@ -1,346 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from types import SimpleNamespace
import cache_dit
import pytest
import torch
from cache_dit import ForwardPattern
from torch import nn
import vllm_omni.diffusion.cache.cachedit.backend as cachedit_backend_module
import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
from tests.diffusion.models.sana_video.test_transformer_sana_video import ( # noqa: F401
_TINY_CONFIG,
_init_distributed,
)
from tests.diffusion.offloader.test_layerwise_backend import DummyEvent, DummyStream, dummy_stream
from vllm_omni.diffusion.attention import selector as attention_selector
from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
from vllm_omni.diffusion.models.sana_video import pipeline_sana_video as pipeline_module
from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import (
SanaVideoPipeline,
_validate_cache_offload_parallelism,
)
from vllm_omni.diffusion.models.sana_video.transformer_sana_video import SanaVideoTransformer3DModel
from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy
from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend, LayerwiseOffloadHook
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@dataclass
class _PipelineWithTransformer:
transformer: nn.Module
def _tiny_transformer(monkeypatch) -> SanaVideoTransformer3DModel:
monkeypatch.setattr(
attention_selector,
"_cached_get_backend_cls",
lambda *_args, **_kwargs: SDPABackend,
)
return SanaVideoTransformer3DModel(**(_TINY_CONFIG | {"num_layers": 2}))
def _record_cache_adapters(monkeypatch):
enabled_adapters = []
original_enable_cache = cachedit_backend_module.cache_dit.enable_cache
def record_enable_cache(block_adapter, **kwargs):
enabled_adapters.append(block_adapter)
return original_enable_cache(block_adapter, **kwargs)
monkeypatch.setattr(cachedit_backend_module.cache_dit, "enable_cache", record_enable_cache)
return enabled_adapters
def test_sana_video_declares_cache_and_layerwise_metadata():
adapter_config = SanaVideoTransformer3DModel._cache_dit_adapter_config
assert adapter_config.block_forward_patterns == {
"transformer_blocks": ForwardPattern.Pattern_3,
}
assert adapter_config.has_separate_cfg is False
assert adapter_config.cached_adapter_cls is None
assert adapter_config.check_forward_pattern is True
assert SanaVideoTransformer3DModel._layerwise_offload_blocks_attrs == ["transformer_blocks"]
assert SanaVideoPipeline.default_num_inference_steps == 50
def test_two_layer_tiny_transformer_enables_cache_dit(monkeypatch):
transformer = _tiny_transformer(monkeypatch)
pipeline = _PipelineWithTransformer(transformer=transformer)
enabled_adapters = _record_cache_adapters(monkeypatch)
backend = CacheDiTBackend()
try:
backend.enable(pipeline)
assert backend.is_enabled()
assert transformer._is_cached is True
assert len(enabled_adapters) == 1
selected_blocks = cache_dit.BlockAdapter.flatten(enabled_adapters[0].blocks)
assert selected_blocks == [transformer.transformer_blocks]
assert len(selected_blocks[0]) == 2
finally:
if enabled_adapters:
cache_dit.disable_cache(enabled_adapters[0])
def test_second_engine_enables_cache_dit_after_undisabled_first(monkeypatch):
"""An engine never disables Cache-DiT on shutdown, and cache_dit marks the
shared wrapper-pipe class as cached; a second engine in the same process
must still get a full enable with a working cache context."""
enabled_adapters = _record_cache_adapters(monkeypatch)
first = _PipelineWithTransformer(transformer=_tiny_transformer(monkeypatch))
second = _PipelineWithTransformer(transformer=_tiny_transformer(monkeypatch))
backend = CacheDiTBackend()
try:
CacheDiTBackend().enable(first)
backend.enable(second)
assert backend.is_enabled()
assert second.transformer._is_cached is True
finally:
for adapter in enabled_adapters:
cache_dit.disable_cache(adapter)
def test_cache_dit_pattern_mismatch_fails_closed(monkeypatch):
transformer = _tiny_transformer(monkeypatch)
transformer.transformer_blocks[1].forward = lambda unexpected_input: unexpected_input
pipeline = _PipelineWithTransformer(transformer=transformer)
backend = CacheDiTBackend()
with pytest.raises(AssertionError, match="No block forward pattern matched"):
backend.enable(pipeline)
assert backend.is_enabled() is False
assert not getattr(transformer, "_is_cached", False)
@pytest.mark.parametrize(
"feature_flags",
[
{"cache_backend": "cache_dit"},
{"enable_cpu_offload": True},
{"enable_layerwise_offload": True},
{"enable_distributed_layerwise_offload": True},
],
)
@pytest.mark.parametrize(
"parallel_config",
[
DiffusionParallelConfig(tensor_parallel_size=2),
DiffusionParallelConfig(cfg_parallel_size=2),
DiffusionParallelConfig(ulysses_degree=2),
],
)
def test_cache_offload_distributed_combinations_fail_closed(feature_flags, parallel_config):
config = OmniDiffusionConfig(parallel_config=parallel_config, **feature_flags)
with pytest.raises(NotImplementedError, match="supported only with TP1, CFG1, and SP1"):
_validate_cache_offload_parallelism(config)
def test_sana_video_rejects_unvalidated_cache_backends():
with pytest.raises(NotImplementedError, match="Cache backend 'tea_cache' is not supported"):
_validate_cache_offload_parallelism(OmniDiffusionConfig(cache_backend="tea_cache"))
def test_sana_video_rejects_cache_dit_with_distributed_layerwise_offload():
config = OmniDiffusionConfig(cache_backend="cache_dit", enable_distributed_layerwise_offload=True)
with pytest.raises(NotImplementedError, match="Cache-DiT with distributed layerwise offload"):
_validate_cache_offload_parallelism(config)
def test_sana_video_accepts_distributed_layerwise_offload():
_validate_cache_offload_parallelism(OmniDiffusionConfig(enable_distributed_layerwise_offload=True))
@pytest.mark.parametrize(
("od_config", "match"),
[
(
OmniDiffusionConfig(
cache_backend="cache_dit",
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
),
"tensor_parallel_size",
),
(
OmniDiffusionConfig(cache_backend="cache_dit", enable_distributed_layerwise_offload=True),
"Cache-DiT with distributed layerwise offload",
),
],
)
def test_validation_fails_before_component_loading(monkeypatch, od_config, match):
load_calls = []
monkeypatch.setattr(pipeline_module, "get_local_device", lambda: torch.device("cpu"))
monkeypatch.setattr(
SanaVideoPipeline,
"_load_components",
lambda *args, **kwargs: load_calls.append((args, kwargs)),
)
with pytest.raises(NotImplementedError, match=match):
SanaVideoPipeline(od_config=od_config)
assert load_calls == []
class _TrackingModule:
def __init__(self):
self.to_calls = []
def to(self, *args, **kwargs):
self.to_calls.append((args, kwargs))
return self
def test_component_loading_uses_loader_device_not_runtime_device(monkeypatch):
component_load_device = torch.device("cpu")
text_encoder = _TrackingModule()
vae = _TrackingModule()
transformer = _TrackingModule()
tokenizer = object()
scheduler = object()
loaded_components = iter([text_encoder, vae])
monkeypatch.setattr(torch, "get_default_device", lambda: component_load_device)
monkeypatch.setattr(pipeline_module, "prefetch_subfolders", lambda *args, **kwargs: None)
monkeypatch.setattr(pipeline_module, "_load_sana_tokenizer", lambda *args, **kwargs: tokenizer)
monkeypatch.setattr(
pipeline_module,
"from_pretrained_with_prefetch",
lambda *args, **kwargs: next(loaded_components),
)
monkeypatch.setattr(
pipeline_module,
"_load_json",
lambda _model, filename, _local: ({"vae": [None, "FakeVAE"]} if filename == "model_index.json" else {}),
)
monkeypatch.setattr(
pipeline_module,
"_resolve_vae_class_and_dtype",
lambda *_args: (SimpleNamespace(from_pretrained=None), torch.float32),
)
monkeypatch.setattr(
pipeline_module.SanaVideoTransformer3DModel,
"from_config",
classmethod(lambda _cls, _config: transformer),
)
monkeypatch.setattr(
pipeline_module.DPMSolverMultistepScheduler,
"from_pretrained",
lambda *args, **kwargs: scheduler,
)
pipeline = object.__new__(SanaVideoPipeline)
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cuda:7")
pipeline.weights_sources = []
config = OmniDiffusionConfig(model="local-sana-video", dtype=torch.bfloat16)
loaded = pipeline._load_components(config, prefix="")
assert loaded == (tokenizer, text_encoder, vae, transformer, scheduler)
assert pipeline.device == torch.device("cuda:7")
assert text_encoder.to_calls == [((component_load_device,), {})]
assert vae.to_calls == [((component_load_device,), {})]
assert transformer.to_calls == [
((), {"dtype": torch.bfloat16, "device": component_load_device}),
]
def _patch_layerwise_platform(monkeypatch) -> None:
monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
monkeypatch.setattr(
layerwise_backend_module.current_omni_platform,
"current_stream",
lambda: DummyStream(),
)
monkeypatch.setattr(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
def _tiny_pipeline(monkeypatch) -> SanaVideoPipeline:
pipeline = object.__new__(SanaVideoPipeline)
nn.Module.__init__(pipeline)
pipeline.transformer = _tiny_transformer(monkeypatch)
pipeline.text_encoder = nn.Linear(2, 2)
pipeline.vae = nn.Linear(2, 2)
return pipeline
def _enable_layerwise_offload(pipeline: SanaVideoPipeline) -> LayerWiseOffloadBackend:
backend = LayerWiseOffloadBackend(
OffloadConfig(strategy=OffloadStrategy.LAYER_WISE, pin_cpu_memory=False),
device=torch.device("cpu"),
)
backend.enable(pipeline)
return backend
def test_two_layer_offload_recovers_from_cached_block_skip_and_cleans_up(monkeypatch):
_patch_layerwise_platform(monkeypatch)
pipeline = _tiny_pipeline(monkeypatch)
discovered = ModuleDiscovery.discover(pipeline)
assert discovered.dits == [pipeline.transformer]
assert discovered.encoders == [pipeline.text_encoder]
assert discovered.vaes == [pipeline.vae]
backend = _enable_layerwise_offload(pipeline)
blocks = list(pipeline.transformer.transformer_blocks)
assert backend.is_enabled()
assert len(blocks) == 2
block_1_hook = blocks[1]._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME)
assert block_1_hook is not None
assert block_1_hook.is_materialized is False
# Cache-DiT may skip block 0, leaving block 1 without the normal prefetch.
block_1_hook.pre_forward(blocks[1])
assert block_1_hook.is_materialized is True
block_1_hook.post_forward(blocks[1], None)
assert block_1_hook.is_materialized is False
backend.disable()
assert backend.is_enabled() is False
for block in blocks:
assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is None
def test_layerwise_offload_then_cache_dit_install_and_cleanup(monkeypatch):
_patch_layerwise_platform(monkeypatch)
pipeline = _tiny_pipeline(monkeypatch)
offload_backend = _enable_layerwise_offload(pipeline)
enabled_adapters = _record_cache_adapters(monkeypatch)
cache_backend = CacheDiTBackend()
blocks = list(pipeline.transformer.transformer_blocks)
try:
cache_backend.enable(pipeline)
assert offload_backend.is_enabled()
assert cache_backend.is_enabled()
assert pipeline.transformer._is_cached is True
for block in blocks:
assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is not None
finally:
if enabled_adapters:
cache_dit.disable_cache(enabled_adapters[0])
offload_backend.disable()
assert not getattr(pipeline.transformer, "_is_cached", False)
assert offload_backend.is_enabled() is False
for block in blocks:
assert block._hook_registry.get_hook(LayerwiseOffloadHook._HOOK_NAME) is None
@@ -1,582 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""CFG dispatch tests for the native SANA-Video pipelines.
These exercise the CFG world-size-1 batch path, the CFG world-size-2 dispatch,
and the fp32 combine semantics on CPU / mocked distributed. The real all-gather
path is covered by the multi-process suite, not here.
"""
import pytest
import torch
from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import SanaVideoPipeline
from vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v import SanaImageToVideoPipeline
from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
SanaVideoTransformerOutput,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_predict_noise_calls_transformer_with_return_dict_false():
"""predict_noise must call the transformer with return_dict=False.
The transformer returns a SanaVideoTransformerOutput dataclass when
return_dict is True (its default), which is not subscriptable. The CFG
mixin's default predict_noise does ``result[0]`` and would therefore
raise TypeError on SANA; the pipeline override must force return_dict=False
so the tuple form is returned and indexed.
"""
captured = {}
expected = torch.full((1, 2, 3), 7.0)
class _SpyTransformer:
def __call__(self, *args, return_dict=True, **kwargs):
captured["return_dict"] = return_dict
# Mirror the real transformer: tuple on False, dataclass on True.
return (expected,) if not return_dict else SanaVideoTransformerOutput(sample=expected)
pipe = SanaVideoPipeline.__new__(SanaVideoPipeline) # bypass heavy __init__
pipe.transformer = _SpyTransformer()
out = pipe.predict_noise(
torch.zeros(1, 2, 3),
encoder_hidden_states=torch.zeros(1, 1, 4),
encoder_attention_mask=None,
timestep=torch.zeros(1),
)
assert captured["return_dict"] is False
assert torch.equal(out, expected)
class _IdentityScheduler:
"""Scheduler stub whose step is a no-op, so diffuse tests isolate the CFG
dispatch from real scheduler math."""
order = 1
def step(self, noise_pred, t, latents, **kwargs):
return (latents,)
def _set_cfg_world_size(monkeypatch, n):
"""Pipelines always run inside an initialized CFG group; unit tests stub its
world size instead of spinning up torch.distributed. Patch both pipeline
modules since each resolves the symbol in its own namespace."""
import vllm_omni.diffusion.distributed.cfg_parallel as cfg
import vllm_omni.diffusion.models.sana_video.pipeline_sana_video as t2v
import vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v as i2v
monkeypatch.setattr(t2v, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
monkeypatch.setattr(i2v, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
monkeypatch.setattr(cfg, "get_classifier_free_guidance_world_size", lambda: n, raising=False)
monkeypatch.setattr(t2v, "model_parallel_is_initialized", lambda: True, raising=False)
monkeypatch.setattr(i2v, "model_parallel_is_initialized", lambda: True, raising=False)
def _bare_t2v_pipeline(transformer, guidance_scale):
pipe = SanaVideoPipeline.__new__(SanaVideoPipeline) # bypass heavy __init__
pipe.transformer = transformer
pipe.scheduler = _IdentityScheduler()
pipe._interrupt = False
pipe._guidance_scale = guidance_scale
return pipe
def test_cfg_world_size_1_does_single_batch2_forward_per_step(monkeypatch):
"""CFG world size 1 with guidance must do exactly one batch-2 transformer
forward per step (the batch path), not the mixin's sequential two-forward
path."""
_set_cfg_world_size(monkeypatch, 1)
batch_dims = []
class _SpyTransformer:
dtype = torch.float32
def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
batch_dims.append(hidden_states.shape[0])
# out_channels == in_channels here (no learned-sigma split).
return (torch.zeros_like(hidden_states),)
b = 1
pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
out = pipe.diffuse(
latents=torch.randn(b, 4, 2, 4, 4),
timesteps=torch.tensor([1000.0]),
prompt_embeds=torch.randn(b, 3, 8),
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=torch.randn(b, 3, 8),
negative_prompt_attention_mask=torch.ones(b, 3),
guidance_scale=6.0,
extra_step_kwargs={},
dtype=torch.float32,
output_slice=None,
)
assert len(batch_dims) == 1, f"expected 1 forward this step, got {len(batch_dims)}"
assert batch_dims[0] == 2 * b, f"expected batch-2 (neg+pos) forward, got batch {batch_dims[0]}"
assert out.shape == (b, 4, 2, 4, 4)
class _DeterministicTransformer:
"""Learned-sigma transformer whose output depends on the row's encoder
embeds and timestep, so uncond != text and the CFG combine is non-trivial."""
dtype = torch.float32
def __init__(self, latent_channels):
self.out_channels = 2 * latent_channels # learned sigma
def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
b, c = hidden_states.shape[0], hidden_states.shape[1]
sig = encoder_hidden_states.float().mean(dim=(1, 2)).view(b, 1, 1, 1, 1)
base = hidden_states.float() + sig + timestep.float().view(b, 1, 1, 1, 1)
return (base.repeat(1, self.out_channels // c, 1, 1, 1),)
class _LinearScheduler:
"""Deterministic stepping scheduler so multi-step divergence compounds."""
order = 1
def step(self, noise_pred, t, latents, return_dict=False, **kwargs):
return (latents - 0.1 * noise_pred,)
def _pr1_reference_loop(
transformer,
scheduler,
latents,
timesteps,
prompt_embeds,
prompt_mask,
negative_embeds,
negative_mask,
guidance_scale,
latent_channels,
dtype,
):
"""The original CFG1 denoise math, verbatim, as the parity oracle."""
prompt_embeds = torch.cat([negative_embeds, prompt_embeds], dim=0)
prompt_mask = torch.cat([negative_mask, prompt_mask], dim=0)
for t in timesteps:
latent_model_input = torch.cat([latents] * 2)
timestep = t.expand(latent_model_input.shape[0])
noise_pred = transformer(
latent_model_input.to(dtype=dtype),
encoder_hidden_states=prompt_embeds.to(dtype=dtype),
encoder_attention_mask=prompt_mask,
timestep=timestep,
return_dict=False,
)[0]
noise_pred = noise_pred.float()
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
if transformer.out_channels // 2 == latent_channels:
noise_pred = noise_pred.chunk(2, dim=1)[0]
latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0]
return latents
def test_cfg1_refactored_diffuse_is_bit_identical_to_pr1_loop(monkeypatch):
"""The refactored diffuse must reproduce the original CFG1 latents bit-for-bit
across multiple steps, preserving the TP1/CFG1 golden output."""
_set_cfg_world_size(monkeypatch, 1)
torch.manual_seed(0)
b, c = 2, 4
latents0 = torch.randn(b, c, 2, 4, 4)
timesteps = torch.tensor([900.0, 600.0, 300.0])
prompt_embeds = torch.randn(b, 3, 8)
negative_embeds = torch.randn(b, 3, 8)
prompt_mask = torch.ones(b, 3)
negative_mask = torch.ones(b, 3)
guidance_scale = 6.0
transformer = _DeterministicTransformer(latent_channels=c)
scheduler = _LinearScheduler()
ref = _pr1_reference_loop(
transformer,
scheduler,
latents0.clone(),
timesteps,
prompt_embeds,
prompt_mask,
negative_embeds,
negative_mask,
guidance_scale,
latent_channels=c,
dtype=torch.float32,
)
pipe = _bare_t2v_pipeline(transformer, guidance_scale=guidance_scale)
pipe.scheduler = scheduler
got = pipe.diffuse(
latents=latents0.clone(),
timesteps=timesteps,
prompt_embeds=prompt_embeds,
prompt_attention_mask=prompt_mask,
negative_prompt_embeds=negative_embeds,
negative_prompt_attention_mask=negative_mask,
guidance_scale=guidance_scale,
extra_step_kwargs={},
dtype=torch.float32,
output_slice=c,
)
assert torch.equal(got, ref)
def test_combine_cfg_noise_upcasts_to_fp32_before_combine():
"""combine_cfg_noise must upcast positive/negative to fp32 before the CFG
formula, so the CFG2 (bf16 all-gather) combine is bit-identical to CFG1's
fp32 combine."""
pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)
torch.manual_seed(0)
pos = torch.randn(1, 4, 2, 4, 4).bfloat16()
neg = torch.randn(1, 4, 2, 4, 4).bfloat16()
scale = 6.0
out = pipe.combine_cfg_noise(pos, neg, scale, cfg_normalize=False)
expected = neg.float() + scale * (pos.float() - neg.float())
assert out.dtype == torch.float32
assert torch.equal(out, expected)
def test_combine_cfg_noise_accepts_tuple_predictions_from_gather():
"""The CFG-parallel gather path hands combine_cfg_noise single-element
tuples, not bare tensors; the override must upcast and combine them."""
pipe = SanaVideoPipeline.__new__(SanaVideoPipeline)
torch.manual_seed(0)
pos = torch.randn(1, 4, 2, 4, 4).bfloat16()
neg = torch.randn(1, 4, 2, 4, 4).bfloat16()
scale = 6.0
out = pipe.combine_cfg_noise((pos,), (neg,), scale, cfg_normalize=False)
expected = neg.float() + scale * (pos.float() - neg.float())
assert out.dtype == torch.float32
assert torch.equal(out, expected)
def test_cfg2_dispatch_builds_independent_branch_kwargs(monkeypatch):
"""In CFG world size 2, diffuse must build separate batch-B positive/negative
kwargs (no pre-cat), and call the mixin with cfg_normalize=False and the
learned-sigma output_slice."""
_set_cfg_world_size(monkeypatch, 2)
b, c = 1, 4
captured = {}
def _spy_cfg(**kwargs):
captured.update(kwargs)
return torch.zeros(b, c, 2, 4, 4)
prompt_embeds = torch.randn(b, 3, 8)
negative_embeds = torch.randn(b, 3, 8)
dtype = torch.float32
pipe = _bare_t2v_pipeline(transformer=None, guidance_scale=6.0)
pipe.predict_noise_maybe_with_cfg = _spy_cfg
pipe.diffuse(
latents=torch.randn(b, c, 2, 4, 4),
timesteps=torch.tensor([1000.0]),
prompt_embeds=prompt_embeds,
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=negative_embeds,
negative_prompt_attention_mask=torch.zeros(b, 3),
guidance_scale=6.0,
extra_step_kwargs={},
dtype=dtype,
output_slice=c,
)
assert captured, "predict_noise_maybe_with_cfg was not called in CFG world size 2"
assert captured["do_true_cfg"] is True
assert captured["true_cfg_scale"] == 6.0
assert captured["cfg_normalize"] is False
assert captured["output_slice"] == c
pos = captured["positive_kwargs"]
neg = captured["negative_kwargs"]
assert pos["hidden_states"].shape[0] == b, "positive branch must be batch-B, not pre-concatenated"
assert neg["hidden_states"].shape[0] == b, "negative branch must be batch-B, not pre-concatenated"
assert pos["timestep"].shape[0] == b
assert torch.equal(pos["encoder_hidden_states"], prompt_embeds.to(dtype))
assert torch.equal(neg["encoder_hidden_states"], negative_embeds.to(dtype))
# ── I2V CFG dispatch ──
class _I2VSpyScheduler:
"""Records the frame count it steps on and mutates non-first frames so
first-frame preservation is observable."""
order = 1
def __init__(self):
self.stepped_frame_counts = []
def step(self, noise_pred, t, latents, return_dict=False, **kwargs):
self.stepped_frame_counts.append(latents.shape[2])
return (latents + 1.0,)
def _bare_i2v_pipeline(transformer, scheduler, guidance_scale):
pipe = SanaImageToVideoPipeline.__new__(SanaImageToVideoPipeline)
pipe.transformer = transformer
pipe.scheduler = scheduler
pipe._interrupt = False
pipe._guidance_scale = guidance_scale
return pipe
def _conditioning_mask(b, frames):
mask = torch.zeros(b, 1, frames, 1, 1)
mask[:, :, 0] = 1
return mask
def test_i2v_cfg1_preserves_first_frame_and_steps_only_tail(monkeypatch):
"""I2V CFG1 must keep the conditioning first frame bit-exact and only run the
scheduler on frames [1:]."""
_set_cfg_world_size(monkeypatch, 1)
b, c, frames = 1, 4, 3
class _SpyTransformer:
dtype = torch.float32
out_channels = 2 * c
def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
return (torch.zeros(hidden_states.shape[0], 2 * c, frames, 4, 4),)
scheduler = _I2VSpyScheduler()
pipe = _bare_i2v_pipeline(_SpyTransformer(), scheduler, guidance_scale=6.0)
latents0 = torch.randn(b, c, frames, 4, 4)
out = pipe.diffuse(
latents=latents0.clone(),
timesteps=torch.tensor([1000.0]),
prompt_embeds=torch.randn(b, 3, 8),
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=torch.randn(b, 3, 8),
negative_prompt_attention_mask=torch.ones(b, 3),
guidance_scale=6.0,
conditioning_mask=_conditioning_mask(b, frames),
extra_step_kwargs={},
dtype=torch.float32,
output_slice=c,
)
assert torch.equal(out[:, :, :1], latents0[:, :, :1]), "first (conditioning) frame must be unchanged"
assert torch.equal(out[:, :, 1:], latents0[:, :, 1:] + 1.0), "tail frames must be the scheduler output"
assert scheduler.stepped_frame_counts == [frames - 1], "scheduler must step only on frames [1:]"
def test_i2v_cfg2_dispatch_builds_masked_branch_kwargs(monkeypatch):
"""I2V CFG world size 2 must build batch-B per-branch kwargs with the
conditioning-frame timestep zeroed (mask not concatenated) and preserve the
first frame."""
_set_cfg_world_size(monkeypatch, 2)
b, c, frames = 1, 4, 3
captured = {}
def _spy_cfg(**kwargs):
captured.update(kwargs)
return torch.zeros(b, c, frames, 4, 4)
scheduler = _I2VSpyScheduler()
pipe = _bare_i2v_pipeline(transformer=None, scheduler=scheduler, guidance_scale=6.0)
pipe.predict_noise_maybe_with_cfg = _spy_cfg
prompt_embeds = torch.randn(b, 3, 8)
negative_embeds = torch.randn(b, 3, 8)
latents0 = torch.randn(b, c, frames, 4, 4)
mask = _conditioning_mask(b, frames)
out = pipe.diffuse(
latents=latents0.clone(),
timesteps=torch.tensor([1000.0]),
prompt_embeds=prompt_embeds,
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=negative_embeds,
negative_prompt_attention_mask=torch.zeros(b, 3),
guidance_scale=6.0,
conditioning_mask=mask,
extra_step_kwargs={},
dtype=torch.float32,
output_slice=c,
)
assert captured["do_true_cfg"] is True
assert captured["cfg_normalize"] is False
assert captured["output_slice"] == c
pos = captured["positive_kwargs"]
neg = captured["negative_kwargs"]
assert pos["hidden_states"].shape[0] == b, "positive branch must be batch-B"
assert neg["hidden_states"].shape[0] == b, "negative branch must be batch-B"
assert tuple(pos["timestep"].shape) == tuple(mask.shape), "mask must not be concatenated in CFG2"
assert torch.equal(pos["timestep"][:, :, 0], torch.zeros(b, 1, 1, 1)), "conditioning frame timestep must be 0"
assert torch.equal(pos["encoder_hidden_states"], prompt_embeds)
assert torch.equal(neg["encoder_hidden_states"], negative_embeds)
assert torch.equal(out[:, :, :1], latents0[:, :, :1]), "first frame must be preserved"
def test_diffuse_runs_cfg_parallel_validity_check(monkeypatch, mocker):
"""diffuse must run the CFG-parallel validity check so a scale<=1 or missing
negative prompt warns the user instead of silently wasting a rank."""
_set_cfg_world_size(monkeypatch, 2)
b = 1
pipe = _bare_t2v_pipeline(transformer=None, guidance_scale=6.0)
pipe.predict_noise_maybe_with_cfg = lambda **kwargs: torch.zeros(b, 4, 2, 4, 4)
check = mocker.patch.object(SanaVideoPipeline, "check_cfg_parallel_validity", return_value=True)
pipe.diffuse(
latents=torch.randn(b, 4, 2, 4, 4),
timesteps=torch.tensor([1000.0]),
prompt_embeds=torch.randn(b, 3, 8),
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=torch.randn(b, 3, 8),
negative_prompt_attention_mask=torch.ones(b, 3),
guidance_scale=6.0,
extra_step_kwargs={},
dtype=torch.float32,
output_slice=None,
)
check.assert_called_once_with(6.0, True)
def test_i2v_diffuse_runs_cfg_parallel_validity_check(monkeypatch, mocker):
"""I2V has its own diffuse; it must also run the CFG-parallel validity check."""
_set_cfg_world_size(monkeypatch, 2)
b, c, frames = 1, 4, 3
pipe = _bare_i2v_pipeline(None, _I2VSpyScheduler(), guidance_scale=6.0)
pipe.predict_noise_maybe_with_cfg = lambda **kwargs: torch.zeros(b, c, frames, 4, 4)
check = mocker.patch.object(SanaVideoPipeline, "check_cfg_parallel_validity", return_value=True)
pipe.diffuse(
latents=torch.randn(b, c, frames, 4, 4),
timesteps=torch.tensor([1000.0]),
prompt_embeds=torch.randn(b, 3, 8),
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=torch.randn(b, 3, 8),
negative_prompt_attention_mask=torch.ones(b, 3),
guidance_scale=6.0,
conditioning_mask=_conditioning_mask(b, frames),
extra_step_kwargs={},
dtype=torch.float32,
output_slice=c,
)
check.assert_called_once_with(6.0, True)
def test_diffuse_runs_without_initialized_parallel_groups():
"""Single-process runs (unit tests, offline scripts) never initialize the
distributed groups; diffuse must fall back to the serial path instead of
tripping the parallel-state assertion. No world-size patching on purpose."""
class _SpyTransformer:
dtype = torch.float32
def __call__(self, hidden_states, *, encoder_hidden_states, encoder_attention_mask, timestep, return_dict):
return (torch.zeros_like(hidden_states),)
b = 1
pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
out = pipe.diffuse(
latents=torch.randn(b, 4, 2, 4, 4),
timesteps=torch.tensor([1000.0]),
prompt_embeds=torch.randn(b, 3, 8),
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=torch.randn(b, 3, 8),
negative_prompt_attention_mask=torch.ones(b, 3),
guidance_scale=6.0,
extra_step_kwargs={},
dtype=torch.float32,
output_slice=None,
)
assert out.shape == (b, 4, 2, 4, 4)
def test_diffuse_propagates_group_error_when_parallel_initialized(monkeypatch):
"""Inside an initialized parallel run a failing CFG-group lookup is a real
error and must propagate instead of silently flipping to the serial path."""
import vllm_omni.diffusion.models.sana_video.pipeline_sana_video as t2v
monkeypatch.setattr(t2v, "model_parallel_is_initialized", lambda: True, raising=False)
def _broken_group_lookup():
raise AssertionError("classifier_free_guidance group is not initialized")
monkeypatch.setattr(t2v, "get_classifier_free_guidance_world_size", _broken_group_lookup)
b = 1
pipe = _bare_t2v_pipeline(None, guidance_scale=6.0)
with pytest.raises(AssertionError, match="classifier_free_guidance"):
pipe.diffuse(
latents=torch.randn(b, 4, 2, 4, 4),
timesteps=torch.tensor([1000.0]),
prompt_embeds=torch.randn(b, 3, 8),
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=torch.randn(b, 3, 8),
negative_prompt_attention_mask=torch.ones(b, 3),
guidance_scale=6.0,
extra_step_kwargs={},
dtype=torch.float32,
output_slice=None,
)
def test_t2v_forwards_eta_and_generator_into_scheduler(monkeypatch):
"""T2V must keep passing eta/generator (extra_step_kwargs) into scheduler.step."""
_set_cfg_world_size(monkeypatch, 1)
captured = {}
class _KwargScheduler:
order = 1
def step(self, noise_pred, t, latents, **kwargs):
captured.update(kwargs)
return (latents,)
class _SpyTransformer:
dtype = torch.float32
def __call__(self, hidden_states, **kwargs):
return (torch.zeros_like(hidden_states),)
b = 1
pipe = _bare_t2v_pipeline(_SpyTransformer(), guidance_scale=6.0)
pipe.scheduler = _KwargScheduler()
generator = torch.Generator().manual_seed(0)
pipe.diffuse(
latents=torch.randn(b, 4, 2, 4, 4),
timesteps=torch.tensor([1000.0]),
prompt_embeds=torch.randn(b, 3, 8),
prompt_attention_mask=torch.ones(b, 3),
negative_prompt_embeds=torch.randn(b, 3, 8),
negative_prompt_attention_mask=torch.ones(b, 3),
guidance_scale=6.0,
extra_step_kwargs={"eta": 0.5, "generator": generator},
dtype=torch.float32,
output_slice=None,
)
assert captured.get("eta") == 0.5
assert captured.get("generator") is generator
assert captured.get("return_dict") is False
@@ -1,90 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Topology validation for SANA-Video parallelism.
The native pipeline only supports tensor/CFG parallel of size 1 or 2 and rejects
sequence parallel and text-encoder TP before any large weight is allocated.
"""
from types import SimpleNamespace
import pytest
from vllm_omni.diffusion.models.sana_video.pipeline_sana_video import SanaVideoPipeline
from vllm_omni.diffusion.models.sana_video.pipeline_sana_video_i2v import SanaImageToVideoPipeline
from vllm_omni.diffusion.models.sana_video.transformer_sana_video import (
validate_sana_video_parallel_config,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _pc(
tp: int = 1,
cfg: int = 1,
sp: int | None = None,
text_tp: int = 1,
pp: int = 1,
hsdp: bool = False,
) -> SimpleNamespace:
return SimpleNamespace(
tensor_parallel_size=tp,
cfg_parallel_size=cfg,
sequence_parallel_size=sp,
text_encoder_tp_size=text_tp,
pipeline_parallel_size=pp,
use_hsdp=hsdp,
)
@pytest.mark.parametrize("tp", [1, 2])
def test_supported_tp_sizes_pass(tp: int) -> None:
validate_sana_video_parallel_config(_pc(tp=tp))
@pytest.mark.parametrize("tp", [3, 4, 8])
def test_unsupported_tp_sizes_rejected(tp: int) -> None:
with pytest.raises(NotImplementedError, match="tensor_parallel_size"):
validate_sana_video_parallel_config(_pc(tp=tp))
def test_cfg3_rejected() -> None:
with pytest.raises(NotImplementedError, match="cfg_parallel_size"):
validate_sana_video_parallel_config(_pc(cfg=3))
def test_sequence_parallel_rejected() -> None:
with pytest.raises(NotImplementedError, match="[Ss]equence parallel"):
validate_sana_video_parallel_config(_pc(sp=2))
def test_pipeline_parallel_rejected() -> None:
with pytest.raises(NotImplementedError, match="pipeline parallel"):
validate_sana_video_parallel_config(_pc(pp=2))
def test_hsdp_rejected() -> None:
with pytest.raises(NotImplementedError, match="HSDP"):
validate_sana_video_parallel_config(_pc(hsdp=True))
def test_text_encoder_tp_rejected() -> None:
with pytest.raises(NotImplementedError, match="text encoder"):
validate_sana_video_parallel_config(_pc(text_tp=2))
def test_tp2_cfg2_combo_passes() -> None:
validate_sana_video_parallel_config(_pc(tp=2, cfg=2))
@pytest.mark.parametrize("pipeline_cls", [SanaVideoPipeline, SanaImageToVideoPipeline])
def test_pipeline_rejects_unsupported_topology_before_loading(pipeline_cls, mocker) -> None:
"""The pipeline must reject an unsupported topology before it loads weights."""
load = mocker.patch.object(SanaVideoPipeline, "_load_components")
od_config = SimpleNamespace(parallel_config=_pc(tp=3))
with pytest.raises(NotImplementedError, match="tensor_parallel_size"):
pipeline_cls(od_config=od_config)
load.assert_not_called()
Loaded 100 of 259 files, more files were not shown because too many files have changed in this diff. Show more