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

|
||||
|
||||
FlashAttention computes attention in a single streaming pass over KV tiles. Per query row the
|
||||
kernel maintains three running values:
|
||||
|
||||
- `m` — the largest score seen so far,
|
||||
- `l` — the running denominator `Σ exp(sⱼ − m)`,
|
||||
- `O` — the running numerator `Σ exp(sⱼ − m)·vⱼ`,
|
||||
|
||||
and returns `O / l` at the end. For each key tile it computes the tile's scores `Q · K_jᵀ · scale`,
|
||||
updates `m`, and accumulates that tile's contribution into `l` and `O`. Rescaling `l` and `O` when
|
||||
`m` grows keeps the single pass numerically exact.
|
||||
and returns `O / l` at the end. For each KV tile it computes the scores `QK_j^T`, updates `m`, and
|
||||
accumulates the tile's contribution into `l` and `O`. Rescaling `l` and `O` when `m` grows keeps the
|
||||
dense pass numerically exact.
|
||||
|
||||
## The skip test
|
||||
|
||||
Once a tile's scores are known, its largest score `tile_max` is compared against the running
|
||||
maximum:
|
||||
A tile is a block of query rows against a block of keys, and the kernel processes it as a unit.
|
||||
BLASST inserts one test between computing a tile's scores and accumulating them. For each query
|
||||
row `i` in the tile, let `tile_max[i]` be the row's largest score within the tile, `running_max[i]`
|
||||
the row's running maximum, and `λ` the effective threshold:
|
||||
|
||||
```text
|
||||
if exp(tile_max − running_max) < threshold:
|
||||
skip this tile # do not compute its softmax weights or its P·V accumulation
|
||||
if all(exp(tile_max[i] - running_max[i]) < λ for every row i in the tile):
|
||||
skip this tile # do not compute its Softmax or PV contribution
|
||||
```
|
||||
|
||||
`exp(tile_max − running_max)` is an upper bound on the softmax weight any key in the tile can
|
||||
receive: if even the tile's best key is far below the current maximum, every key in the tile is
|
||||
negligible, and both the softmax (the exponentials) and the `P·V` accumulation for that tile can be
|
||||
dropped. The tile's contribution to `l` and `O` is simply skipped.
|
||||
`exp(tile_max[i] − running_max[i])` is an upper bound on the softmax weight any key in the tile can
|
||||
receive in row `i`: if even the row's best key in this tile is far below the row's current maximum,
|
||||
the tile is unimportant for that row. Only when this holds for every row in the tile can the
|
||||
Softmax and `PV` work for the tile be skipped, in which case its contribution to `l` and `O` is
|
||||
simply omitted. A larger `λ` makes the test more aggressive. The figure states the same test in
|
||||
log space, `m̃ − m < ln(λ)`, which is how the kernel evaluates it.
|
||||
|
||||
## What this bounds
|
||||
The test needs the tile's scores, so `QK_j^T` always runs and only the Softmax and `PV` work can
|
||||
be skipped. Since the two matmuls have the same FLOP count, skipping every eligible tile removes at
|
||||
most roughly half of the attention arithmetic, and the kernel-level speedup stays well under 2×.
|
||||
|
||||
Two properties of the test shape the achievable speedup:
|
||||
## From configuration to the kernel threshold
|
||||
|
||||
- **`Q · K_jᵀ` always runs.** The test needs `tile_max`, which comes from the tile's scores, so the
|
||||
score matmul is never skipped — only the softmax and the `P·V` accumulation are. The score matmul
|
||||
and the value matmul are comparable in cost, so skipping every eligible tile removes at most
|
||||
roughly half the attention work; the kernel-level speedup is bounded well under 2×, not unbounded.
|
||||
|
||||
- **The decision is per tile, not per key.** A tile is skipped only when *all* of its keys are
|
||||
collectively negligible; a single important key keeps the whole tile. How many tiles actually
|
||||
qualify depends on the data and rounds down to tile granularity, so `target_sparsity` selects an
|
||||
operating point on a calibrated curve — it is not a promise that a fixed fraction of tiles is
|
||||
skipped.
|
||||
|
||||
## The threshold
|
||||
|
||||
The per-tile threshold is normalized by sequence length:
|
||||
The FlashInfer kernel does not take `λ` directly. It takes a `threshold_scale_factor` and divides
|
||||
it by the KV sequence length of the call:
|
||||
|
||||
```text
|
||||
threshold = factor / seqlen
|
||||
λ = threshold_scale_factor / max_kv_len
|
||||
```
|
||||
|
||||
`factor` comes from one of two sources:
|
||||
vLLM-Omni resolves the factor from one of two mutually exclusive user controls:
|
||||
|
||||
- **Calibrated** — `factor = a · exp(b · target_sparsity)`, where `a` and `b` are fit per model (and
|
||||
per expert, for a multi-expert model) so that a given `target_sparsity` lands near that fraction of
|
||||
skipped tiles on the calibration data.
|
||||
- **Direct** — `factor = skip_softmax.threshold · seqlen`, the calibration-free path; the user sets
|
||||
the threshold themselves.
|
||||
| Control | Factor passed to the kernel | Resulting `λ` |
|
||||
| --- | --- | --- |
|
||||
| `threshold` | `threshold * max_kv_len` | `threshold` |
|
||||
| `target_sparsity=s` | `a * exp(b * s)` | `a * exp(b * s) / max_kv_len` |
|
||||
|
||||
`threshold` is therefore `λ` itself and is independent of sequence length: the same value yields
|
||||
the same per-tile test at any resolution or frame count. `threshold=0` skips no tiles. Because the
|
||||
left-hand side of the skip test lies in `(0, 1]`, values in `(0, 1)` are the meaningful range; the
|
||||
schema accepts any finite non-negative value.
|
||||
|
||||
`target_sparsity` selects a point on a curve fitted per model by
|
||||
[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity),
|
||||
so that `s` lands near that fraction of skipped tiles on the calibration data. The coefficients
|
||||
`a` and `b` come from the checkpoint; the achieved sparsity on another prompt, shape, or layer can
|
||||
differ.
|
||||
|
||||
### Skip-Softmax calibration config
|
||||
|
||||
ModelOpt stores the calibration in each transformer component's `config.json` under
|
||||
`sparse_attention_config`. The following is taken from the `transformer/config.json` of
|
||||
[`nvidia/Wan2.2-T2V-A14B-Diffusers-FP8`](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8),
|
||||
with the `ignore` list abridged:
|
||||
|
||||
```json
|
||||
{
|
||||
"sparse_attention_config": {
|
||||
"config_groups": {
|
||||
"group_0": {
|
||||
"algorithm": "skip_softmax",
|
||||
"targets": ["WanAttention"],
|
||||
"ignore": ["blocks.0.attn1", "blocks.0.attn2", "blocks.1.attn1", "..."],
|
||||
"threshold_scale_factor": {
|
||||
"formula": "a * exp(b * target_sparsity)",
|
||||
"coefficients": {"a": 2822.6209952221557, "b": 4.388250623520793}
|
||||
},
|
||||
"target_sparsity": 0.7,
|
||||
"disabled_until_timestep": 0.93
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
A checkpoint with several transformer components carries one such config per component, each
|
||||
with its own coefficients; in this example, `transformer_2/config.json` holds the low-noise
|
||||
expert's curve.
|
||||
|
||||
vLLM-Omni consumes `formula`, `coefficients`, and `ignore`, and ignores `targets`,
|
||||
`target_sparsity`, and `disabled_until_timestep`:
|
||||
|
||||
- Only the formula `a * exp(b * target_sparsity)` is accepted; a checkpoint with any other formula
|
||||
is rejected at startup. If a spec requests `target_sparsity` and the checkpoint carries no
|
||||
calibration, startup fails and the error points to `threshold` as the calibration-free
|
||||
alternative.
|
||||
- Each component's curve is applied to that component's attention layers. If the config for a
|
||||
secondary component such as `transformer_2` cannot be read, that component stays dense and a
|
||||
warning is logged.
|
||||
- `ignore` holds fnmatch patterns matched against both the full module name and the name relative
|
||||
to the component, so `blocks.0.attn1` matches `transformer.blocks.0.attn1` as well as
|
||||
`transformer_2.blocks.0.attn1`. Matching layers receive no coefficients and therefore stay dense
|
||||
regardless of the user configuration.
|
||||
|
||||
At each attention call, a layer with a `threshold` uses it directly; a layer with
|
||||
`target_sparsity` and stamped coefficients evaluates the curve; a layer with `target_sparsity` but
|
||||
no coefficients runs dense. The timestep gate below is applied to the result before it reaches the
|
||||
kernel.
|
||||
|
||||
## Timestep gating
|
||||
|
||||
`disabled_until_timestep = D` keeps the mode off during the early, high-noise denoise steps and
|
||||
turns it on once the normalized timestep `t` drops to `t ≤ D` (`t` runs `1.0` → `0.0` over the
|
||||
schedule). The early steps set the global structure of the output and their errors propagate through
|
||||
every later step, so keeping them dense costs a few skipped-tile opportunities but protects fidelity.
|
||||
The early, high-noise denoising steps set the global structure of the output, and their errors
|
||||
propagate through every later step. Keeping just these steps dense markedly improves the quality
|
||||
of the generated video at a small cost in skipped work. `disabled_until_timestep` implements this:
|
||||
on each attention call the backend compares the current normalized timestep `t` against the
|
||||
configured cutoff and passes no skip factor to the kernel while `t` is above it. Skip-Softmax
|
||||
becomes active on the first step whose `t` is at or below the cutoff and stays active for the rest
|
||||
of the run, since `t` only decreases.
|
||||
|
||||
`t` is the scheduler's own timestep divided by `num_train_timesteps`, published by the pipeline via
|
||||
`DenoiseProgressMixin.record_denoise_step`. It is deliberately not derived from the step index,
|
||||
because schedulers space their steps non-uniformly. A pipeline that does not publish a timestep
|
||||
stays dense when `disabled_until_timestep` is set.
|
||||
The default `disabled_until_timestep = 0` is a sentinel rather than a cutoff: the gate is off, the
|
||||
timestep is never read, and Skip-Softmax runs on every step. Any positive value enables the gate.
|
||||
`disabled_until_timestep = 1.0` therefore also yields no dense steps on a pipeline that publishes
|
||||
`t`, but on a pipeline that does not publish it the backend stays dense and logs a warning once
|
||||
rather than guessing from the step index.
|
||||
|
||||
### Where `t` comes from
|
||||
|
||||
The pipeline publishes `t` once per denoising step before running the transformer. It is the
|
||||
scheduler's own position in `[0, 1]`, decreasing from near `1.0` (pure noise) to `0.0` (clean
|
||||
sample): scheduler-based pipelines publish the scheduler timestep divided by `num_train_timesteps`,
|
||||
and rectified-flow pipelines publish the current sigma. `t` is deliberately not the step index
|
||||
divided by the step count. Schedulers place their steps non-uniformly in `t`, so expressing the
|
||||
cutoff in `t` ties it to the same noise level across step counts and schedulers.
|
||||
|
||||
### Mapping the cutoff to denoising steps
|
||||
|
||||
For a run whose published sequence is `t[0], ..., t[N-1]`, the number of dense steps is
|
||||
|
||||
```text
|
||||
dense_steps = count(t[i] > disabled_until_timestep)
|
||||
skip_softmax_steps = N - dense_steps
|
||||
```
|
||||
|
||||
This count **depends on the schedule**, and schedules differ from model to model: the
|
||||
scheduler family, the number of steps, and any time shift all change where the steps land in `t`.
|
||||
The same cutoff can therefore gate a very different fraction of the run on two models. Take
|
||||
MiniMax-H3 as an example. Its video branch uses a flow-shifted rectified-flow schedule that applies
|
||||
a shift `s` to `N + 1` uniform positions `u` from `1` to `0`,
|
||||
|
||||
```text
|
||||
t = s * u / (1 + (s - 1) * u)
|
||||
```
|
||||
|
||||
and a large `s` pushes most positions toward `t ≈ 1`. With the default `s = 12` and
|
||||
`num_inference_steps = 50` (50 sigma points, `N = 49` denoiser forwards), the published sequence
|
||||
starts `1.000, 0.998, 0.996, 0.995, 0.993, ...` and stays above `0.9` for 28 steps, giving:
|
||||
|
||||
| `disabled_until_timestep` | Dense steps | Skip-Softmax steps |
|
||||
| :---: | ---: | ---: |
|
||||
| `1.00` | 0 | 49 |
|
||||
| `0.99` | 6 | 43 |
|
||||
| `0.97` | 14 | 35 |
|
||||
| `0.95` | 19 | 30 |
|
||||
| `0.90` | 28 | 21 |
|
||||
| `0.86` | 33 | 16 |
|
||||
|
||||
These numbers hold only for that schedule. A model with a smaller shift, a different step count,
|
||||
or a scheduler that is not rectified flow produces a different sequence and a different table; a
|
||||
request that overrides `flow_shift` or the step count changes it for MiniMax-H3 too.
|
||||
Choose the cutoff by counting against the schedule actually served, not by reusing a value from
|
||||
another model.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 562 KiB |
@@ -3,4 +3,6 @@
|
||||
vLLM-Omni's examples are split into two categories:
|
||||
|
||||
- If you are using vLLM-Omni from within Python code, see the *Offline Inference* section.
|
||||
- If you are using vLLM-Omni from an HTTP application or client, see the *Online Serving* section.
|
||||
- If you are using vLLM-Omni from an HTTP application or client, start with
|
||||
the [API Server guide](../serving/README.md) to choose an endpoint, then use
|
||||
the *Online Serving* section for model-specific examples.
|
||||
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
### 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, 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).
|
||||
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).
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
# --8<-- [end:requirements]
|
||||
# --8<-- [start:set-up-using-python]
|
||||
|
||||
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.26.x release line uses vLLM 0.26.x.
|
||||
vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.28.x release line uses vLLM 0.28.x.
|
||||
|
||||
!!! 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.
|
||||
@@ -95,6 +95,13 @@ 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.
|
||||
|
||||
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:
|
||||
```bash
|
||||
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
|
||||
# vLLM-Ascend image
|
||||
# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.26.0
|
||||
# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3
|
||||
# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.28.0
|
||||
# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.28.0-a3
|
||||
docker run --rm \
|
||||
--name vllm-omni-npu \
|
||||
--shm-size=1g \
|
||||
@@ -28,11 +28,11 @@ docker run --rm \
|
||||
-v /etc/ascend_install.info:/etc/ascend_install.info \
|
||||
-v /root/.cache:/root/.cache \
|
||||
-p 8000:8000 \
|
||||
-it quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 bash
|
||||
-it quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 bash
|
||||
|
||||
# Inside the container, install vLLM-Omni from source
|
||||
cd /vllm-workspace
|
||||
git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git
|
||||
git clone -b v0.28.0 https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
pip install -v -e . --no-build-isolation
|
||||
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
|
||||
@@ -50,21 +50,11 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
|
||||
|
||||
# --8<-- [start:installation-main]
|
||||
|
||||
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.)
|
||||
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.
|
||||
|
||||
```bash
|
||||
# Pin vLLM and vLLM-Ascend to the v0.26 release line
|
||||
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
|
||||
# Inside the vLLM-Ascend v0.28.0 container, install vLLM-Omni from main
|
||||
cd /vllm-workspace
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
pip install -v -e . --no-build-isolation
|
||||
@@ -88,8 +78,8 @@ Supported images as following.
|
||||
Here's an example deployment command that has been verified on 4 x NPUs:
|
||||
|
||||
```bash
|
||||
# Atlas A2: quay.io/ascend/vllm-omni:v0.26.0
|
||||
# Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3
|
||||
# Atlas A2: quay.io/ascend/vllm-omni:v0.28.0
|
||||
# Atlas A3: quay.io/ascend/vllm-omni:v0.28.0-a3
|
||||
docker run --rm \
|
||||
--name vllm-omni-a3 \
|
||||
--shm-size=4g \
|
||||
@@ -114,7 +104,7 @@ docker run --rm \
|
||||
-v /etc/ascend_install.info:/etc/ascend_install.info \
|
||||
-v ~/.cache:/root/.cache \
|
||||
-p 8091:8091 \
|
||||
-it quay.io/ascend/vllm-omni:v0.26.0-a3 bash
|
||||
-it quay.io/ascend/vllm-omni:v0.28.0-a3 bash
|
||||
```
|
||||
|
||||
!!! tip
|
||||
|
||||
@@ -91,7 +91,7 @@ if __name__ == "__main__":
|
||||
|
||||
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
|
||||
|
||||
## Online Serving with OpenAI-Completions API
|
||||
## Online Serving with the API Server
|
||||
|
||||
Text-to-image generation quickstart with vLLM-Omni:
|
||||
|
||||
@@ -100,20 +100,16 @@ vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091
|
||||
```
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
curl -s http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "a cup of coffee on the table"}
|
||||
],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > coffee.png
|
||||
"prompt": "a cup of coffee on the table",
|
||||
"size": "1024x1024",
|
||||
"response_format": "b64_json",
|
||||
"seed": 42
|
||||
}' | jq -r '.data[0].b64_json' | base64 -d > coffee.png
|
||||
```
|
||||
|
||||
For more details, please refer to [online serving](../user_guide/examples/online_serving/text_to_image.md).
|
||||
See the [API Server guide](../serving/README.md) to choose an endpoint. For
|
||||
model-specific details, refer to the [text-to-image online serving
|
||||
example](../user_guide/examples/online_serving/text_to_image.md).
|
||||
@@ -21,7 +21,7 @@ th {
|
||||
|--------------|--------|-------------------|------------|---------|-----|-----------|--------|
|
||||
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Qwen/Qwen3-Omni.md) |
|
||||
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
|
||||
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | — |
|
||||
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | — |
|
||||
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | — |
|
||||
@@ -62,7 +62,7 @@ th {
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
|
||||
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit` | ✅︎ | | | ✅︎ | — |
|
||||
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
|
||||
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
@@ -95,7 +95,7 @@ th {
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — |
|
||||
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — |
|
||||
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | — |
|
||||
| `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) |
|
||||
| `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` | ✅︎ | ✅︎ | | | — |
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
# API Server
|
||||
|
||||
vLLM-Omni exposes OpenAI-compatible HTTP APIs plus vLLM-Omni extensions for
|
||||
image, audio, video, realtime, and robot-policy workloads. Use this page to
|
||||
choose an endpoint; follow the linked reference page for model-specific fields,
|
||||
wire protocols, and response details.
|
||||
|
||||
Each server instance hosts one model. An endpoint is usable only when the
|
||||
loaded model supports that task.
|
||||
|
||||
This guide focuses on vLLM-Omni's primary task endpoints. Compatible models
|
||||
may also expose standard routes inherited from vLLM, such as
|
||||
`POST /v1/completions`; those routes remain model- and task-dependent.
|
||||
|
||||
## Start and Verify the Server
|
||||
|
||||
Start a model with the shared serving command:
|
||||
|
||||
```bash
|
||||
vllm serve <model> --omni --port 8091
|
||||
```
|
||||
|
||||
Some models require additional flags or a deployment configuration. Use the
|
||||
model-specific guide when one is provided.
|
||||
|
||||
After the server starts, check its health and served model name:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_BASE_URL=http://localhost:8091
|
||||
|
||||
curl "$VLLM_OMNI_BASE_URL/health"
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/models" | jq .
|
||||
```
|
||||
|
||||
Use `http://localhost:8091/v1` as the `base_url` for an OpenAI SDK client. If
|
||||
the server was started with `--api-key`, include an `Authorization` header with
|
||||
`Bearer <api-key>` in requests.
|
||||
|
||||
## Choose a Core HTTP Endpoint
|
||||
|
||||
Prefer a task-specific endpoint when one matches your workload. Use Chat
|
||||
Completions for conversational or heterogeneous omni pipelines, rather than as
|
||||
the default wrapper for every generation task.
|
||||
|
||||
| Task | Endpoint | Request | Response | Details |
|
||||
|------|----------|---------|----------|---------|
|
||||
| Conversation or multimodal understanding/generation | `POST /v1/chat/completions` | JSON | OpenAI-style JSON or SSE | [Chat Completions](chat_completions_api.md) |
|
||||
| Text-to-speech | `POST /v1/audio/speech` | JSON | Audio bytes or SSE | [Speech](speech_api.md) |
|
||||
| Sound, music, or ambient audio generation | `POST /v1/audio/generate` | JSON | Audio bytes | [Audio Generation](audio_generate_api.md) |
|
||||
| Text-to-image | `POST /v1/images/generations` | JSON | Base64 JSON or an image file | [Image Generation](image_generation_api.md) |
|
||||
| Image editing | `POST /v1/images/edits` | Multipart form | Base64 JSON or SSE | [Image Edit](image_edit_api.md) |
|
||||
| Video generation (recommended) | `POST /v1/videos` | Multipart form | Asynchronous job JSON | [Videos](videos_api.md) |
|
||||
| Video generation (blocking) | `POST /v1/videos/sync` | Multipart form | Video bytes | [Videos](videos_api.md#synchronous-response) |
|
||||
|
||||
`POST /v1/videos/sync` is intended for tests and simple scripts. Use the
|
||||
asynchronous job API for production requests so clients can poll, download,
|
||||
and delete generated videos independently.
|
||||
|
||||
## Minimal Requests
|
||||
|
||||
The examples below assume a compatible model is already running. TTS voices,
|
||||
media inputs, generation controls, and output capabilities vary by model.
|
||||
|
||||
=== "Chat"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/chat/completions" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Describe vLLM-Omni briefly."}],
|
||||
"modalities": ["text"]
|
||||
}'
|
||||
```
|
||||
|
||||
=== "Speech"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/audio/speech" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello from vLLM-Omni.",
|
||||
"voice": "vivian"
|
||||
}' \
|
||||
--output speech.wav
|
||||
```
|
||||
|
||||
=== "Audio"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/audio/generate" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Ocean waves on a quiet beach",
|
||||
"audio_length": 5.0
|
||||
}' \
|
||||
--output audio.wav
|
||||
```
|
||||
|
||||
=== "Image"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/images/generations" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "A small robot reading beside a window",
|
||||
"size": "1024x1024",
|
||||
"response_format": "file"
|
||||
}' \
|
||||
--output image.png
|
||||
```
|
||||
|
||||
=== "Image edit"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/images/edits" \
|
||||
-F "image=@input.png" \
|
||||
-F "prompt=Turn the background into a snowy mountain" \
|
||||
| jq -r '.data[0].b64_json' \
|
||||
| base64 --decode > edited.png
|
||||
```
|
||||
|
||||
=== "Video"
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/videos" \
|
||||
-F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \
|
||||
-F "size=1280x720" \
|
||||
-F "seconds=5"
|
||||
```
|
||||
|
||||
The response contains a video job `id`. Poll and download it with:
|
||||
|
||||
```bash
|
||||
curl "$VLLM_OMNI_BASE_URL/v1/videos/<video_id>"
|
||||
curl -L "$VLLM_OMNI_BASE_URL/v1/videos/<video_id>/content" --output video.mp4
|
||||
```
|
||||
|
||||
## Choose a Streaming or Realtime Endpoint
|
||||
|
||||
These WebSocket APIs have different event schemas and cannot be used
|
||||
interchangeably.
|
||||
|
||||
| Workload | Endpoint | Interaction model | Details |
|
||||
|----------|----------|-------------------|---------|
|
||||
| Incremental text input for speech synthesis | `WS /v1/audio/speech/stream` | Send text events and receive audio | [Streaming Text to Speech](speech_api.md#streaming-text-input-websocket) |
|
||||
| Live video understanding | `WS /v1/video/chat/stream` | Send video frames and receive text/audio | [Streaming Video Input](video_stream_api.md) |
|
||||
| Turn-based realtime audio | `WS /v1/realtime` | Stream one audio input and receive transcript/audio events | [Realtime Audio](realtime_api.md) |
|
||||
| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` or `WS /v1/duplex` | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) |
|
||||
| Generated video chunks | `WS /v1/realtime/video` | Start a diffusion request and receive fragmented MP4 | [Streaming Video Output](streaming_video_output_api.md) |
|
||||
| Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) |
|
||||
|
||||
All six routes are model- or configuration-dependent. In particular,
|
||||
`/v1/realtime` is not full duplex unless the client sets `duplex=1` and the
|
||||
deployment explicitly enables duplex sessions. Clients should also verify the
|
||||
duplex capability payload because the query-parameter form falls back to the
|
||||
ordinary realtime handler when duplex is unavailable.
|
||||
|
||||
## Related Endpoints
|
||||
|
||||
| Purpose | Endpoints | Reference |
|
||||
|---------|-----------|-----------|
|
||||
| Discovery and readiness | `GET /health`, `GET /v1/models` | This page |
|
||||
| Batched conversations | `POST /v1/chat/completions/batch` | [Batch requests](chat_completions_api.md#batch-requests) |
|
||||
| Batched speech | `POST /v1/audio/speech/batch` | [Batch speech generation](speech_api.md#batch-speech-generation) |
|
||||
| TTS voice management | `GET/POST /v1/audio/voices`, `DELETE /v1/audio/voices/{name}` | [Voices](speech_api.md#voices-endpoint) |
|
||||
| Video job lifecycle | `GET /v1/videos`, `GET/DELETE /v1/videos/{video_id}`, `GET /v1/videos/{video_id}/content` | [Video endpoints](videos_api.md#endpoints) |
|
||||
| Release and restore stage memory | `POST /v1/omni/sleep`, `POST /v1/omni/wakeup` | [Sleep Mode](../features/sleep_mode.md) |
|
||||
|
||||
## Standalone Experimental Servers
|
||||
|
||||
PersonaPlex also provides a Moshi-compatible WebSocket server, while JoyVL
|
||||
provides a stateful interaction orchestrator in front of another model server.
|
||||
These are separate processes with their own routes; they are not additional
|
||||
paths on every `vllm serve --omni` instance. See [Standalone Experimental
|
||||
Servers](standalone_servers.md) before deploying either one.
|
||||
|
||||
## Compatibility Notes
|
||||
|
||||
- A registered path does not mean every model supports it. Unsupported tasks
|
||||
return an OpenAI-style error or a WebSocket `error` event.
|
||||
- Use JSON for Chat, Speech, Audio Generation, and Image Generation. Use
|
||||
`multipart/form-data` when uploading images, videos, or other files.
|
||||
- Generation can take minutes for large diffusion models. Configure client
|
||||
timeouts accordingly; prefer asynchronous video jobs over long blocking
|
||||
requests.
|
||||
- Model-specific parameters are vLLM-Omni extensions. Check the endpoint
|
||||
reference before assuming an OpenAI SDK exposes them directly.
|
||||
@@ -4,14 +4,14 @@ vLLM-Omni provides an API for text-to-audio generation using diffusion-based mod
|
||||
|
||||
Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
|
||||
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Start the Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve stabilityai/stable-audio-open-1.0 \
|
||||
vllm serve stabilityai/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
@@ -326,7 +326,7 @@ curl http://localhost:8091/health
|
||||
Enable debug logging:
|
||||
|
||||
```bash
|
||||
vllm-omni serve stabilityai/stable-audio-open-1.0 \
|
||||
vllm serve stabilityai/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
|
||||
@@ -1,78 +1,110 @@
|
||||
# Chat Completions API
|
||||
|
||||
vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
|
||||
endpoint using diffusion models. This page explains how to pass generation
|
||||
parameters (such as `num_inference_steps`, `height`, `width`) to diffusion
|
||||
models through this endpoint.
|
||||
Use `/v1/chat/completions` for conversational and multimodal pipelines that
|
||||
follow the OpenAI Chat API schema. Depending on the loaded model, a request can
|
||||
contain text, image, audio, or video input and can return text, audio, images,
|
||||
or other model-specific output.
|
||||
|
||||
!!! tip
|
||||
For dedicated endpoints that accept generation parameters as top-level
|
||||
fields, see [Image Generation API](image_generation_api.md) and
|
||||
[Image Edit API](image_edit_api.md).
|
||||
Some diffusion pipelines also support image generation and editing through
|
||||
Chat Completions. Prefer the dedicated [Image Generation
|
||||
API](image_generation_api.md) or [Image Edit API](image_edit_api.md) for
|
||||
task-specific clients; their request fields and response contracts are more
|
||||
direct.
|
||||
|
||||
## Passing Generation Parameters
|
||||
|
||||
The `/v1/chat/completions` endpoint follows the OpenAI Chat API schema, which
|
||||
does not natively include diffusion-specific fields like `num_inference_steps`
|
||||
or `height`. How you pass these extra fields depends on your client.
|
||||
|
||||
### curl / Python `requests`
|
||||
|
||||
Wrap generation parameters inside an `"extra_body"` key in the JSON body:
|
||||
## Basic Request
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "A beautiful landscape painting"}
|
||||
{"role": "user", "content": "Describe vLLM-Omni briefly."}
|
||||
],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42
|
||||
}
|
||||
"modalities": ["text"]
|
||||
}'
|
||||
```
|
||||
|
||||
### OpenAI Python SDK
|
||||
Use `stream: true` for Server-Sent Events. Input media syntax, supported output
|
||||
`modalities`, and response choices depend on the model; see the
|
||||
[model-specific online serving examples](../user_guide/examples/online_serving/qwen3_omni.md).
|
||||
|
||||
Use the `extra_body` **keyword argument**. The SDK automatically merges these
|
||||
fields into the top-level request body:
|
||||
!!! tip
|
||||
Each server hosts one model. Query `GET /v1/models` for its served model
|
||||
name when your client requires the `model` field.
|
||||
|
||||
```python
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen-Image",
|
||||
messages=[{"role": "user", "content": "A beautiful landscape painting"}],
|
||||
extra_body={
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
## vLLM Extension Parameters
|
||||
|
||||
Send standard and vLLM-specific Chat parameters as top-level fields in a
|
||||
direct HTTP request. When using the OpenAI Python SDK, pass vLLM-specific
|
||||
fields through the SDK's `extra_body` keyword; the SDK merges them into the
|
||||
top-level JSON sent to the server.
|
||||
|
||||
=== "curl"
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Write a short haiku."}],
|
||||
"top_k": 40
|
||||
}'
|
||||
```
|
||||
|
||||
=== "OpenAI Python SDK"
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="your-served-model",
|
||||
messages=[{"role": "user", "content": "Write a short haiku."}],
|
||||
extra_body={"top_k": 40},
|
||||
)
|
||||
```
|
||||
|
||||
??? note "Diffusion compatibility"
|
||||
|
||||
Diffusion models exposed through Chat Completions can accept supported
|
||||
fields such as `num_inference_steps`, `seed`, `height`, and `width`. Send
|
||||
them at the top level with direct HTTP, or in the SDK's `extra_body`
|
||||
keyword argument.
|
||||
|
||||
A literal nested `"extra_body"` JSON object is accepted for compatibility,
|
||||
but it is not recommended for new direct HTTP clients. Do not provide the
|
||||
same parameter in multiple locations; the server returns a `400` error for
|
||||
duplicate diffusion parameters. Prefer the dedicated [Image Generation
|
||||
API](image_generation_api.md) or [Image Edit API](image_edit_api.md) when
|
||||
either endpoint matches the task.
|
||||
|
||||
## Batch Requests
|
||||
|
||||
`POST /v1/chat/completions/batch` accepts the same shared generation fields,
|
||||
but `messages` is a list of conversations. The response contains one choice
|
||||
per conversation in input order.
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions/batch \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
[{"role": "user", "content": "Summarize vLLM in one sentence."}],
|
||||
[{"role": "user", "content": "Summarize vLLM-Omni in one sentence."}]
|
||||
],
|
||||
"max_tokens": 64
|
||||
}'
|
||||
```
|
||||
|
||||
!!! note "SDK `extra_body` vs. JSON `extra_body`"
|
||||
These two `extra_body` usages look similar but work differently under the
|
||||
hood. The SDK flattens the dict into the top-level request JSON, while the
|
||||
curl/requests approach sends it as a nested `"extra_body"` key. Both are
|
||||
handled correctly by the server.
|
||||
|
||||
!!! note "About the `ignored fields` warning"
|
||||
You may see a log message like:
|
||||
|
||||
```
|
||||
WARNING: The following fields were present in the request but ignored: {'height', 'width', ...}
|
||||
```
|
||||
|
||||
This is **harmless**. It is emitted by vLLM's request validation layer
|
||||
because these fields are not part of the standard OpenAI
|
||||
`ChatCompletionRequest` schema. The fields are still stored internally
|
||||
and correctly forwarded to the diffusion pipeline.
|
||||
Batch Chat Completions does not support streaming, tools, beam search, or
|
||||
`n > 1`.
|
||||
|
||||
## Model-Specific Examples
|
||||
|
||||
For complete examples with full request/response details, see the model-specific
|
||||
guides:
|
||||
For complete examples with model-specific inputs and outputs, see:
|
||||
|
||||
- [Qwen3-Omni](../user_guide/examples/online_serving/qwen3_omni.md)
|
||||
- [Qwen2.5-Omni](../user_guide/examples/online_serving/qwen2_5_omni.md)
|
||||
- [Text-to-Image (Qwen-Image)](../user_guide/examples/online_serving/text_to_image.md)
|
||||
- [Image-to-Image (Qwen-Image-Edit, Qwen-Image-Layered)](../user_guide/examples/online_serving/image_to_image.md)
|
||||
- [GLM-Image](../user_guide/examples/online_serving/glm_image.md)
|
||||
@@ -0,0 +1,116 @@
|
||||
# Full-Duplex WebSocket API
|
||||
|
||||
vLLM-Omni provides an experimental full-duplex runtime for models that can
|
||||
continue receiving speech while producing speech. It adds persistent session
|
||||
state, model-specific turn policy, overlap handling, playback acknowledgement,
|
||||
and optional session resume.
|
||||
|
||||
Full duplex is distinct from the turn-based [Realtime Audio API](realtime_api.md).
|
||||
|
||||
## Choose an Endpoint
|
||||
|
||||
| Endpoint | Protocol | Recommended use |
|
||||
|----------|----------|-----------------|
|
||||
| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style event projection | Applications and browser clients |
|
||||
| `WS /v1/duplex` | Native vLLM-Omni duplex events | Runtime integration and low-level testing |
|
||||
|
||||
Both endpoints use the same duplex engine and require the same model-side
|
||||
adapter. Prefer the Realtime projection unless the native lifecycle events are
|
||||
specifically required.
|
||||
|
||||
## Enable Full Duplex
|
||||
|
||||
The route becomes usable only when the deployment configuration explicitly
|
||||
sets:
|
||||
|
||||
```yaml
|
||||
session_mode: duplex
|
||||
```
|
||||
|
||||
The selected model pipeline must also provide a duplex serving adapter.
|
||||
Model-native deployments configure the corresponding engine runtime extension
|
||||
and control plane as part of their registered pipeline.
|
||||
|
||||
!!! warning
|
||||
|
||||
`WS /v1/duplex` fails with `Duplex API is not available` when duplex is not
|
||||
enabled. By contrast, `/v1/realtime?duplex=1` falls back to the ordinary
|
||||
turn-based realtime handler when no duplex handler exists. Confirm that
|
||||
`session.created.session.capabilities` is present before treating the
|
||||
connection as full duplex.
|
||||
|
||||
The current unified-runtime integrations are:
|
||||
|
||||
- MiniCPM-o 4.5, using `vllm_omni/deploy/minicpmo_4_5_duplex.yaml`;
|
||||
- PersonaPlex, whose default `vllm_omni/deploy/personaplex.yaml` enables duplex.
|
||||
|
||||
JoyVL is a separate HTTP interaction orchestrator and does not use these
|
||||
WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md).
|
||||
|
||||
## MiniCPM-o Quick Start
|
||||
|
||||
Start the duplex deployment:
|
||||
|
||||
```bash
|
||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
|
||||
--trust-remote-code \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Stream a mono, PCM16, 16 kHz WAV file with the provided client:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/minicpmo/realtime_duplex_demo.py \
|
||||
--url 'ws://localhost:8091/v1/realtime?duplex=1' \
|
||||
--model openbmb/MiniCPM-o-4_5 \
|
||||
--input-wav input_16k_mono.wav \
|
||||
--ref-audio reference_voice.wav \
|
||||
--output-dir /tmp/minicpmo-duplex
|
||||
```
|
||||
|
||||
## Realtime Event Lifecycle
|
||||
|
||||
A typical `/v1/realtime?duplex=1` session follows this lifecycle:
|
||||
|
||||
1. Send `session.update` with the model, modalities, audio formats, and session
|
||||
options.
|
||||
2. Wait for `session.created`; inspect `session.capabilities` instead of
|
||||
assuming every model supports the same controls.
|
||||
3. Send `input_audio_buffer.append` events while microphone audio arrives.
|
||||
4. Send `input_audio_buffer.commit` at a user-turn boundary when required by
|
||||
the model policy.
|
||||
5. Consume `response.created`, transcript deltas, `response.audio.delta`, and
|
||||
`response.done` or `response.listen` events.
|
||||
6. Send `playback.ack` after audio has been played when the session advertises
|
||||
playback acknowledgement support.
|
||||
7. Send `session.close` and wait for `session.closed`.
|
||||
|
||||
Unlike the turn-based realtime endpoint, input may continue while a response
|
||||
is active. The server can emit `overlap.decision` to describe whether input was
|
||||
deferred, treated as a short acknowledgement, or used to interrupt output.
|
||||
|
||||
## Capabilities and Model Differences
|
||||
|
||||
The `session.created` payload includes capability fields such as
|
||||
`supports_barge_in`, `supports_playback_ack`, `supports_multi_session`,
|
||||
`supports_session_resume`, and `chunk_period_ms`. Treat this payload as the
|
||||
runtime contract.
|
||||
|
||||
For example, PersonaPlex supports native overlapping speech but currently
|
||||
advertises `supports_barge_in=false`; destructive output interruption and
|
||||
model-state rewind have not been validated for that integration. Capacity and
|
||||
session-resume behavior also depend on the selected deployment configuration.
|
||||
|
||||
## Native Protocol
|
||||
|
||||
`WS /v1/duplex` exposes lower-level lifecycle names including
|
||||
`session.create`, `input_audio_buffer.append`, `turn.signal`, `playback.ack`,
|
||||
and `session.close`. It returns native session, input, response, overlap, and
|
||||
error events. This protocol is experimental and may evolve with the runtime;
|
||||
applications should use the provided Realtime client where possible.
|
||||
|
||||
See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo),
|
||||
[PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex),
|
||||
and [full-duplex runtime design](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/experimental/fullduplex/DESIGN.md)
|
||||
for model-specific validation and architecture details.
|
||||
@@ -0,0 +1,109 @@
|
||||
# OpenPI Robot Policy WebSocket API
|
||||
|
||||
Use `WS /v1/realtime/robot/openpi` for low-latency robot-policy inference with
|
||||
OpenPI-compatible clients. A persistent WebSocket carries observations to the
|
||||
loaded policy model and returns action arrays.
|
||||
|
||||
Despite the `/realtime` prefix, this endpoint does not use the OpenAI Realtime
|
||||
event schema. It uses binary MessagePack frames with NumPy extensions.
|
||||
|
||||
## Availability
|
||||
|
||||
The endpoint is available only when the loaded diffusion policy configuration
|
||||
contains `policy_server_config`. The server sends that model-specific
|
||||
configuration to the client as the first binary frame.
|
||||
|
||||
For example, DreamZero declares:
|
||||
|
||||
```yaml
|
||||
policy_server_config:
|
||||
image_resolution: [180, 320]
|
||||
n_external_cameras: 2
|
||||
needs_wrist_camera: true
|
||||
needs_stereo_camera: false
|
||||
needs_session_id: true
|
||||
action_space: joint_position
|
||||
```
|
||||
|
||||
Without this configuration, the WebSocket returns `Robot policy not available`
|
||||
and closes.
|
||||
|
||||
## DreamZero Quick Start
|
||||
|
||||
Start the bundled DreamZero deployment:
|
||||
|
||||
```bash
|
||||
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8091 \
|
||||
--served-model-name dreamzero-droid \
|
||||
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
|
||||
--enforce-eager --disable-log-stats
|
||||
```
|
||||
|
||||
Install the optional client dependencies, download the sample camera inputs,
|
||||
and run the client:
|
||||
|
||||
```bash
|
||||
pip install openpi-client websockets opencv-python
|
||||
|
||||
hf download YangshenDeng/vllm-omni-dreamzero-assets \
|
||||
--repo-type dataset \
|
||||
--local-dir outputs/dreamzero/assets
|
||||
|
||||
python examples/online_serving/dreamzero/openpi_client.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--video-dir outputs/dreamzero/assets
|
||||
```
|
||||
|
||||
## Protocol
|
||||
|
||||
The connection is request-response after an initial server handshake:
|
||||
|
||||
```text
|
||||
connect
|
||||
<- msgpack(policy_server_config)
|
||||
|
||||
infer
|
||||
-> msgpack({"endpoint": "infer", "session_id": "...", ...observation})
|
||||
<- msgpack(ndarray | dict[str, ndarray])
|
||||
|
||||
reset
|
||||
-> msgpack({"endpoint": "reset"})
|
||||
<- msgpack({"status": "reset successful"})
|
||||
```
|
||||
|
||||
If `endpoint` is omitted, the server treats the message as `infer`. Each
|
||||
inference message produces one action response after the engine request
|
||||
completes; action tokens or intermediate tensors are not streamed.
|
||||
|
||||
NumPy arrays use the marker format implemented by `openpi-client`. The server
|
||||
also accepts the legacy vLLM NumPy marker representation on input. JSON text
|
||||
frames are not observation messages.
|
||||
|
||||
## Sessions and Reset
|
||||
|
||||
- The API layer's current-session and first-call counters are scoped to each
|
||||
WebSocket connection.
|
||||
- `session_id` identifies model-side state across observations on that
|
||||
connection. If omitted, it defaults to `default`.
|
||||
- The first inference for a session is sent to the model with `reset=true`.
|
||||
- Changing `session_id`, or sending the `reset` command, causes the next
|
||||
inference to start with `reset=true`.
|
||||
- The policy pipeline owns observation transforms and persistent model state,
|
||||
normally keyed by `session_id`; the API layer forwards the raw observation
|
||||
dictionary.
|
||||
|
||||
## Limits and Errors
|
||||
|
||||
- Maximum inbound payload size is 64 MiB.
|
||||
- The server closes an idle connection after 30 seconds.
|
||||
- Invalid binary input returns a MessagePack
|
||||
`{"type":"error","message":"Invalid request payload"}` response.
|
||||
- Inference failures return a generic `Internal inference error` without
|
||||
exposing an internal traceback.
|
||||
- Unsupported NumPy object, structured, and complex dtypes are rejected.
|
||||
|
||||
Observation keys, camera layout, state tensors, and action shapes are defined
|
||||
by the loaded policy rather than by this transport. Read the handshake before
|
||||
constructing observations. See the [DreamZero example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/dreamzero)
|
||||
for a complete OpenPI client and DROID simulation loop.
|
||||
@@ -0,0 +1,82 @@
|
||||
# Realtime Audio WebSocket API
|
||||
|
||||
Use `WS /v1/realtime` to stream PCM audio into a compatible omni model and
|
||||
receive incremental transcription and synthesized response audio. This is a
|
||||
turn-based streaming API: it starts one generation over an audio stream and
|
||||
finishes after the client closes that input stream.
|
||||
|
||||
For a model that listens while it is already speaking, use the experimental
|
||||
[Full-Duplex API](full_duplex_api.md) instead.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Start a compatible model. The currently documented example uses Qwen3-Omni:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
Run the provided client with a mono, 16-bit PCM, 16 kHz WAV file:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/qwen3_omni/openai_realtime_client.py \
|
||||
--url ws://localhost:8091/v1/realtime \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--input-wav input_16k_mono.wav \
|
||||
--output-wav response.wav \
|
||||
--output-text response.txt
|
||||
```
|
||||
|
||||
The client dependency is `websockets`.
|
||||
|
||||
## Protocol
|
||||
|
||||
Messages in both directions are JSON text frames. Audio carried inside an
|
||||
event is base64-encoded raw PCM16.
|
||||
|
||||
| Direction | Event | Purpose |
|
||||
|-----------|-------|---------|
|
||||
| Server to client | `session.created` | Confirms that the WebSocket connection is ready |
|
||||
| Client to server | `session.update` | Selects and validates the served model |
|
||||
| Client to server | `input_audio_buffer.commit` with `final: false` | Starts generation over the incoming stream |
|
||||
| Client to server | `input_audio_buffer.append` | Appends base64 PCM16 audio |
|
||||
| Client to server | `input_audio_buffer.commit` with `final: true` | Marks the end of input |
|
||||
| Server to client | `transcription.delta` | Carries incremental response text |
|
||||
| Server to client | `transcription.done` | Carries final text and token usage |
|
||||
| Server to client | `response.audio.delta` | Carries incremental PCM16 response audio |
|
||||
| Server to client | `response.audio.done` | Marks the end of response audio |
|
||||
| Server to client | `error` | Reports an invalid event, model, or audio payload |
|
||||
|
||||
A minimal client sends events in this order:
|
||||
|
||||
```json
|
||||
{"type":"session.update","model":"Qwen/Qwen3-Omni-30B-A3B-Instruct"}
|
||||
{"type":"input_audio_buffer.commit","final":false}
|
||||
{"type":"input_audio_buffer.append","audio":"<base64-pcm16>"}
|
||||
{"type":"input_audio_buffer.commit","final":true}
|
||||
```
|
||||
|
||||
The initial non-final commit intentionally starts generation before all audio
|
||||
has arrived. Continue sending `input_audio_buffer.append` events while the
|
||||
engine is consuming the stream.
|
||||
|
||||
## Audio Handling
|
||||
|
||||
- Input is mono PCM16 at 16 kHz for the Qwen3-Omni example.
|
||||
- `response.audio.delta.audio` contains base64-encoded PCM16 bytes.
|
||||
- Read `sample_rate_hz` from each audio event instead of assuming an output
|
||||
rate. Qwen3-Omni output is typically 24 kHz.
|
||||
- Concatenate audio deltas in receive order to construct the output waveform.
|
||||
- Set `input_audio_buffer.commit.final` to `true` after the last input chunk so
|
||||
the server can terminate the streaming request.
|
||||
|
||||
## Availability and Limitations
|
||||
|
||||
The path is registered on the unified API server, but it is usable only when
|
||||
the loaded pipeline implements realtime audio input and produces compatible
|
||||
audio output. Unsupported deployments return an `error` event. The endpoint
|
||||
does not provide duplex session resume, playback acknowledgement, overlap
|
||||
policy, or barge-in controls.
|
||||
|
||||
See the [Qwen3-Omni online serving example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni)
|
||||
for concurrency options, chunk pacing, and per-delta audio debugging.
|
||||
+62
-32
@@ -2,7 +2,7 @@
|
||||
|
||||
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. Output: 24 kHz.
|
||||
- **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.
|
||||
- **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.
|
||||
- **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
|
||||
|
||||
```
|
||||
```text
|
||||
POST /v1/audio/speech
|
||||
Content-Type: application/json
|
||||
```
|
||||
@@ -105,7 +105,7 @@ Content-Type: application/json
|
||||
#### OpenAI Standard Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `input` | string | **required** | The text to synthesize into speech |
|
||||
| `model` | string | server's model | Model to use (optional, should match server if specified) |
|
||||
| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
|
||||
@@ -115,8 +115,9 @@ Content-Type: application/json
|
||||
#### vLLM-Omni Extension Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `task_type` | string | "CustomVoice" | TTS task type: CustomVoice, VoiceDesign, or Base |
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `task_type` | string | null (inferred) | TTS task type: CustomVoice, VoiceDesign, or Base. For Qwen3-TTS, `ref_audio` or `ref_text` infers Base when this field is omitted, while an uploaded/precomputed voice always selects Base. Other omitted values select CustomVoice; VoiceDesign must be specified explicitly. |
|
||||
| `sample_rate` | integer | model native | Target output sample rate. Qwen3-TTS supports 8000 or 24000 Hz; the model remains native 24 kHz internally and is resampled before encoding. |
|
||||
| `language` | string | "Auto" | Language (see supported languages below) |
|
||||
| `instructions` | string | "" | Voice style/emotion instructions |
|
||||
| `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
|
||||
@@ -130,7 +131,7 @@ Content-Type: application/json
|
||||
#### Voice Clone Parameters (Base task)
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`). Local files fold `mtime_ns` and `size` into cache keys to automatically reload on-disk edits; HTTP URLs and base64 URIs remain cached by string locator. |
|
||||
| `ref_text` | string | null | Transcript of reference audio |
|
||||
| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
|
||||
@@ -204,7 +205,7 @@ The `usage` object on `speech.audio.done` is the same shape returned per item by
|
||||
|
||||
### Voices Endpoint
|
||||
|
||||
```
|
||||
```text
|
||||
GET /v1/audio/voices
|
||||
```
|
||||
|
||||
@@ -229,7 +230,7 @@ Lists available voices for the loaded model.
|
||||
|
||||
`uploaded_voices` is always present (empty list when no custom voices have been uploaded). Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
|
||||
|
||||
```
|
||||
```text
|
||||
POST /v1/audio/voices
|
||||
Content-Type: multipart/form-data
|
||||
```
|
||||
@@ -239,7 +240,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
**Form Parameters:**
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------|
|
||||
| ----------- | ------ | ---------- | ------------- |
|
||||
| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
|
||||
| `consent` | string | Yes | Consent recording ID |
|
||||
| `name` | string | Yes | Name for the new voice |
|
||||
@@ -288,7 +289,7 @@ The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and
|
||||
Client -> Server:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| --------- | ------------- |
|
||||
| `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) |
|
||||
| `{"type": "input.text", "text": "..."}` | Text chunk |
|
||||
| `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open |
|
||||
@@ -297,7 +298,7 @@ Client -> Server:
|
||||
Server -> Client:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| --------- | ------------- |
|
||||
| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
|
||||
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
|
||||
| `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
|
||||
@@ -311,17 +312,17 @@ text, emits `session.done`, and then waits on the same connection for the next
|
||||
utterance, so a client that speaks repeatedly (for example one driven by an
|
||||
upstream LLM) pays the WebSocket handshake once instead of once per utterance.
|
||||
|
||||
* The session config is sticky. Send `input.text` again straight after
|
||||
- The session config is sticky. Send `input.text` again straight after
|
||||
`session.done` to reuse it, or send another `session.config` first to change
|
||||
voice, format, or reference audio. A `session.config` sent while text is
|
||||
still buffered is rejected so no pending input is silently dropped.
|
||||
* An utterance is the flush unit, not a linguistic one: it is whatever text was
|
||||
- An utterance is the flush unit, not a linguistic one: it is whatever text was
|
||||
buffered when `input.done` arrived, of any length, synthesized as one request.
|
||||
`utterance_index` counts those flushes across the connection, so it tells you
|
||||
which `input.done` a frame belongs to. `sentence_index` counts within one
|
||||
flush and so pairs with `total_sentences`, which means every utterance reports
|
||||
`sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer).
|
||||
* End the connection with `session.close`, or by closing the socket. An idle
|
||||
- End the connection with `session.close`, or by closing the socket. An idle
|
||||
connection is still closed after the server's idle timeout, which now also
|
||||
applies to the gap between utterances.
|
||||
|
||||
@@ -330,10 +331,9 @@ upstream LLM) pays the WebSocket handshake once instead of once per utterance.
|
||||
All REST API parameters are supported, plus:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |
|
||||
|
||||
|
||||
```bash
|
||||
DELETE /v1/audio/voices/{name}
|
||||
```
|
||||
@@ -343,7 +343,7 @@ Delete an uploaded voice sample.
|
||||
**Path Parameters:**
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------|
|
||||
| ----------- | ------ | ---------- | ------------- |
|
||||
| `name` | string | Yes | Name of the voice to delete |
|
||||
|
||||
**Response Example:**
|
||||
@@ -433,6 +433,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
### Upload Voice
|
||||
|
||||
Upload voice (speaker embedding only):
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
@@ -441,6 +442,7 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
```
|
||||
|
||||
Upload voice with transcript (in-context cloning, higher quality):
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
@@ -450,6 +452,7 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
```
|
||||
|
||||
### Use Uploaded Voice
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
@@ -459,6 +462,11 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
For Qwen3-TTS, an uploaded voice is a Base voice-cloning input. The server
|
||||
infers `task_type="Base"` when `voice` names an uploaded entry, so the request
|
||||
must be sent to a Base checkpoint. Built-in presets such as `vivian` and
|
||||
`ryan` remain CustomVoice speakers and require a CustomVoice checkpoint.
|
||||
|
||||
### Voice Storage & Caching
|
||||
|
||||
Uploaded voices are persisted to disk as a single `.safetensors` file per voice
|
||||
@@ -478,6 +486,9 @@ once.
|
||||
### Precomputed Custom Voices
|
||||
|
||||
Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
|
||||
Qwen3-TTS precomputed voices follow the same Base task and checkpoint-matching
|
||||
rules as uploaded voices.
|
||||
|
||||
Generate a directory containing `custom_voice_manifest.json` plus one
|
||||
`.safetensors` file per voice, then set the pipeline-wide deploy config field:
|
||||
|
||||
@@ -516,7 +527,7 @@ by `GET /v1/audio/voices`. Valid precomputed voices can be used in
|
||||
**Configuration (environment variables):**
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| ---------- | --------- | ------------- |
|
||||
| `SPEAKER_SAMPLES_DIR` | `~/.cache/vllm-omni/speakers` | Directory for persisted uploaded speakers (`.safetensors` files). |
|
||||
| `SPEAKER_MAX_UPLOADED` | `1000` | Maximum number of uploaded speakers kept on disk. Upload requests past the cap return 400. |
|
||||
|
||||
@@ -528,7 +539,7 @@ The batch endpoint synthesizes multiple texts in a single request, returning all
|
||||
|
||||
### Endpoint
|
||||
|
||||
```
|
||||
```text
|
||||
POST /v1/audio/speech/batch
|
||||
Content-Type: application/json
|
||||
```
|
||||
@@ -536,7 +547,7 @@ Content-Type: application/json
|
||||
### Request Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `items` | array | **required** | List of items to synthesize (1–32) |
|
||||
| `model` | string | server's model | Model to use |
|
||||
| `voice` | string | null | Default voice for all items |
|
||||
@@ -553,7 +564,7 @@ Content-Type: application/json
|
||||
Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| ------- | ------ | ------------- |
|
||||
| `input` | string | **required** — text to synthesize |
|
||||
| `voice` | string | Override voice for this item |
|
||||
| `response_format` | string | Override format for this item |
|
||||
@@ -685,7 +696,7 @@ for result in response.json()["results"]:
|
||||
### Configuration
|
||||
|
||||
| Parameter | Source | Default | Description |
|
||||
|-----------|--------|---------|-------------|
|
||||
| ----------- | -------- | --------- | ------------- |
|
||||
| `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request |
|
||||
|
||||
All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_num_seqs` and queues the rest — no client-side throttling needed.
|
||||
@@ -708,17 +719,17 @@ The bundled config also sets `initial_codec_chunk_frames: 1`. This emits only th
|
||||
### Qwen3-TTS
|
||||
|
||||
| Model | Task Type | Description |
|
||||
|-------|-----------|-------------|
|
||||
| ------- | ----------- | ------------- |
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio or an uploaded/precomputed voice |
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster Base variant for voice cloning, including uploaded/precomputed voices |
|
||||
|
||||
### Fish Speech S2 Pro
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| ------- | ------------- |
|
||||
| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. |
|
||||
|
||||
Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech section of the online TTS hub](../user_guide/examples/online_serving/text_to_speech.md#fish-speech-s2-pro) for details.
|
||||
@@ -726,31 +737,31 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
|
||||
### Voxtral TTS
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| ------- | ------------- |
|
||||
| `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
|
||||
|
||||
### CosyVoice3
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| ------- | ------------- |
|
||||
| `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | Voice cloning from `ref_audio` + `ref_text`. No built-in voice presets — upload a voice or pass `ref_audio`/`ref_text` per request. |
|
||||
|
||||
### OmniVoice
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| ------- | ------------- |
|
||||
| `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. |
|
||||
|
||||
### VoxCPM2
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| ------- | ------------- |
|
||||
| `openbmb/VoxCPM2` | TTS + voice cloning with built-in speaker presets and uploaded-voice support. Accepts `voice` (preset or uploaded) or `ref_audio` + optional `ref_text`. |
|
||||
|
||||
### MOSS-TTS-Nano
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| ------- | ------------- |
|
||||
| `OpenMOSS-Team/MOSS-TTS-Nano` | Voice cloning only. Requires `ref_audio` (or an uploaded `voice`); no built-in voice presets. `ref_text` is accepted but ignored — upstream's `voice_clone` mode does not consume a transcript. |
|
||||
|
||||
## Error Responses
|
||||
@@ -770,6 +781,19 @@ Invalid parameters:
|
||||
}
|
||||
```
|
||||
|
||||
Qwen3-TTS task/checkpoint mismatch:
|
||||
|
||||
```json
|
||||
{
|
||||
"error": {
|
||||
"message": "Qwen3-TTS CustomVoice checkpoint does not support task_type='Base'. Use task_type='CustomVoice' or load the matching Base checkpoint.",
|
||||
"type": "BadRequestError",
|
||||
"param": null,
|
||||
"code": 400
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 404 Not Found
|
||||
|
||||
Model not found:
|
||||
@@ -790,10 +814,15 @@ Model not found:
|
||||
### "TTS model did not produce audio output"
|
||||
|
||||
Ensure you're using the correct model variant for your task type:
|
||||
|
||||
- CustomVoice task → CustomVoice model
|
||||
- VoiceDesign task → VoiceDesign model
|
||||
- Base task → Base model
|
||||
|
||||
Uploaded and precomputed Qwen3-TTS voices select the Base task, even when the
|
||||
client supplies a different `task_type`; serve them with a Base checkpoint.
|
||||
Built-in Qwen3-TTS presets remain CustomVoice voices.
|
||||
|
||||
### Server Not Running
|
||||
|
||||
```bash
|
||||
@@ -804,6 +833,7 @@ curl http://localhost:8091/v1/audio/voices
|
||||
### Out of Memory
|
||||
|
||||
If you encounter OOM errors:
|
||||
|
||||
1. Use smaller model variant: `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`
|
||||
2. Reduce `--gpu-memory-utilization`
|
||||
|
||||
@@ -822,7 +852,7 @@ final-output drain to a condition-variable wakeup at the same time.
|
||||
**Configuration (environment variables):**
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| --- | --- | --- |
|
||||
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `0` (off) | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. Enabled by `1`, `true`, `yes`, or `on`, matched case-insensitively after surrounding whitespace is stripped; any other value leaves it off. |
|
||||
|
||||
Set it on the process that runs the orchestrator (stage 0 of an omni
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
# Standalone Experimental Servers
|
||||
|
||||
Most vLLM-Omni endpoints run on the unified server started with
|
||||
`vllm serve <model> --omni`. The experimental full-duplex package also contains
|
||||
two separate user-facing server processes. Their routes are not added to the
|
||||
unified API server.
|
||||
|
||||
| Server | Transport | Role |
|
||||
|--------|-----------|------|
|
||||
| PersonaPlex compatibility server | WebSocket | Hosts the official PersonaPlex browser protocol and a raw-PCM alternative |
|
||||
| JoyVL interaction server | HTTP | Adds state, memory, proactive decisions, and delegation in front of a separate OpenAI-compatible model backend |
|
||||
|
||||
## PersonaPlex Compatibility Server
|
||||
|
||||
Use this server when a client requires the official Moshi/PersonaPlex wire
|
||||
protocol or the bundled browser UI. For the engine-owned unified duplex path,
|
||||
use [`WS /v1/realtime?duplex=1`](full_duplex_api.md) instead.
|
||||
|
||||
Start it with:
|
||||
|
||||
```bash
|
||||
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \
|
||||
vllm_omni.experimental.fullduplex.personaplex.serving.server \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
| Route | Purpose |
|
||||
|-------|---------|
|
||||
| `GET /` | Official PersonaPlex browser client, when its assets are available |
|
||||
| `GET /health` | Readiness check |
|
||||
| `WS /api/chat` | Moshi-compatible tagged binary protocol with Opus audio |
|
||||
| `WS /v1/audio/duplex` | Simpler float32 raw-PCM protocol |
|
||||
|
||||
### `/api/chat` Protocol
|
||||
|
||||
The first byte of every binary frame is a message tag:
|
||||
|
||||
| Direction | Tag | Payload |
|
||||
|-----------|-----|---------|
|
||||
| Server to client | `0x00` | Ready handshake |
|
||||
| Client to server | `0x01` | Opus microphone audio at 24 kHz |
|
||||
| Server to client | `0x01` | Opus response audio |
|
||||
| Server to client | `0x02` | UTF-8 inner-monologue text |
|
||||
|
||||
The optional `text_prompt` and `voice_prompt` query parameters select a persona
|
||||
and voice. The default server admits one conversation; use `--batch-size N` to
|
||||
configure multiple lockstep slots. Connections beyond capacity close with
|
||||
WebSocket code 1013.
|
||||
|
||||
### `/v1/audio/duplex` Protocol
|
||||
|
||||
Send `{"type":"open","persona":"...","voice":"..."}` as a JSON text frame,
|
||||
then exchange binary float32 PCM frames. The server returns generated float32
|
||||
audio as binary frames and text as `{"type":"text","text":"..."}`. Send
|
||||
`{"type":"close"}` and wait for `{"type":"done"}` to flush and finish.
|
||||
|
||||
See the [PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex)
|
||||
for the browser and headless clients.
|
||||
|
||||
## JoyVL Interaction Server
|
||||
|
||||
JoyVL is an orchestration layer, not a model-serving engine. It calls a
|
||||
separate OpenAI-compatible backend and adds per-session frame history, memory,
|
||||
persona policy, and optional delegation.
|
||||
|
||||
Start the model backend and orchestrator separately:
|
||||
|
||||
```bash
|
||||
vllm serve jdopensource/JoyAI-VL-Interaction-Preview \
|
||||
--served-model-name JoyAI-VL-Interaction-Preview \
|
||||
--port 8092 \
|
||||
--max-model-len 131072 \
|
||||
--enable-prefix-caching \
|
||||
--limit-mm-per-prompt '{"image":256,"video":1}'
|
||||
|
||||
python -m vllm_omni.experimental.fullduplex.joyvl.serving.server \
|
||||
--port 8091 \
|
||||
--main-backend-url http://127.0.0.1:8092/v1 \
|
||||
--main-model JoyAI-VL-Interaction-Preview
|
||||
```
|
||||
|
||||
| Method and route | Purpose |
|
||||
|------------------|---------|
|
||||
| `GET /health` | Readiness check |
|
||||
| `GET /v1/models` | Reports the configured interaction model |
|
||||
| `POST /v1/chat/completions` | Processes one frame or interaction tick |
|
||||
| `POST /reset` | Resets a session |
|
||||
| `POST /v1/streaming/reset` | Alias for `/reset` |
|
||||
| `POST /v1/streaming/persona` | Changes the session persona |
|
||||
|
||||
Send one multimodal Chat Completions request per video frame, normally around
|
||||
one frame per second, and identify the session with `x-session-id`. Responses
|
||||
include an `interaction` block whose action is `silence`, `response`, or
|
||||
`delegate`.
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8091/v1/chat/completions \
|
||||
-H 'content-type: application/json' \
|
||||
-H 'x-session-id: demo' \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": [
|
||||
{"type": "text", "text": "Alert me if a fire appears"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
|
||||
]}]
|
||||
}'
|
||||
```
|
||||
|
||||
See the [JoyVL recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/JD/JoyAI-VL-Interaction.md)
|
||||
for memory, personas, delegation backends, and the browser UI.
|
||||
@@ -0,0 +1,92 @@
|
||||
# Streaming Video Output WebSocket API
|
||||
|
||||
Use `WS /v1/realtime/video` to receive fragmented MP4 chunks while a compatible
|
||||
diffusion video model is generating. This endpoint streams generated video
|
||||
output; it is unrelated to [`WS /v1/video/chat/stream`](video_stream_api.md),
|
||||
which streams live video input into a conversational model.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Start a model that supports step execution and enable diffusion streaming
|
||||
output:
|
||||
|
||||
```bash
|
||||
vllm serve BestWishYsh/Helios-Distilled --omni \
|
||||
--diffusion-streaming-output \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Install the client dependency and run the provided client:
|
||||
|
||||
```bash
|
||||
pip install av websockets
|
||||
|
||||
python examples/online_serving/streaming_video_generation/streaming_video_client.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--prompt "A serene lakeside sunrise with mist over the water" \
|
||||
--output output.mp4
|
||||
```
|
||||
|
||||
## Protocol
|
||||
|
||||
Control events are JSON text frames. Each `video.chunk_metadata` event is
|
||||
immediately followed by one binary WebSocket frame containing fragmented MP4
|
||||
bytes.
|
||||
|
||||
| Direction | Event or frame | Purpose |
|
||||
|-----------|----------------|---------|
|
||||
| Client to server | `session.start` | Starts generation with a prompt and video parameters |
|
||||
| Server to client | `video.start` | Confirms the request ID, format, and accepted configuration |
|
||||
| Server to client | `video.chunk_metadata` | Describes the next binary media frame |
|
||||
| Server to client | Binary frame | Carries one fragmented MP4 chunk |
|
||||
| Client to server | `session.interaction` | Queues an optional prompt update during generation |
|
||||
| Server to client | `session.interaction.queued` | Confirms that an interaction was queued |
|
||||
| Client to server | `session.ping` | Refreshes the stall clock |
|
||||
| Server to client | `session.pong` | Acknowledges a ping |
|
||||
| Client to server | `session.stop` | Cancels the active generation |
|
||||
| Server to client | `session.done` | Reports completion or cancellation and the chunk count |
|
||||
| Server to client | `error` | Reports validation, generation, or timeout failures |
|
||||
|
||||
Start a minimal session with:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "session.start",
|
||||
"model": "BestWishYsh/Helios-Distilled",
|
||||
"prompt": "A serene lakeside sunrise",
|
||||
"format": "m4s",
|
||||
"width": 640,
|
||||
"height": 384,
|
||||
"fps": 16,
|
||||
"num_frames": 99
|
||||
}
|
||||
```
|
||||
|
||||
Only `m4s` output is currently accepted. Concatenate the binary frames in
|
||||
receive order, then remux the fragmented stream if the target player requires
|
||||
a progressive MP4 file.
|
||||
|
||||
## Mid-Generation Interaction
|
||||
|
||||
Models that implement interaction updates can change the active prompt:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "session.interaction",
|
||||
"interaction": {
|
||||
"event_id": "sunlight-update",
|
||||
"event": {"prompt": "Sunlight breaks through the mist"},
|
||||
"transition_chunks": 2
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
This feature is pipeline-dependent. An enabled endpoint does not imply that
|
||||
the loaded model supports midway prompt updates.
|
||||
|
||||
The default start timeout is 10 seconds and the server reports a stall after
|
||||
about 60 seconds without generation progress or a `session.ping`. See the
|
||||
[streaming video generation example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/streaming_video_generation)
|
||||
for prompt-update scheduling and browser playback.
|
||||
@@ -20,7 +20,7 @@ For the internal selector, registry, and platform contract, see
|
||||
| Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
|
||||
| Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
|
||||
| Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
|
||||
| Use FastVideo VSA with FastWan2.2-TI2V-5B or FastH3 MiniMax-H3 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
|
||||
|
||||
## Backend options
|
||||
|
||||
@@ -36,7 +36,7 @@ For the internal selector, registry, and platform contract, see
|
||||
| `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
|
||||
| `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
|
||||
| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
|
||||
| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B and FastH3 MiniMax-H3 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
|
||||
|
||||
## Configuration
|
||||
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
# Dense Attention Backends
|
||||
|
||||
Use these backends when you want dense attention without backend-specific
|
||||
sparsity. Start with `TORCH_SDPA` for correctness comparisons, then benchmark
|
||||
the fastest compatible kernel for your model, shape, and hardware.
|
||||
Use these backends to run dense attention. Optional quantized or sparse modes
|
||||
remain inactive unless they are explicitly configured. Start with
|
||||
`TORCH_SDPA` for correctness comparisons, then benchmark the fastest compatible
|
||||
kernel for your model, shape, and hardware.
|
||||
|
||||
For selection precedence and per-role configuration, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
@@ -20,10 +21,11 @@ vllm-omni serve <model> --diffusion-attention-backend TORCH_SDPA
|
||||
## `FLASH_ATTN`
|
||||
|
||||
`FLASH_ATTN` uses the installed FlashAttention implementation. On Blackwell it
|
||||
prefers the CuTe-based FlashAttention 4 path and falls back to
|
||||
FlashAttention 3 or 2 when that path is unavailable. On Hopper, Ada, and
|
||||
Ampere it is the preferred automatic route when a compatible package is
|
||||
installed.
|
||||
is FA4 only (`flash_attn.cute` / `vllm-omni[fa4]`). If FA4 is unavailable,
|
||||
explicit `FLASH_ATTN` is rejected and automatic selection continues to another
|
||||
compatible Blackwell backend. Hopper-only FA2/FA3 wheels are not used, even if
|
||||
they import. On Hopper, Ada, and Ampere it is the preferred automatic route
|
||||
when a compatible package is installed.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN
|
||||
@@ -38,8 +40,25 @@ pip install 'vllm-omni[fa4]'
|
||||
```
|
||||
|
||||
Version `4.0.0b18` is required; earlier beta wheels had known JIT failures on
|
||||
Blackwell. If the CuTe path is unavailable, the backend falls back to the
|
||||
compatible FlashAttention 3 or 2 path.
|
||||
Blackwell. On Blackwell, `FLASH_ATTN` is FA4 only: Hopper-only FA2/FA3
|
||||
wheels are not used, even if they import. If FA4 is missing, explicit
|
||||
`FLASH_ATTN` is rejected and automatic selection continues to other
|
||||
Blackwell backends.
|
||||
|
||||
## `TRTLLM_ATTN`
|
||||
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels and is the platform
|
||||
default on datacenter Blackwell for models that declare a compatible attention
|
||||
path. Selected without a `quant` or `skip_softmax` block, it runs dense BF16
|
||||
attention at FA4-level performance.
|
||||
|
||||
```bash
|
||||
vllm serve <model> --omni \
|
||||
--diffusion-attention-backend TRTLLM_ATTN
|
||||
```
|
||||
|
||||
See [TRTLLM Attention](trtllm.md) for its requirements and for the optional
|
||||
SAGE quantization and Skip-Softmax modes.
|
||||
|
||||
## `CUDNN_ATTN`
|
||||
|
||||
@@ -73,6 +92,13 @@ DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN \
|
||||
option on CUDA platforms and an automatic Blackwell fallback when FlashInfer
|
||||
is installed but cuDNN is too old for `CUDNN_ATTN`.
|
||||
|
||||
On Blackwell, `auto` resolves to FlashInfer cute-dsl, which cannot run a
|
||||
nontrivial custom mask. Automatic selection may fall back to SDPA for those
|
||||
masks. An explicit `FLASHINFER_ATTN` selection does not: sequence-parallel
|
||||
auto-padding that would create a padding mask is rejected at capability
|
||||
preflight. Use `TORCH_SDPA`, pin `quant.flashinfer_backend` to `fa2`/`fa3`
|
||||
where those kernels exist, or choose a mask-capable backend.
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend FLASHINFER_ATTN
|
||||
```
|
||||
|
||||
@@ -7,11 +7,18 @@ attention only against the selected top-k blocks.
|
||||
|
||||
The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
|
||||
|
||||
MiniMax-H3 served with a FastH3 VSA adapter uses a second, H3-specific route
|
||||
through the same backend. See the
|
||||
[MiniMax-H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for its geometry,
|
||||
supported topologies, and commands.
|
||||
|
||||
VSA is a CUDA-only, explicitly selected backend. It requires the
|
||||
`fastvideo-kernel` package and currently supports non-causal self-attention
|
||||
with equal query and key/value sequence lengths. Unsupported shapes, masks,
|
||||
dtypes, sequence-parallel execution, or kernel failures fall back to
|
||||
`TORCH_SDPA` and emit a warning with the reason.
|
||||
`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
|
||||
sequence-parallel context is one of those fallbacks; the H3 route supports pure
|
||||
Ulysses and rejects ring or all-gather sequence parallelism at startup.
|
||||
|
||||
## Enable the backend
|
||||
|
||||
@@ -119,7 +126,8 @@ the backend guarantees sparse execution:
|
||||
- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
|
||||
executed; the warning includes the reason.
|
||||
|
||||
The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks,
|
||||
The Wan route requires CUDA tensors in FP16 or BF16, 256-token blocks,
|
||||
standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
|
||||
and no active sequence-parallel context. NPU and XPU paths do not execute the
|
||||
FastVideo VSA CUDA kernel.
|
||||
and no active sequence-parallel context. The MiniMax-H3 route uses 64-token
|
||||
`(4, 4, 4)` blocks and supports pure Ulysses. NPU and XPU paths do not execute
|
||||
the FastVideo VSA CUDA kernel.
|
||||
@@ -12,12 +12,12 @@ implementations to reduce numerical drift and sampling divergence.
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
pip install kernels==0.14.1
|
||||
pip install kernels==0.15.2
|
||||
```
|
||||
|
||||
If `kernels` is unavailable, vLLM-Omni warns and falls back to the
|
||||
corresponding local FlashAttention implementation. On CUDA devices older than
|
||||
Hopper, `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`.
|
||||
If `kernels` is unavailable, an explicit `FLASH_ATTN_HUB` or
|
||||
`FLASH_ATTN_3_HUB` selection raises. Automatic / `FLASH_ATTN` selection in the
|
||||
Diffusers adapter may still walk hub then local FlashAttention backends.
|
||||
|
||||
## Usage
|
||||
|
||||
|
||||
@@ -1,103 +1,274 @@
|
||||
# TRTLLM Attention
|
||||
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels. It supports dense
|
||||
BF16 attention plus optional Skip-Softmax sparsity or SAGE quantization.
|
||||
`TRTLLM_ATTN` runs FlashInfer's trtllm-gen FMHA kernels on datacenter
|
||||
Blackwell GPUs. Selected on its own it computes dense BF16 attention, and its
|
||||
dense performance is on par with
|
||||
[FlashAttention-4](dense_backends.md#flashattention-4-on-blackwell). It also
|
||||
provides two opt-in, lossy acceleration modes that can be enabled independently
|
||||
or together:
|
||||
|
||||
For common selection and per-role configuration, see the
|
||||
[attention backend overview](../attention_backends.md).
|
||||
| Mode | Config key | What it changes |
|
||||
| --- | --- | --- |
|
||||
| [Skip-Softmax](#skip-softmax) | `skip_softmax` | Skips the Softmax and `PV` work of KV tiles whose scores are too low to matter |
|
||||
| [SAGE quantization](#sage-quantization) | `quant` | Runs `QK^T` in INT8 or FP8 and `PV` in FP8 instead of BF16 |
|
||||
|
||||
Both keys are set through `--diffusion-attention-config`, as JSON or as
|
||||
vLLM-style dotted flags; the
|
||||
[attention backend overview](../attention_backends.md#configuration) covers
|
||||
both syntaxes and per-role resolution. The examples on this page use JSON.
|
||||
|
||||
Attention computes scores `S = QK^T`, probabilities `P = softmax(S)`, and the
|
||||
output `O = PV`. SAGE lowers the precision of both matrix multiplications.
|
||||
Skip-Softmax keeps `QK^T` dense and removes the Softmax and `PV` work for
|
||||
tiles that `QK^T` shows to be unimportant. The two modes therefore compose:
|
||||
SAGE makes every tile cheaper, and Skip-Softmax reduces the number of tiles
|
||||
that reach the second half of the kernel.
|
||||
|
||||
## Requirements
|
||||
|
||||
Selecting `TRTLLM_ATTN` requires:
|
||||
`TRTLLM_ATTN` requires all of the following:
|
||||
|
||||
- datacenter Blackwell, sm_100 or sm_103;
|
||||
- an `sm100a` or `sm103a` GPU (B200, B300, GB200, GB300); workstation Blackwell
|
||||
(`sm120`/`sm121`) is not supported;
|
||||
- `head_dim=128`;
|
||||
- FlashInfer; and
|
||||
- a model-declared packed or mask-free compatible path.
|
||||
- a FlashInfer build that exposes the trtllm-gen kernels (0.6.16rc1 or newer
|
||||
for SAGE);
|
||||
- an attention path that is mask-free or provides packed-padding metadata.
|
||||
Structural suffix padding is expressed through that metadata rather than an
|
||||
`attn_mask` tensor.
|
||||
|
||||
It is not supported on workstation Blackwell (sm_120/sm_121). An incompatible
|
||||
explicit selection raises instead of silently changing the backend.
|
||||
An explicit selection that violates these requirements raises at startup
|
||||
instead of silently falling back to another backend.
|
||||
|
||||
## Dense mode
|
||||
### Sequence parallelism
|
||||
|
||||
Select `TRTLLM_ATTN` without `skip_softmax` or `quant` to run dense BF16:
|
||||
`TRTLLM_ATTN` runs under no sequence parallelism or under pure Ulysses.
|
||||
Ulysses redistributes the sequence and attention heads around the attention
|
||||
call, but the local computation still goes through the configured backend, so
|
||||
both optional modes work unchanged. Ring and AllGather-KV do not:
|
||||
|
||||
- Ring runs its own distributed attention and bypasses the backend. Combining
|
||||
Ring with a `skip_softmax` key raises; a `quant` key would be silently
|
||||
ignored, so do not combine them either.
|
||||
- AllGather-KV changes the Q/KV distribution and is rejected when
|
||||
`TRTLLM_ATTN` is selected.
|
||||
|
||||
Configure eight-way Ulysses alone as:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> --diffusion-attention-backend TRTLLM_ATTN
|
||||
--usp 8 --ring 1 --allgather-degree 1
|
||||
```
|
||||
|
||||
A degree of `1` disables that sequence-parallel mode. It does not limit the
|
||||
server to one GPU or affect tensor, pipeline, or VAE parallelism.
|
||||
|
||||
## Quick start
|
||||
|
||||
On datacenter Blackwell the platform selects `TRTLLM_ATTN` by default when the
|
||||
model declares a compatible path. To select it explicitly, for dense BF16:
|
||||
|
||||
```bash
|
||||
vllm serve <model> --omni \
|
||||
--diffusion-attention-backend TRTLLM_ATTN
|
||||
```
|
||||
|
||||
To enable both optimizations with a conservative Skip-Softmax setting and FP8
|
||||
SAGE:
|
||||
|
||||
```bash
|
||||
vllm serve <model> --omni \
|
||||
--diffusion-attention-config '{
|
||||
"default": {
|
||||
"backend": "TRTLLM_ATTN",
|
||||
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16},
|
||||
"skip_softmax": {"threshold": 0.05, "disabled_until_timestep": 0.97}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
The startup log reports the selection; look for one of:
|
||||
|
||||
```text
|
||||
Defaulting to diffusion attention backend TRTLLM_ATTN (datacenter Blackwell ..., head_dim 128)
|
||||
Resolved diffusion attention backend 'TRTLLM_ATTN' for role='self' via attention_config.default
|
||||
```
|
||||
|
||||
## Skip-Softmax
|
||||
|
||||
Skip-Softmax trades some fidelity for speed by skipping low-contribution
|
||||
attention work. See the [feature design](../../../design/feature/skip_softmax.md)
|
||||
for the algorithm and implementation contract.
|
||||
Skip-Softmax (BLASST) skips the Softmax and `PV` work of KV tiles whose
|
||||
scores fall far below the running row maximum. `QK^T` always runs, and how
|
||||
many tiles qualify depends on the input. The
|
||||
[feature design](../../../design/feature/skip_softmax.md) gives the skip test.
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
### Configuration keys
|
||||
|
||||
| Key | Range | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `target_sparsity` | finite, `[0, 1]` | Point on a model-calibrated curve |
|
||||
| `threshold` | finite, `>= 0` | Direct threshold; mutually exclusive with `target_sparsity` |
|
||||
| `disabled_until_timestep` | finite, `[0, 1]` | Keep early high-noise steps dense until normalized `t <= D` |
|
||||
| `threshold` | `>= 0`; useful values in `(0, 1)` | Skip threshold, independent of sequence length. Calibration-free. |
|
||||
| `target_sparsity` | `[0, 1]` | Requested operating point on the checkpoint's calibrated curve. Requires calibration metadata. |
|
||||
| `disabled_until_timestep` | `[0, 1]`; default `0` | Keeps attention dense while the normalized timestep `t > D`; `0` disables the gate. |
|
||||
|
||||
`threshold` and `target_sparsity` are two ways to set the same kernel
|
||||
threshold; setting both is a configuration error. Exactly one of them enables
|
||||
Skip-Softmax.
|
||||
|
||||
### Direct threshold
|
||||
|
||||
Set `threshold` when the checkpoint carries no calibration. `threshold=0`
|
||||
skips nothing; larger values skip more tiles and lower output fidelity. Values
|
||||
around `0.05` are a reasonable first try for video DiTs; tune against dense
|
||||
output on the same prompt and seed. The value is independent of sequence
|
||||
length; see the
|
||||
[feature design](../../../design/feature/skip_softmax.md#from-configuration-to-the-kernel-threshold)
|
||||
for how it reaches the kernel.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
|
||||
"skip_softmax":{"target_sparsity":0.65,"disabled_until_timestep":0.86}}}'
|
||||
vllm serve <model> --omni \
|
||||
--diffusion-attention-config \
|
||||
'{"default":{"backend":"TRTLLM_ATTN","skip_softmax":{
|
||||
"threshold":0.05,"disabled_until_timestep":0.97}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, SkipSoftmaxSpec
|
||||
### Calibrated target sparsity
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
skip_softmax=SkipSoftmaxSpec(
|
||||
target_sparsity=0.65,
|
||||
disabled_until_timestep=0.86,
|
||||
),
|
||||
),
|
||||
)
|
||||
A fixed `threshold` does not produce a fixed fraction of skipped tiles because
|
||||
the score distribution changes with the model, prompt, and shape.
|
||||
[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer)
|
||||
can calibrate a per-model curve from sparsity to threshold and store it in the
|
||||
checkpoint's transformer `config.json` under `sparse_attention_config`.
|
||||
`target_sparsity` selects a point on that curve. The achieved sparsity still
|
||||
varies per prompt, layer, and denoising step; the calibration makes the
|
||||
requested value a meaningful target, not a guarantee.
|
||||
|
||||
What vLLM-Omni takes from the checkpoint:
|
||||
|
||||
- The curve coefficients. Only the `a * exp(b * target_sparsity)` form ModelOpt
|
||||
writes is supported; another formula is rejected at startup.
|
||||
- The `ignore` list: attention layers matching those patterns stay dense.
|
||||
- For multi-expert Diffusers checkpoints, `transformer_2/config.json` is read
|
||||
separately; if it is missing, `transformer_2` stays dense with a warning.
|
||||
|
||||
Checkpoint-level `target_sparsity` and `disabled_until_timestep` defaults are
|
||||
not consumed; set them in the vLLM-Omni configuration. The
|
||||
[feature design](../../../design/feature/skip_softmax.md#skip-softmax-calibration-config)
|
||||
documents the checkpoint format and how the coefficients reach the kernel.
|
||||
|
||||
To calibrate a checkpoint yourself, follow the
|
||||
[ModelOpt Skip-Softmax example](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity),
|
||||
which exports a copy of the checkpoint with the coefficients embedded. The
|
||||
fitted curve depends on the attention statistics of the calibration shape, so
|
||||
calibrate at the resolution and frame count you will serve.
|
||||
|
||||
Requesting `target_sparsity` for a checkpoint without calibration is a startup
|
||||
error that names the `threshold` alternative. The
|
||||
[ModelOpt Wan2.2 FP8 checkpoint](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8/blob/main/transformer/config.json)
|
||||
is a calibrated example:
|
||||
|
||||
```bash
|
||||
vllm serve nvidia/Wan2.2-T2V-A14B-Diffusers-FP8 --omni \
|
||||
--diffusion-attention-config \
|
||||
'{"default":{"backend":"TRTLLM_ATTN","skip_softmax":{
|
||||
"target_sparsity":0.75,"disabled_until_timestep":0.86}}}'
|
||||
```
|
||||
|
||||
Start with `target_sparsity=0.65` and `disabled_until_timestep=0.86`. Increase
|
||||
sparsity only after comparing quality against dense output at the same seed.
|
||||
The benefit grows with sequence length because only attention is accelerated.
|
||||
### Timestep gating
|
||||
|
||||
The early, high-noise denoising steps fix the global layout of the output, and
|
||||
their errors propagate through every later step. `disabled_until_timestep`
|
||||
keeps those steps dense and enables Skip-Softmax once the normalized timestep
|
||||
`t` reaches the configured cutoff. The default `0` is a sentinel that turns the
|
||||
gate off: Skip-Softmax runs on every step and no timestep needs to be
|
||||
published. `1.0` also leaves no step dense, but goes through the gate and
|
||||
therefore requires the pipeline to publish `t`.
|
||||
|
||||
`t` is the scheduler's own timestep normalized to `[0, 1]`, published by the
|
||||
pipeline for each denoising step; for rectified-flow models it is the current
|
||||
sigma. It runs from near `1.0` down to `0.0`, but not linearly in the step
|
||||
index: flow-shifted schedules spend many of their steps at high `t`. The number
|
||||
of dense steps a given cutoff produces therefore depends on the model's
|
||||
schedule and step count, and follows from the actual timestep sequence:
|
||||
|
||||
```text
|
||||
dense_steps = count(t[i] > disabled_until_timestep)
|
||||
```
|
||||
|
||||
Derive the cutoff from the schedule of the model you are serving rather than
|
||||
reusing another model's value; the
|
||||
[feature design](../../../design/feature/skip_softmax.md#mapping-the-cutoff-to-denoising-steps)
|
||||
works through an example.
|
||||
|
||||
A pipeline that does not publish `t` stays dense whenever
|
||||
`disabled_until_timestep > 0` is set, and logs a warning once. Pipelines
|
||||
publish it through `DenoiseProgressMixin.record_denoise_step`.
|
||||
|
||||
## SAGE quantization
|
||||
|
||||
TRTLLM SAGE quantizes Q/K per block and V per channel so both attention
|
||||
matrix multiplications use low precision.
|
||||
SAGE quantization runs both attention matrix multiplications in low precision:
|
||||
Q and K are quantized to INT8 or FP8 E4M3 for `QK^T`, and P and V use FP8 E4M3
|
||||
for `PV`. P is quantized inside the FMHA kernel; V is quantized per channel
|
||||
before the kernel call.
|
||||
vLLM-Omni exposes the Q/K dtype and the Q/K scale granularity. The P and V
|
||||
formats are fixed by the kernel, so the `quant` key has no V dtype for this
|
||||
backend. This mode is distinct from the standalone
|
||||
[SageAttention backends](sage.md), which use their own kernels.
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
FP8 Q/K kernels exist on `sm100a` and `sm103a`; INT8 Q/K kernels exist on
|
||||
`sm100a` only.
|
||||
|
||||
| Key | Values | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype; absent means dense |
|
||||
| `q_block_size` | `1`, `4`, `16` | Q scale block size; default `1` |
|
||||
| `k_block_size` | `1`, `4`, `16` | K scale block size; default `16` |
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Q/K quantization dtype. Setting it enables SAGE. |
|
||||
| `q_block_size` | `1`, `4`, `16` | Consecutive query tokens sharing one Q scale; default `1` |
|
||||
| `k_block_size` | `1`, `4`, `16` | Consecutive key tokens sharing one K scale; default `16` |
|
||||
|
||||
V is always quantized per channel to FP8 E4M3 and K-smoothing is internal.
|
||||
Only block sizes with compiled kernels are accepted.
|
||||
Start with the defaults, `q_block_size=1` and `k_block_size=16`, and try
|
||||
smaller K blocks only if output quality does not meet expectations; smaller
|
||||
blocks give finer scales at some cost in speed. When a KV sequence in a call is
|
||||
shorter than `k_block_size`, that call falls back to dense attention and a
|
||||
warning is logged once.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default":{"backend":"TRTLLM_ATTN",\
|
||||
"quant":{"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
|
||||
vllm serve <model> --omni \
|
||||
--diffusion-attention-config \
|
||||
'{"default":{"backend":"TRTLLM_ATTN","quant":{
|
||||
"dtype_qk":"fp8_e4m3","q_block_size":1,"k_block_size":16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
The `quant` key is shared with `FLASHINFER_ATTN`, but each backend validates
|
||||
its own fields: `float16`/`bfloat16` Q/K dtypes and `dtype_vo` are
|
||||
`FLASHINFER_ATTN` options and are rejected here.
|
||||
|
||||
config = AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(
|
||||
dtype_qk="fp8_e4m3",
|
||||
q_block_size=1,
|
||||
k_block_size=16,
|
||||
),
|
||||
),
|
||||
)
|
||||
## Composing both modes
|
||||
|
||||
`skip_softmax` and `quant` may appear in the same `AttentionSpec`. Their
|
||||
quality effects compound, so establish a dense baseline, enable one mode at a
|
||||
time, and then evaluate the combination on the same prompts and seeds.
|
||||
|
||||
Modes configured in `default` apply to every attention role that has no
|
||||
`per_role` entry. A `per_role` spec replaces the whole spec for that role and
|
||||
does not inherit `quant` or `skip_softmax` from `default`, so
|
||||
`{"backend":"TRTLLM_ATTN"}` is the way to keep a short or sensitive attention
|
||||
site dense while the long DiT sequence uses both modes:
|
||||
|
||||
```bash
|
||||
vllm serve MiniMaxAI/MiniMax-H3 --omni \
|
||||
--diffusion-attention-config '{
|
||||
"default": {
|
||||
"backend": "TRTLLM_ATTN",
|
||||
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16},
|
||||
"skip_softmax": {"threshold": 0.05, "disabled_until_timestep": 0.97}
|
||||
},
|
||||
"per_role": {
|
||||
"minimax_h3.token_refiner": {"backend": "TRTLLM_ATTN"}
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
This path requires FlashInfer 0.6.16rc1 or newer. FP8 Q/K kernels are
|
||||
available on sm_100 and sm_103; INT8 Q/K kernels are compiled for sm_100 only.
|
||||
The shared `AttnQuantSpec` is backend typed: values intended for
|
||||
`FLASHINFER_ATTN` are rejected here and vice versa.
|
||||
Role names are declared by each model; the
|
||||
[attention backend overview](../attention_backends.md#configuration) covers
|
||||
the resolution order and the equivalent Python API.
|
||||
|
||||
End-to-end speedup depends on the share of step time spent in attention, the
|
||||
sequence length, the chosen Q/K precision, and the tile sparsity the input
|
||||
actually yields. Benchmark the exact workload rather than extrapolating from
|
||||
another model.
|
||||
@@ -94,7 +94,7 @@ In `DiffusionParallelConfig`
|
||||
|
||||
|
||||
!!! info
|
||||
Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** is an exception: it supports `cfg_parallel_size=3`, which adds a third branch on rank 2 for full three-way CFG parallelism.
|
||||
Most models support `cfg_parallel_size=2` (positive branch on rank 0, negative branch on rank 1). **Bagel** and **Boogu-Image-Edit double guidance** also support `cfg_parallel_size=3` for full three-way branch execution. Boogu double guidance additionally supports size 2 via round-robin assignment of its three branches.
|
||||
|
||||
|
||||
---
|
||||
|
||||
@@ -117,6 +117,7 @@ The following tables show which models support each feature:
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Boogu-Image** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
@@ -138,7 +139,7 @@ The following tables show which models support each feature:
|
||||
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **SenseNova-U1** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **SenseNova-U1 / U1.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
||||
@@ -162,7 +163,7 @@ The following tables show which models support each feature:
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
|
||||
| **SANA-Video-2B T2V I2V** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ |
|
||||
| **SANA-Video-2B T2V I2V** | ❌ | ❌ | ❌ | ✅ | ✅ (TP=2 only) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/text_to_speech>.
|
||||
|
||||
|
||||
vLLM-Omni exposes TTS models through the OpenAI-compatible
|
||||
[`POST /v1/audio/speech`](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md) endpoint,
|
||||
launched with `vllm serve <model> --omni`. Each TTS model has its own
|
||||
@@ -16,7 +15,7 @@ For the full list of supported architectures across all modalities, see
|
||||
## Supported Models
|
||||
|
||||
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
|
||||
|---|---|---|---|---|---|
|
||||
| --- | --- | --- | --- | --- | --- |
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
|
||||
@@ -104,6 +103,7 @@ For full request-shape documentation (all parameters, response formats, error co
|
||||
2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
|
||||
```
|
||||
@@ -111,6 +111,7 @@ huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
|
||||
If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class).
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
|
||||
# or:
|
||||
@@ -120,13 +121,16 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
|
||||
Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
|
||||
|
||||
### CLI client
|
||||
|
||||
The client defaults to the official upstream zero-shot prompt, so it runs without extra flags:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/cosyvoice3
|
||||
python speech_client.py --text "收到好友从远方寄来的生日礼物。"
|
||||
```
|
||||
|
||||
Pass your own reference clip and transcript for a different voice:
|
||||
|
||||
```bash
|
||||
python speech_client.py --text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
@@ -134,6 +138,7 @@ python speech_client.py --text "Hello, this is a cloned voice." \
|
||||
```
|
||||
|
||||
Stream PCM instead of WAV:
|
||||
|
||||
```bash
|
||||
python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
```
|
||||
@@ -141,6 +146,7 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`.
|
||||
|
||||
### Notes
|
||||
|
||||
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
|
||||
- For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
|
||||
@@ -152,19 +158,23 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
# or:
|
||||
./fish_speech/run_server.sh
|
||||
```
|
||||
|
||||
The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`).
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
@@ -177,6 +187,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
```
|
||||
|
||||
### CLI client
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/fish_speech
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
@@ -184,12 +195,14 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
```
|
||||
|
||||
### Gradio demo
|
||||
|
||||
```bash
|
||||
./fish_speech/run_gradio_demo.sh # launches server + Gradio
|
||||
python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server already running
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 44.1 kHz mono.
|
||||
- Streaming PCM player command must use `-r 44100`.
|
||||
|
||||
@@ -208,12 +221,15 @@ pip install -U "transformers>=5.3.0"
|
||||
```
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
|
||||
```
|
||||
|
||||
Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# Plain TTS
|
||||
python higgs_audio_v2/batch_speech_client.py \
|
||||
@@ -231,14 +247,17 @@ python higgs_audio_v2/batch_speech_client.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 24 kHz mono.
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
|
||||
- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
|
||||
|
||||
## GLM-TTS
|
||||
|
||||
2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
|
||||
# or:
|
||||
@@ -246,6 +265,7 @@ bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-T
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# Voice cloning (required)
|
||||
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
@@ -262,11 +282,13 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
```
|
||||
|
||||
### Gradio demo
|
||||
|
||||
```bash
|
||||
bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 24 kHz mono WAV via HiFT vocoder.
|
||||
- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
|
||||
- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
|
||||
@@ -278,12 +300,15 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
|
||||
Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
huggingface-cli download k2-fsa/OmniVoice
|
||||
```
|
||||
|
||||
Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
|
||||
# or:
|
||||
@@ -291,6 +316,7 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
|
||||
```
|
||||
|
||||
### CLI client
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/omnivoice
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
@@ -300,6 +326,7 @@ python speech_client.py --text "Bonjour, comment allez-vous?" --language French
|
||||
The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`.
|
||||
|
||||
### Notes
|
||||
|
||||
- Voice cloning and voice design require offline inference; see the [offline OmniVoice section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#omnivoice).
|
||||
|
||||
---
|
||||
@@ -309,7 +336,7 @@ The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--l
|
||||
Three model variants exposed via separate checkpoints:
|
||||
|
||||
| Variant | HF repo | Use |
|
||||
|---|---|---|
|
||||
| --- | --- | --- |
|
||||
| CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions |
|
||||
| VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description |
|
||||
| Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio |
|
||||
@@ -317,6 +344,7 @@ Three model variants exposed via separate checkpoints:
|
||||
Each variant ships smaller `0.6B` companions where available.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
|
||||
# or:
|
||||
@@ -326,27 +354,30 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
|
||||
```
|
||||
|
||||
### Executor backend
|
||||
|
||||
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
|
||||
|
||||
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
|
||||
|
||||
### Tuning stage 1 `max_num_seqs` per task type
|
||||
|
||||
The bundled `qwen3_tts.yaml` ships stage 1 (Code2Wav) at `max_num_seqs: 10`, tuned for Base voice cloning: stage-1 lifetimes are long (~3 s/req), so admitting up to 10 concurrent codec sequences lets requests progress in parallel in the scheduler — ~2× TTFA p95 at c=4 / c=8 (1× H100, 1.7B-Base, seed-tts) at an 8–12 % audio-throughput cost.
|
||||
|
||||
CustomVoice / VoiceDesign have much shorter stage-1 lifetimes (~50–200 ms) and are TTFA-optimal at `max_num_seqs: 1`. Override the default when serving those task types:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
|
||||
--stage-overrides '{"1": {"max_num_seqs": 1}}'
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# CustomVoice with a predefined speaker
|
||||
python qwen3_tts/openai_speech_client.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--text "今天天气真好" \
|
||||
--voice ryan \
|
||||
--speaker ryan \
|
||||
--instructions "用开心的语气说"
|
||||
|
||||
# VoiceDesign with a style description
|
||||
@@ -366,7 +397,9 @@ python qwen3_tts/openai_speech_client.py \
|
||||
```
|
||||
|
||||
### Voices endpoint
|
||||
|
||||
List available voices, or upload a custom one for Base cloning:
|
||||
|
||||
```bash
|
||||
# List
|
||||
curl http://localhost:8091/v1/audio/voices
|
||||
@@ -379,10 +412,18 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "ref_text=The exact transcript of the audio sample." \
|
||||
-F "speaker_description=warm narrator"
|
||||
```
|
||||
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
|
||||
|
||||
For Qwen3-TTS, uploaded voices are Base voice-cloning inputs and require a Base
|
||||
checkpoint. When a request names an uploaded voice, the server infers
|
||||
`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
|
||||
CustomVoice speakers and require a CustomVoice checkpoint.
|
||||
|
||||
### Precomputed custom voices
|
||||
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
|
||||
|
||||
For reused Base voice-cloning speakers, precompute the reference artifacts once
|
||||
and load them at server startup. Precomputed voices use the same Base task and
|
||||
checkpoint-matching rules as uploaded voices:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/precompute_custom_voice.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
@@ -392,11 +433,15 @@ python qwen3_tts/precompute_custom_voice.py \
|
||||
--mode icl \
|
||||
--output-dir /path/to/custom_voices
|
||||
```
|
||||
|
||||
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
|
||||
|
||||
```yaml
|
||||
custom_voice_dir: /path/to/custom_voices
|
||||
```
|
||||
|
||||
Then start the server with that config and call the Speech API with only the voice name:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
|
||||
|
||||
@@ -407,6 +452,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
```
|
||||
|
||||
### Streaming PCM
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
@@ -419,16 +465,20 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
|
||||
|
||||
### Streaming WebSocket
|
||||
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "First utterance." \
|
||||
@@ -436,6 +486,7 @@ python qwen3_tts/streaming_speech_client.py \
|
||||
```
|
||||
|
||||
### Gradio demos
|
||||
|
||||
```bash
|
||||
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
|
||||
./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign
|
||||
@@ -447,12 +498,15 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
```
|
||||
|
||||
### Speaker embedding interpolation
|
||||
|
||||
`qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage.
|
||||
|
||||
### Batch client
|
||||
|
||||
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
|
||||
|
||||
### Notes
|
||||
|
||||
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
|
||||
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
|
||||
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
||||
@@ -467,12 +521,15 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
Single-stage native AR TTS at 48 kHz.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000
|
||||
```
|
||||
|
||||
Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks.
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# Zero-shot synthesis
|
||||
python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2."
|
||||
@@ -482,10 +539,13 @@ python voxcpm2/openai_speech_client.py \
|
||||
--text "This should sound like the reference speaker." \
|
||||
--ref-audio /path/to/reference.wav
|
||||
```
|
||||
|
||||
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
|
||||
|
||||
### Precomputed custom voices
|
||||
|
||||
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
|
||||
|
||||
```bash
|
||||
python voxcpm2/precompute_custom_voice.py \
|
||||
--model openbmb/VoxCPM2 \
|
||||
@@ -495,16 +555,21 @@ python voxcpm2/precompute_custom_voice.py \
|
||||
--prompt-text "Original transcript of the reference audio" \
|
||||
--output-dir /path/to/custom_voices
|
||||
```
|
||||
|
||||
Add the output directory to the deploy config:
|
||||
|
||||
```yaml
|
||||
custom_voice_dir: /path/to/custom_voices
|
||||
```
|
||||
|
||||
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
|
||||
|
||||
### Gradio demo (gapless streaming via AudioWorklet)
|
||||
|
||||
```bash
|
||||
python voxcpm2/gradio_demo.py
|
||||
```
|
||||
|
||||
Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`.
|
||||
|
||||
---
|
||||
@@ -514,24 +579,31 @@ Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free p
|
||||
Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Latest `mistral_common` with `SpeechRequest` support:
|
||||
|
||||
```bash
|
||||
pip install -e /path/to/mistral-common # or upgrade from PyPI when available
|
||||
```
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
|
||||
```
|
||||
|
||||
Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`.
|
||||
|
||||
### Gradio demo
|
||||
|
||||
```bash
|
||||
python voxtral_tts/gradio_demo.py
|
||||
```
|
||||
|
||||
The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients).
|
||||
|
||||
### Notes
|
||||
|
||||
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
|
||||
- Voice cloning is gated upstream and may require a recent `mistral_common`.
|
||||
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
|
||||
|
||||
@@ -46,20 +46,20 @@ More failure modes (for example, network jitter and network interruption) will b
|
||||
## Fault Injection Scenario Matrix
|
||||
|
||||
| Scenario | Fault Type | System Behavior | Current Status |
|
||||
|------|----------------------------|----------|----------|
|
||||
| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | |
|
||||
| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | |
|
||||
| ------ | ---------------------------- | ---------- | ---------- |
|
||||
| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | |
|
||||
| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | |
|
||||
| No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725) <br>[#43060](https://github.com/vllm-project/vllm/issues/43060) |
|
||||
| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released | |
|
||||
| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released | |
|
||||
| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly | |
|
||||
| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed | |
|
||||
| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released | |
|
||||
| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released | |
|
||||
| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly | |
|
||||
| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed | |
|
||||
| Under load | Send `SIGKILL` to Worker | In-flight requests are hard interrupted (5xx/connection drop); main service becomes unavailable | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable | |
|
||||
| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable | |
|
||||
| Under load | Send `SIGKILL` to serve main process | serve is hard-killed and current connections are interrupted; in-flight requests fail; after cleanup there are no residual processes and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGTERM` to serve main process | serve stops accepting new requests and executes shutdown flow; in-flight requests fail; no residue remains and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | `Ctrl+C`-style serve shutdown; in-flight requests fail (5xx/connection interruption); service unavailable; after exit there is no residue and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly | |
|
||||
| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly | |
|
||||
| Under load | Send `SIGTERM` to all related processes | All processes exit gracefully; in-flight requests fail; service unavailable; GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| OOM | Occupy all free GPU memory via an extra process | After OOM injection process starts, GPU memory is continuously saturated; service enters unavailable/degraded state and health check drops to 503; different request types (chat/speech, etc.) fail fast within a fixed time and return 500 (no hanging) | [#4285](https://github.com/vllm-project/vllm-omni/issues/4285) |
|
||||
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Run Mage-VL offline end-to-end inference through the upstream Mage scripts.
|
||||
|
||||
Mage-VL currently ships its custom Mage-ViT and codec processors as
|
||||
Transformers remote code in the Microsoft Mage repository. This example keeps
|
||||
the vLLM-Omni entrypoint stable while delegating the checkpoint-specific forward
|
||||
path to that upstream implementation.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from importlib.util import find_spec
|
||||
from pathlib import Path
|
||||
|
||||
DEFAULT_MODEL = "microsoft/Mage-VL"
|
||||
DEFAULT_IMAGE_QUESTION = "Describe this image in detail."
|
||||
DEFAULT_VIDEO_QUESTION = "Describe this video."
|
||||
|
||||
|
||||
def _default_mage_repo() -> Path | None:
|
||||
value = os.getenv("MAGE_REPO_DIR")
|
||||
if value:
|
||||
return Path(value)
|
||||
sibling = Path(__file__).resolve().parents[3].parent / "Mage"
|
||||
if sibling.exists():
|
||||
return sibling
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_mage_repo(value: str | None) -> Path:
|
||||
repo = Path(value) if value else _default_mage_repo()
|
||||
if repo is None:
|
||||
raise ValueError("Set --mage-repo or MAGE_REPO_DIR to a checkout of https://github.com/microsoft/Mage.")
|
||||
repo = repo.expanduser().resolve()
|
||||
if not (repo / "mage_vl" / "inference_base.py").is_file():
|
||||
raise FileNotFoundError(f"Missing Mage-VL scripts under {repo}. Expected mage_vl/inference_base.py.")
|
||||
return repo
|
||||
|
||||
|
||||
def _default_media_path(mage_repo: Path, task: str) -> Path:
|
||||
if task == "image":
|
||||
return mage_repo / "mage_vl" / "assets" / "examples" / "dog.jpg"
|
||||
return mage_repo / "mage_vl" / "assets" / "examples" / "soccer-broadcast.mp4"
|
||||
|
||||
|
||||
def _resolve_media_path(path: str | None, mage_repo: Path, task: str) -> Path:
|
||||
resolved = Path(path).expanduser().resolve() if path else _default_media_path(mage_repo, task)
|
||||
if not resolved.is_file():
|
||||
raise FileNotFoundError(f"Input media file not found: {resolved}")
|
||||
return resolved
|
||||
|
||||
|
||||
def build_command(args: argparse.Namespace, mage_repo: Path, media_path: Path) -> list[str]:
|
||||
if args.task == "streaming":
|
||||
return [
|
||||
sys.executable,
|
||||
str(mage_repo / "mage_vl" / "inference_streaming.py"),
|
||||
"--video",
|
||||
str(media_path),
|
||||
"--checkpoint",
|
||||
args.model,
|
||||
"--video_backend",
|
||||
args.video_backend,
|
||||
"--num_frames",
|
||||
str(args.num_frames),
|
||||
"--cur_fps",
|
||||
str(args.cur_fps),
|
||||
"--segment_sec",
|
||||
str(args.segment_sec),
|
||||
"--max_new_tokens",
|
||||
str(args.max_new_tokens),
|
||||
"--max_segments",
|
||||
str(args.max_segments),
|
||||
"--gate_threshold",
|
||||
str(args.gate_threshold),
|
||||
"--device",
|
||||
args.device,
|
||||
"--attn_impl",
|
||||
args.attn_impl,
|
||||
]
|
||||
|
||||
cmd = [
|
||||
sys.executable,
|
||||
str(mage_repo / "mage_vl" / "inference_base.py"),
|
||||
"--mode",
|
||||
"offline",
|
||||
"--question",
|
||||
args.question or (DEFAULT_IMAGE_QUESTION if args.task == "image" else DEFAULT_VIDEO_QUESTION),
|
||||
"--model",
|
||||
args.model,
|
||||
"--max-new-tokens",
|
||||
str(args.max_new_tokens),
|
||||
]
|
||||
if args.task == "image":
|
||||
cmd.extend(["--image", str(media_path)])
|
||||
return cmd
|
||||
|
||||
cmd.extend(
|
||||
[
|
||||
"--video",
|
||||
str(media_path),
|
||||
"--video-backend",
|
||||
args.video_backend,
|
||||
"--num-frames",
|
||||
str(args.num_frames),
|
||||
"--max-pixels",
|
||||
str(args.max_pixels),
|
||||
]
|
||||
)
|
||||
if args.video_backend == "codec":
|
||||
cmd.extend(["--codec-engine", args.codec_engine])
|
||||
return cmd
|
||||
|
||||
|
||||
def check_runtime_dependencies(task: str) -> None:
|
||||
required = ["torch"]
|
||||
if task == "streaming":
|
||||
required.append("mamba_ssm")
|
||||
missing = [name for name in required if find_spec(name) is None]
|
||||
if not missing:
|
||||
return
|
||||
|
||||
install_hint = (
|
||||
"Install Mage-VL runtime dependencies in this environment first:\n"
|
||||
" pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n"
|
||||
" grep -v '^mamba-ssm' /path/to/Mage/mage_vl/requirements.txt > /tmp/mage_vl_requirements_no_mamba.txt\n"
|
||||
" pip install -r /tmp/mage_vl_requirements_no_mamba.txt\n"
|
||||
" pip install 'causal-conv1d>=1.4.0' --no-build-isolation\n"
|
||||
" pip install 'mamba-ssm>=2.2' --no-build-isolation"
|
||||
)
|
||||
raise RuntimeError(f"Missing Python packages: {', '.join(missing)}.\n{install_hint}")
|
||||
|
||||
|
||||
def run_command(cmd: list[str], cwd: Path, timeout: int) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
cmd,
|
||||
cwd=str(cwd),
|
||||
input="y\n",
|
||||
text=True,
|
||||
capture_output=True,
|
||||
timeout=timeout,
|
||||
check=False,
|
||||
)
|
||||
|
||||
|
||||
def write_outputs(
|
||||
*,
|
||||
args: argparse.Namespace,
|
||||
mage_repo: Path,
|
||||
media_path: Path,
|
||||
cmd: list[str],
|
||||
result: subprocess.CompletedProcess[str] | None,
|
||||
elapsed_s: float | None,
|
||||
) -> Path:
|
||||
output_dir = Path(args.output_dir).expanduser().resolve()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
summary = {
|
||||
"task": args.task,
|
||||
"model": args.model,
|
||||
"mage_repo": str(mage_repo),
|
||||
"media_path": str(media_path),
|
||||
"command": cmd,
|
||||
"dry_run": args.dry_run,
|
||||
"elapsed_s": elapsed_s,
|
||||
"returncode": None if result is None else result.returncode,
|
||||
}
|
||||
if result is not None:
|
||||
stdout = result.stdout.strip()
|
||||
stderr = result.stderr.strip()
|
||||
summary["stdout"] = stdout
|
||||
summary["stderr"] = stderr
|
||||
(output_dir / "stdout.txt").write_text(result.stdout, encoding="utf-8")
|
||||
(output_dir / "stderr.txt").write_text(result.stderr, encoding="utf-8")
|
||||
summary_path = output_dir / "summary.json"
|
||||
summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
return summary_path
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Mage-VL offline end-to-end inference example.")
|
||||
parser.add_argument(
|
||||
"--mage-repo",
|
||||
default=None,
|
||||
help="Path to a Microsoft Mage checkout. Defaults to MAGE_REPO_DIR.",
|
||||
)
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL, help="HF model id or local checkpoint path.")
|
||||
parser.add_argument("--task", choices=("image", "video", "streaming"), default="image")
|
||||
parser.add_argument("--image-path", default=None, help="Image input. Defaults to the Mage sample image.")
|
||||
parser.add_argument("--video-path", default=None, help="Video input. Defaults to the Mage sample video.")
|
||||
parser.add_argument("--question", default=None, help="Question for image/video offline QA.")
|
||||
parser.add_argument("--video-backend", choices=("frames", "codec"), default="frames")
|
||||
parser.add_argument("--codec-engine", choices=("traditional", "neural"), default="traditional")
|
||||
parser.add_argument("--num-frames", type=int, default=32)
|
||||
parser.add_argument("--max-pixels", type=int, default=150000)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=256)
|
||||
parser.add_argument("--cur-fps", type=float, default=2.0)
|
||||
parser.add_argument("--segment-sec", type=float, default=8.0)
|
||||
parser.add_argument("--max-segments", type=int, default=0)
|
||||
parser.add_argument("--gate-threshold", type=float, default=0.5)
|
||||
parser.add_argument("--device", default="cuda")
|
||||
parser.add_argument("--attn-impl", choices=("flash_attention_2", "sdpa", "eager"), default="flash_attention_2")
|
||||
parser.add_argument("--timeout", type=int, default=1800)
|
||||
parser.add_argument("--output-dir", default="outputs/mage_vl")
|
||||
parser.add_argument("--dry-run", action="store_true", help="Build the command and summary without executing it.")
|
||||
args = parser.parse_args()
|
||||
if args.num_frames <= 0:
|
||||
parser.error("--num-frames must be positive")
|
||||
if args.segment_sec <= 0:
|
||||
parser.error("--segment-sec must be positive")
|
||||
if args.task == "image" and args.video_path:
|
||||
parser.error("--video-path cannot be used with --task image")
|
||||
if args.task in {"video", "streaming"} and args.image_path:
|
||||
parser.error("--image-path cannot be used with video tasks")
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
mage_repo = _resolve_mage_repo(args.mage_repo)
|
||||
media_arg = args.image_path if args.task == "image" else args.video_path
|
||||
media_path = _resolve_media_path(media_arg, mage_repo, args.task)
|
||||
cmd = build_command(args, mage_repo, media_path)
|
||||
|
||||
result = None
|
||||
elapsed_s = None
|
||||
if not args.dry_run:
|
||||
check_runtime_dependencies(args.task)
|
||||
start = time.perf_counter()
|
||||
result = run_command(cmd, cwd=mage_repo, timeout=args.timeout)
|
||||
elapsed_s = time.perf_counter() - start
|
||||
|
||||
summary_path = write_outputs(
|
||||
args=args,
|
||||
mage_repo=mage_repo,
|
||||
media_path=media_path,
|
||||
cmd=cmd,
|
||||
result=result,
|
||||
elapsed_s=elapsed_s,
|
||||
)
|
||||
print(f"summary saved to {summary_path}")
|
||||
if result is not None:
|
||||
if result.stdout:
|
||||
print(result.stdout, end="")
|
||||
if result.returncode != 0:
|
||||
if result.stderr:
|
||||
print(result.stderr, file=sys.stderr, end="")
|
||||
raise SystemExit(result.returncode)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,79 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import websockets
|
||||
|
||||
|
||||
async def run(args: argparse.Namespace) -> None:
|
||||
headers = {"Authorization": f"Bearer {args.auth_token}"} if args.auth_token else None
|
||||
started = time.perf_counter()
|
||||
async with websockets.connect(
|
||||
f"{args.url}?session_id={args.session_id}",
|
||||
additional_headers=headers,
|
||||
max_size=args.max_message_bytes,
|
||||
open_timeout=args.timeout,
|
||||
) as websocket:
|
||||
created = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout))
|
||||
print(json.dumps(created, ensure_ascii=False))
|
||||
await websocket.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.append",
|
||||
"modality": "text",
|
||||
"data": args.prompt,
|
||||
}
|
||||
)
|
||||
)
|
||||
encoded = base64.b64encode(Path(args.video).read_bytes()).decode("ascii")
|
||||
await websocket.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.append",
|
||||
"modality": "video",
|
||||
"data": {
|
||||
"kind": "video/mp4",
|
||||
"segment_id": args.segment_id,
|
||||
"pts_ms": args.pts_ms,
|
||||
"duration_ms": args.duration_ms,
|
||||
"data": {"video_base64": encoded},
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
while True:
|
||||
event = json.loads(await asyncio.wait_for(websocket.recv(), args.timeout))
|
||||
print(json.dumps(event, ensure_ascii=False))
|
||||
if event.get("type") in {"response.done", "error"}:
|
||||
break
|
||||
await websocket.send(json.dumps({"type": "close"}))
|
||||
print(f"elapsed_s={time.perf_counter() - started:.3f}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Send one Mage-VL video window over the duplex WebSocket.")
|
||||
parser.add_argument("--url", default="ws://127.0.0.1:8090/v1/mage-vl/duplex")
|
||||
parser.add_argument("--video", required=True)
|
||||
parser.add_argument("--prompt", default="Describe this video segment in detail.")
|
||||
parser.add_argument("--session-id", default="mage-vl-example")
|
||||
parser.add_argument("--segment-id", default="segment-0")
|
||||
parser.add_argument("--pts-ms", type=int, default=0)
|
||||
parser.add_argument("--duration-ms", type=int, default=4000)
|
||||
parser.add_argument("--auth-token")
|
||||
parser.add_argument("--timeout", type=float, default=120.0)
|
||||
parser.add_argument("--max-message-bytes", type=int, default=16 * 1024 * 1024)
|
||||
args = parser.parse_args()
|
||||
asyncio.run(run(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -49,7 +49,7 @@ scenario below for live barge-in validation on the target GPU.
|
||||
Default:
|
||||
|
||||
```bash
|
||||
vllm-omni serve openbmb/MiniCPM-o-4_5 \
|
||||
vllm serve openbmb/MiniCPM-o-4_5 \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \
|
||||
--trust-remote-code \
|
||||
|
||||
@@ -37,7 +37,7 @@ During generation the client normally sends only `session.start` and then receiv
|
||||
## Install Client Dependency
|
||||
|
||||
```bash
|
||||
pip install websockets
|
||||
pip install av websockets
|
||||
# For the Gradio demo:
|
||||
pip install 'vllm-omni[demo]' websockets
|
||||
```
|
||||
|
||||
@@ -8,7 +8,7 @@ the received bytes to disk when the session finishes or is interrupted.
|
||||
Image/reference input is intentionally not supported in this example yet.
|
||||
|
||||
Requirements:
|
||||
pip install websockets
|
||||
pip install av websockets
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
@@ -13,7 +13,7 @@ For the full list of supported architectures across all modalities, see
|
||||
## Supported Models
|
||||
|
||||
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
|
||||
|---|---|---|---|---|---|
|
||||
| --- | --- | --- | --- | --- | --- |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
|
||||
@@ -104,6 +104,7 @@ For full request-shape documentation (all parameters, response formats, error co
|
||||
2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve zai-org/GLM-TTS --omni --trust-remote-code --port 8091
|
||||
# or:
|
||||
@@ -111,6 +112,7 @@ bash examples/online_serving/text_to_speech/glm_tts/run_server.sh /path/to/GLM-T
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# Voice cloning (required)
|
||||
python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
@@ -127,11 +129,13 @@ python examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py \
|
||||
```
|
||||
|
||||
### Gradio demo
|
||||
|
||||
```bash
|
||||
bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 24 kHz mono WAV via HiFT vocoder.
|
||||
- `ref_audio` + `ref_text` are **required** together on every request. Reference audio should be 3-10 seconds.
|
||||
- Voice cloning feature extraction (WhisperVQ, CampPlus, mel) runs on the model side — no external dependency on the serving layer.
|
||||
@@ -151,6 +155,7 @@ pip install 'vllm-omni[indextts2]'
|
||||
```
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092
|
||||
# or, to pass the bundled deploy config explicitly:
|
||||
@@ -164,6 +169,7 @@ bash examples/online_serving/text_to_speech/indextts2/run_server.sh
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# Voice cloning (ref_audio required)
|
||||
python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
@@ -194,6 +200,7 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 22.05 kHz mono WAV.
|
||||
- Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; neither version provides a built-in text-only preset voice.
|
||||
- IndexTTS-2.5 request controls `lang` and `text_normalization` are carried in `extra_params`; the client does this for `--model-version 2.5`.
|
||||
@@ -212,6 +219,7 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
@@ -237,14 +245,17 @@ export VLLM_OMNI_FISH_KVCACHE_ATTN=0
|
||||
```
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
# or:
|
||||
./fish_speech/run_server.sh
|
||||
```
|
||||
|
||||
The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`).
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
@@ -257,6 +268,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
```
|
||||
|
||||
### CLI client
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/fish_speech
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
@@ -264,12 +276,14 @@ python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
```
|
||||
|
||||
### Gradio demo
|
||||
|
||||
```bash
|
||||
./fish_speech/run_gradio_demo.sh # launches server + Gradio
|
||||
python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server already running
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 44.1 kHz mono.
|
||||
- Streaming PCM player command must use `-r 44100`.
|
||||
|
||||
@@ -280,10 +294,13 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
|
||||
Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
bash examples/online_serving/text_to_speech/ming_tts/run_server.sh
|
||||
```
|
||||
|
||||
Equivalent manual command:
|
||||
|
||||
```bash
|
||||
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
@@ -292,12 +309,14 @@ vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "你好,这是 Ming 在线语音合成测试。"
|
||||
```
|
||||
|
||||
Structured dialect control:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "我觉得社会企业同个人都有责任" \
|
||||
@@ -306,6 +325,7 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
```
|
||||
|
||||
Zero-shot cloning:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
|
||||
@@ -314,6 +334,7 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook.
|
||||
- Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples.
|
||||
- Full request details live in [`ming_tts/README.md`](ming_tts/README.md).
|
||||
@@ -325,11 +346,14 @@ python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
# from repo root
|
||||
bash examples/online_serving/text_to_speech/ming_flash_omni_tts/run_server.sh
|
||||
```
|
||||
|
||||
Equivalent manual command:
|
||||
|
||||
```bash
|
||||
vllm serve Jonathan1909/Ming-flash-omni-2.0 \
|
||||
--deploy-config vllm_omni/deploy/ming_flash_omni_tts.yaml \
|
||||
@@ -338,6 +362,7 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 \
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.py \
|
||||
--text "我们当迎着阳光辛勤耕作,去摘取,去制作,去品尝,去馈赠。" \
|
||||
@@ -345,6 +370,7 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
|
||||
```
|
||||
|
||||
ASMR-style caption via `instructions`:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.py \
|
||||
--text "我会一直在这里陪着你,直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗?" \
|
||||
@@ -353,6 +379,7 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Server uses `use_zero_spk_emb=True` and the cookbook decode defaults (`max_decode_steps=200`, `cfg=2.0`, `sigma=0.25`, `temperature=0.0`). For other caption fields (`语速`, `基频`, `IP`, BGM, etc.) or overriding decode args, use the offline example where `additional_information` is set explicitly.
|
||||
- This is the online counterpart of [`examples/offline_inference/text_to_speech/ming_flash_omni_tts/`](../../offline_inference/text_to_speech/ming_flash_omni_tts/).
|
||||
- For multimodal Ming-flash-omni online serving, see [`examples/online_serving/ming_flash_omni/`](../../ming_flash_omni/).
|
||||
@@ -366,14 +393,17 @@ Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every
|
||||
> The OpenAI-schema `voice` and `ref_text` fields are accepted but ignored — `voice_clone` does not consume a transcript, and upstream's `continuation` mode (the only path that accepts `prompt_text`) emits near-silent output, so it is not exposed here. Sample reference clips ship in the upstream repo under [`assets/audio/`](https://github.com/OpenMOSS/MOSS-TTS-Nano/tree/main/assets/audio).
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
|
||||
# or:
|
||||
./moss_tts_nano/run_server.sh
|
||||
```
|
||||
|
||||
The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--deploy-config`, `--trust-remote-code`, or `--enforce-eager` flags are needed.
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# One-off fetch of a sample reference clip; cache under XDG_CACHE_HOME.
|
||||
REF_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/moss-tts-nano"
|
||||
@@ -392,6 +422,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
```
|
||||
|
||||
### Streaming PCM
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
@@ -405,6 +436,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
```
|
||||
|
||||
### Gradio demo
|
||||
|
||||
```bash
|
||||
# Option 1: launch server + Gradio together
|
||||
./moss_tts_nano/run_gradio_demo.sh
|
||||
@@ -412,9 +444,11 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
# Option 2: server already running
|
||||
python moss_tts_nano/gradio_demo.py --api-base http://localhost:8091
|
||||
```
|
||||
Then open http://localhost:7860 in your browser.
|
||||
|
||||
Then open <http://localhost:7860> in your browser.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output is 48 kHz mono PCM (the upstream tokenizer is internally stereo at 48 kHz; the wrapper averages to mono before reaching the engine).
|
||||
- Standard `/v1/audio/speech` request shape: `input`, `ref_audio` (base64 data URL), `response_format`, `stream`, `max_new_tokens`. The `voice` and `ref_text` fields from the OpenAI schema are accepted but ignored.
|
||||
|
||||
@@ -425,12 +459,15 @@ Then open http://localhost:7860 in your browser.
|
||||
Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
huggingface-cli download k2-fsa/OmniVoice
|
||||
```
|
||||
|
||||
Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
|
||||
# or:
|
||||
@@ -438,6 +475,7 @@ vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
|
||||
```
|
||||
|
||||
### CLI client
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/omnivoice
|
||||
# Text-only (auto voice)
|
||||
@@ -463,13 +501,12 @@ python speech_client.py --text "Hello, how are you?" --seed 42
|
||||
|
||||
The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--voice`, `--ref-audio`, `--ref-text`, `--instructions`, `--seed`, and `--output`.
|
||||
|
||||
|
||||
## Qwen3-TTS
|
||||
|
||||
Three model variants exposed via separate checkpoints:
|
||||
|
||||
| Variant | HF repo | Use |
|
||||
|---|---|---|
|
||||
| --- | --- | --- |
|
||||
| CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions |
|
||||
| VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description |
|
||||
| Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio |
|
||||
@@ -477,6 +514,7 @@ Three model variants exposed via separate checkpoints:
|
||||
Each variant ships smaller `0.6B` companions where available.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
|
||||
# or:
|
||||
@@ -486,11 +524,13 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
|
||||
```
|
||||
|
||||
### Executor backend
|
||||
|
||||
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
|
||||
|
||||
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# CustomVoice with a predefined speaker
|
||||
python qwen3_tts/openai_speech_client.py \
|
||||
@@ -516,7 +556,9 @@ python qwen3_tts/openai_speech_client.py \
|
||||
```
|
||||
|
||||
### Voices endpoint
|
||||
|
||||
List available voices, or upload a custom one for Base cloning:
|
||||
|
||||
```bash
|
||||
# List
|
||||
curl http://localhost:8091/v1/audio/voices
|
||||
@@ -529,10 +571,18 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "ref_text=The exact transcript of the audio sample." \
|
||||
-F "speaker_description=warm narrator"
|
||||
```
|
||||
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
|
||||
|
||||
For Qwen3-TTS, uploaded voices are Base voice-cloning inputs and require a Base
|
||||
checkpoint. When a request names an uploaded voice, the server infers
|
||||
`task_type="Base"`. Built-in presets such as `vivian` and `ryan` remain
|
||||
CustomVoice speakers and require a CustomVoice checkpoint.
|
||||
|
||||
### Precomputed custom voices
|
||||
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
|
||||
|
||||
For reused Base voice-cloning speakers, precompute the reference artifacts once
|
||||
and load them at server startup. Precomputed voices use the same Base task and
|
||||
checkpoint-matching rules as uploaded voices:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/precompute_custom_voice.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
@@ -542,11 +592,15 @@ python qwen3_tts/precompute_custom_voice.py \
|
||||
--mode icl \
|
||||
--output-dir /path/to/custom_voices
|
||||
```
|
||||
|
||||
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
|
||||
|
||||
```yaml
|
||||
custom_voice_dir: /path/to/custom_voices
|
||||
```
|
||||
|
||||
Then start the server with that config and call the Speech API with only the voice name:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
|
||||
|
||||
@@ -557,6 +611,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
```
|
||||
|
||||
### Streaming PCM
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
@@ -569,16 +624,20 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` in the deploy config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
|
||||
|
||||
### Streaming WebSocket
|
||||
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "First utterance." \
|
||||
@@ -586,6 +645,7 @@ python qwen3_tts/streaming_speech_client.py \
|
||||
```
|
||||
|
||||
To receive word-level timestamps, launch the server with a forced aligner:
|
||||
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni \
|
||||
@@ -593,7 +653,9 @@ vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--trust-remote-code \
|
||||
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
|
||||
```
|
||||
|
||||
Then request PCM JSON sidecar chunks:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "Hello world. How are you?" \
|
||||
@@ -601,6 +663,7 @@ python qwen3_tts/streaming_speech_client.py \
|
||||
--response-format pcm \
|
||||
--word-timestamps
|
||||
```
|
||||
|
||||
The client writes one PCM file per sentence and a matching
|
||||
`sentence_XXX_timestamps.json` sidecar.
|
||||
|
||||
@@ -613,15 +676,18 @@ still returns — use the WebSocket path for long transcripts.
|
||||
|
||||
To *see* the alignment instead of reading a JSON sidecar, run the
|
||||
word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
|
||||
|
||||
```bash
|
||||
python qwen3_tts/word_timestamps_demo.py --api-base http://localhost:8091
|
||||
```
|
||||
|
||||
Each sentence's audio plays in an `<audio>` element while its text is rendered
|
||||
as inline word spans; the current word highlights as `audio.currentTime`
|
||||
crosses each `start_ms`. The **Stop (barge-in)** button cuts playback and
|
||||
reports the last-spoken word, useful for the voice-agent barge-in case.
|
||||
|
||||
### Gradio demos
|
||||
|
||||
```bash
|
||||
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
|
||||
./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign
|
||||
@@ -629,12 +695,15 @@ reports the last-spoken word, useful for the voice-agent barge-in case.
|
||||
```
|
||||
|
||||
### Speaker embedding interpolation
|
||||
|
||||
`qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage.
|
||||
|
||||
### Batch client
|
||||
|
||||
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
|
||||
|
||||
### Notes
|
||||
|
||||
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
|
||||
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
|
||||
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
||||
@@ -646,12 +715,15 @@ reports the last-spoken word, useful for the voice-agent barge-in case.
|
||||
Single-stage native AR TTS at 48 kHz.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000
|
||||
```
|
||||
|
||||
Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks.
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
# Zero-shot synthesis
|
||||
python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2."
|
||||
@@ -661,10 +733,13 @@ python voxcpm2/openai_speech_client.py \
|
||||
--text "This should sound like the reference speaker." \
|
||||
--ref-audio /path/to/reference.wav
|
||||
```
|
||||
|
||||
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
|
||||
|
||||
### Precomputed custom voices
|
||||
|
||||
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
|
||||
|
||||
```bash
|
||||
python voxcpm2/precompute_custom_voice.py \
|
||||
--model openbmb/VoxCPM2 \
|
||||
@@ -674,16 +749,21 @@ python voxcpm2/precompute_custom_voice.py \
|
||||
--prompt-text "Original transcript of the reference audio" \
|
||||
--output-dir /path/to/custom_voices
|
||||
```
|
||||
|
||||
Add the output directory to the deploy config:
|
||||
|
||||
```yaml
|
||||
custom_voice_dir: /path/to/custom_voices
|
||||
```
|
||||
|
||||
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
|
||||
|
||||
### Gradio demo (gapless streaming via AudioWorklet)
|
||||
|
||||
```bash
|
||||
python voxcpm2/gradio_demo.py
|
||||
```
|
||||
|
||||
Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`.
|
||||
|
||||
---
|
||||
@@ -693,24 +773,31 @@ Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free p
|
||||
Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Latest `mistral_common` with `SpeechRequest` support:
|
||||
|
||||
```bash
|
||||
pip install -e /path/to/mistral-common # or upgrade from PyPI when available
|
||||
```
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
|
||||
```
|
||||
|
||||
Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`.
|
||||
|
||||
### Gradio demo
|
||||
|
||||
```bash
|
||||
python voxtral_tts/gradio_demo.py
|
||||
```
|
||||
|
||||
The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients).
|
||||
|
||||
### Notes
|
||||
|
||||
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
|
||||
- Voice cloning is gated upstream and may require a recent `mistral_common`.
|
||||
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
|
||||
@@ -1,3 +1,6 @@
|
||||
# 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.
|
||||
|
||||
This script demonstrates how to use the OpenAI-compatible speech API
|
||||
@@ -74,6 +77,8 @@ def run_tts_generation(args) -> None:
|
||||
"voice": args.speaker,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
if args.sample_rate is not None:
|
||||
payload["sample_rate"] = args.sample_rate
|
||||
|
||||
# Add optional parameters
|
||||
if args.instructions:
|
||||
@@ -246,6 +251,13 @@ def parse_args():
|
||||
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
|
||||
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(
|
||||
"--output",
|
||||
"-o",
|
||||
|
||||
+173
-20
@@ -1,11 +1,14 @@
|
||||
# Boogu-Image
|
||||
|
||||
> Text-to-image and image-editing online serving (Boogu-Image-0.1-Base / -Edit)
|
||||
> Text-to-image and image-editing online serving
|
||||
(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Boogu
|
||||
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image), `Boogu/Boogu-Image-0.1-Edit` (image editing)
|
||||
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image),
|
||||
`Boogu/Boogu-Image-0.1-Edit` (image editing), and
|
||||
`Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing)
|
||||
- Task: Text-to-image generation and text-guided image editing (TI2I)
|
||||
- Mode: Online serving with the OpenAI-compatible API
|
||||
- Maintainer: Community
|
||||
@@ -13,28 +16,33 @@
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a known-good starting point for serving
|
||||
`Boogu/Boogu-Image-0.1-Base` with vLLM-Omni's native pipeline (no
|
||||
`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or
|
||||
`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no
|
||||
`--diffusion-load-format diffusers`, and the upstream `boogu` package is not
|
||||
required).
|
||||
|
||||
Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model
|
||||
family. This recipe covers the Base text-to-image checkpoint, which pairs a
|
||||
Qwen3-VL multimodal encoder with a Diffusion Transformer (DiT) and a flow-match
|
||||
Euler scheduler with time-shift. It handles photorealistic generation and
|
||||
Chinese/English text rendering.
|
||||
family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder
|
||||
with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with
|
||||
time-shift. It handles photorealistic generation and Chinese/English text
|
||||
rendering. The Edit and Edit-Turbo checkpoints use the same native
|
||||
`BooguImagePipeline` for image editing.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Base>
|
||||
- Edit model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit>
|
||||
- Edit-Turbo 1K hotfix: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit-Turbo/tree/hotfix-1k-20260708>
|
||||
- Project page: <https://boogu.org>
|
||||
- GitHub: <https://github.com/boogu-project/Boogu-Image>
|
||||
- Related example: [`examples/online_serving/text_to_image/`](../../examples/online_serving/text_to_image/README.md)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents tested configurations for CUDA GPU serving. The native
|
||||
pipeline runs single-GPU; multi-GPU parallelism, CPU offload, and cache
|
||||
acceleration are not yet supported for this model (see Notes).
|
||||
This recipe documents supported configurations for CUDA GPU serving. The native
|
||||
pipeline supports single-GPU inference and classifier-free-guidance (CFG)
|
||||
parallelism. CPU offload, cache acceleration, and other multi-GPU dimensions
|
||||
remain unsupported (see Notes).
|
||||
|
||||
## GPU
|
||||
|
||||
@@ -93,7 +101,7 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
- **Memory usage:** ~34.6 GiB on GPU. Use `--vae-use-slicing --vae-use-tiling`
|
||||
to trim peak VRAM if needed.
|
||||
- **Key flags:**
|
||||
- `--omni` — enables vLLM-Omni diffusion serving.
|
||||
- `--omni` — enables vLLM-Omni diffusion serving.
|
||||
- **Guidance:** Boogu-Image uses `guidance_scale` (mapped to the upstream
|
||||
`text_guidance_scale`); the default is `4.0`. Classifier-free guidance is
|
||||
active whenever `guidance_scale > 1.0`.
|
||||
@@ -101,8 +109,24 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
The model's maximum native resolution is 2K.
|
||||
- **Known limitations (not yet supported):** CPU offload
|
||||
(`--enable-cpu-offload` / `--enable-layerwise-offload`), Cache-DiT
|
||||
(`--cache-backend cache_dit`), and multi-GPU parallelism (TP / SP / CFG /
|
||||
HSDP) are planned follow-ups and are not validated for this model yet.
|
||||
(`--cache-backend cache_dit`), and TP / SP / HSDP multi-GPU parallelism.
|
||||
|
||||
### 2 x H100 (CFG parallel, Base T2I)
|
||||
|
||||
Base T2I has two predictions per guided denoising step, so two CFG ranks map
|
||||
one-to-one to the positive and negative-text branches:
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Base \
|
||||
--omni \
|
||||
--cfg-parallel-size 2 \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Use `guidance_scale > 1.0` to activate the parallel two-branch path. Requests
|
||||
with `guidance_scale=1.0` remain valid on the same server: every rank evaluates
|
||||
only the positive branch, no negative embeddings are built, and no guidance is
|
||||
applied.
|
||||
|
||||
## Image editing (Boogu-Image-0.1-Edit)
|
||||
|
||||
@@ -111,13 +135,40 @@ the image-editing (TI2I) path activates automatically when a request carries a
|
||||
reference image. The Base text-to-image path is unaffected (no reference image
|
||||
is sent).
|
||||
|
||||
#### Command
|
||||
### Command
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit --omni --port 8091
|
||||
```
|
||||
|
||||
#### Verification
|
||||
For text-only or image-only guidance, use two CFG ranks:
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit \
|
||||
--omni \
|
||||
--cfg-parallel-size 2 \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Double guidance has three predictions per denoising step. It supports either
|
||||
two ranks (round-robin assignment: rank 0 runs branches 0 and 2) or three ranks
|
||||
(one branch per rank):
|
||||
|
||||
```bash
|
||||
# Practical two-GPU configuration
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit \
|
||||
--omni \
|
||||
--cfg-parallel-size 2 \
|
||||
--port 8091
|
||||
|
||||
# Maximum branch parallelism
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit \
|
||||
--omni \
|
||||
--cfg-parallel-size 3 \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
### Verification
|
||||
|
||||
Edit an image with `/v1/images/edits` (the model-card example — change a photo
|
||||
to a colored-pencil drawing). Diffusion parameters are plain multipart form
|
||||
@@ -157,24 +208,126 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited.png
|
||||
```
|
||||
|
||||
#### Notes
|
||||
### Notes
|
||||
|
||||
- **Single reference image:** only one input image is supported for now (the
|
||||
upstream "Only support 1 reference image for now" limit).
|
||||
- **Guidance semantics:**
|
||||
- `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
|
||||
- `guidance_scale` = text guidance (upstream `text_guidance_scale`, default
|
||||
`4.0`); `> 1.0` enables text CFG. Editing typically uses `5.0`.
|
||||
- `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
|
||||
- `guidance_scale_2` = image guidance (upstream `image_guidance_scale`,
|
||||
default `1.0` = off). Setting it `> 1.0` enables the double-guidance path
|
||||
(3 model predictions per step), steering more strongly toward the reference
|
||||
image.
|
||||
- **CFG-size recommendations:**
|
||||
|
||||
| Request mode | Scales | Predictions per step | Recommended CFG size |
|
||||
| --- | --- | ---: | ---: |
|
||||
| Base T2I text guidance | `guidance_scale > 1`, no reference | 2 | 2 |
|
||||
| Edit text-only guidance | `guidance_scale > 1`, `guidance_scale_2 = 1` | 2 | 2 |
|
||||
| Edit image-only guidance | `guidance_scale = 1`, `guidance_scale_2 > 1` | 2 | 2 |
|
||||
| Edit double guidance | both scales `> 1` | 3 | 2 (practical) or 3 (maximum parallelism) |
|
||||
| CFG off | both scales `= 1` | 1 | Reuse a size-2 server safely; no speedup expected |
|
||||
|
||||
Double guidance preserves Boogu's original combination:
|
||||
|
||||
```text
|
||||
pred = positive_with_reference
|
||||
+ (text_scale - 1) * (positive_with_reference - negative_with_reference)
|
||||
+ (image_scale - 1) * (negative_with_reference - uncond)
|
||||
```
|
||||
- **Output resolution:** the output size follows the reference image (upstream
|
||||
`align_res`, on by default for a single reference), so `height`/`width` are
|
||||
derived from the input and requested sizes are not applied for edits.
|
||||
- **Same limitations** as the Base checkpoint apply (no CPU offload, Cache-DiT,
|
||||
or multi-GPU parallelism yet).
|
||||
- **Same limitations** as the Base checkpoint apply except that CFG parallelism
|
||||
is supported as described above; CPU offload, Cache-DiT, TP, SP, and HSDP are
|
||||
still unsupported.
|
||||
- **Offline editing:** the shared offline example
|
||||
[`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_to_image.md)
|
||||
supports Boogu-Image-Edit directly
|
||||
(`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional
|
||||
`--guidance-scale-2`).
|
||||
|
||||
## Fast image editing (Boogu-Image-0.1-Edit-Turbo)
|
||||
|
||||
Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json`
|
||||
declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit.
|
||||
The upstream 1K hotfix is recommended for stable results and is pinned below
|
||||
to avoid silently loading the older checkpoint from the repository's default
|
||||
revision.
|
||||
|
||||
### Command
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \
|
||||
--omni \
|
||||
--revision hotfix-1k-20260708 \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
### Verification
|
||||
|
||||
Edit an image with the distilled four-step settings:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/images/edits \
|
||||
-F model="Boogu/Boogu-Image-0.1-Edit-Turbo" \
|
||||
-F image="@input.png" \
|
||||
-F prompt="Change the style to a colored pencil drawing." \
|
||||
-F num_inference_steps=4 \
|
||||
-F guidance_scale=1.0 \
|
||||
-F seed=42 \
|
||||
| jq -r '.data[0].b64_json' | base64 -d > edited-turbo.png
|
||||
```
|
||||
|
||||
Or via chat completions:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
|
||||
{"type": "text", "text": "Change the style to a colored pencil drawing."}
|
||||
]}
|
||||
],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 4,
|
||||
"guidance_scale": 1.0,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > edited-turbo.png
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K
|
||||
hotfix over the 1.5K variant for more stable results.
|
||||
- **Recommended settings:** `num_inference_steps=4` and
|
||||
`guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance
|
||||
behavior is already baked into the distilled checkpoint. A scale of `1.0`
|
||||
avoids applying additional classifier-free guidance; unlike the regular Edit
|
||||
checkpoint, it should not be increased to `5.0`.
|
||||
- **Reference images:** the same single-reference and `align_res` behavior as
|
||||
the regular Edit checkpoint applies.
|
||||
- **Known limitations:** the same single-GPU limitations as Base and Edit apply;
|
||||
CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
|
||||
|
||||
## Performance validation
|
||||
|
||||
The checked-in single-device A40 measurements remain the current Boogu
|
||||
baseline: Base T2I is about 16.37 s request latency at concurrency 1 and Edit
|
||||
text-only guidance is about 29.64 s for 512x512, 28-step requests, with roughly
|
||||
36.8 GB peak memory. CFG-parallel A/B numbers must be collected on the same
|
||||
hardware and software revision before reporting a speedup; do not compare a
|
||||
multi-GPU result against these A40 numbers as if it were a controlled A/B.
|
||||
|
||||
Run the dedicated perf rows in
|
||||
`tests/dfx/perf/tests/test_boogu_image_vllm_omni.json` and
|
||||
`tests/dfx/perf/tests/test_boogu_image_edit_vllm_omni.json` on 2x/3x H100, then
|
||||
record latency, throughput, and per-GPU peak memory beside the single-device
|
||||
run. The expected shape is close to the general CFG-parallel range (roughly
|
||||
1.5x-1.8x for balanced two-branch workloads); the three-branch Edit result will
|
||||
depend on whether two-rank round-robin or full three-rank execution is used and
|
||||
must be reported from measurement rather than assumed.
|
||||
@@ -0,0 +1,114 @@
|
||||
# Mage-VL
|
||||
|
||||
> Codec-native image/video understanding plus proactive full-duplex streaming.
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Microsoft
|
||||
- Model: [`microsoft/Mage-VL`](https://huggingface.co/microsoft/Mage-VL)
|
||||
- Task: Image understanding, video understanding, codec-native video understanding,
|
||||
and event-gated streaming commentary
|
||||
- Mode: Full-duplex WebSocket server and session adapter in
|
||||
`vllm_omni/experimental/fullduplex/mage_vl`
|
||||
- Maintainer: Community
|
||||
|
||||
## Status
|
||||
|
||||
This experimental integration provides the Mage-VL adapter contract: bounded causal
|
||||
visual windows, codec-window input, rolling gate evaluation, explicit query handling,
|
||||
per-session state isolation, and stale-output cancellation through the shared
|
||||
full-duplex runtime.
|
||||
|
||||
The production transport exposes `WebSocket /v1/mage-vl/duplex`, with bearer-token
|
||||
authentication, bounded concurrent sessions and message sizes, idle timeouts,
|
||||
disconnect cancellation, per-session state cleanup, and serialized access to the
|
||||
shared Transformers checkpoint. Native vLLM model execution remains follow-up work.
|
||||
|
||||
The default `frames` backend uses decoded frames and has no codec preprocessing
|
||||
requirement. The optional `codec` backend requires `ffmpeg` and `ffprobe`, plus the
|
||||
`cv-preinfer` executable provided by `codec-video-prep>=0.2.5`, on `PATH`. Server
|
||||
startup fails with an actionable error when `cv-preinfer` is unavailable; use
|
||||
`--video-backend frames` when codec preprocessing is not installed.
|
||||
|
||||
## Serving
|
||||
|
||||
```bash
|
||||
python -m vllm_omni.experimental.fullduplex.mage_vl.serving.server \
|
||||
--model microsoft/Mage-VL --host 0.0.0.0 --port 8090 \
|
||||
--auth-token "$MAGE_VL_API_KEY" --max-sessions 32 \
|
||||
--video-backend frames --num-frames 8 --target-fps 1 \
|
||||
--gate-threshold 0.5 --attn-impl sdpa
|
||||
|
||||
python examples/online_serving/mage_vl/duplex_client.py \
|
||||
--url ws://127.0.0.1:8090/v1/mage-vl/duplex \
|
||||
--auth-token "$MAGE_VL_API_KEY" \
|
||||
--video /path/to/segment.mp4 \
|
||||
--prompt "Describe this video segment in detail."
|
||||
```
|
||||
|
||||
Each input window is a JSON `input.append` event whose `data.data.video_base64`
|
||||
contains an MP4 segment. A text `input.append` before the video requests an explicit
|
||||
answer; without text, the StreamMind gate decides whether to respond proactively.
|
||||
The server emits `session.created`, `response.created`, one or more
|
||||
`response.delta`, and `response.done` events. Clients can send `response.cancel` for
|
||||
barge-in and `close` for graceful shutdown.
|
||||
|
||||
The adapter can still be wired to another backend through `--adapter module:factory`,
|
||||
including:
|
||||
|
||||
- a Transformers `trust_remote_code=True` backend for offline/local testing;
|
||||
- the model author's SGLang-compatible serving path;
|
||||
- a future native vLLM `mage_vl` backend.
|
||||
|
||||
## Adapter Shape
|
||||
|
||||
```python
|
||||
from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig
|
||||
from vllm_omni.experimental.fullduplex.mage_vl import MageVLDuplexAdapter, MageVLDuplexRuntime
|
||||
|
||||
|
||||
adapter = MageVLDuplexAdapter(gate=run_mage_gate, generate=run_mage_decoder)
|
||||
session = DuplexSession(
|
||||
"stream-1",
|
||||
DuplexSessionConfig(
|
||||
input_modalities=("codec_window", "video", "image", "text"),
|
||||
output_modalities=("text",),
|
||||
proactive=True,
|
||||
),
|
||||
)
|
||||
runtime = MageVLDuplexRuntime(session, adapter)
|
||||
await runtime.run(input_events, emit_event)
|
||||
```
|
||||
|
||||
Input events use the shared full-duplex protocol. For decoded video/image, pass the frame
|
||||
payload directly. For codec-native streaming, pass `modality="codec_window"` and include
|
||||
fields such as `segment_id`, `pts_ms`, `duration_ms`, `codec`, and `metadata`.
|
||||
|
||||
```python
|
||||
{
|
||||
"type": "input.append",
|
||||
"modality": "codec_window",
|
||||
"data": {
|
||||
"kind": "h264",
|
||||
"segment_id": "seg-42",
|
||||
"pts_ms": 42000,
|
||||
"duration_ms": 1000,
|
||||
"codec": {
|
||||
"motion_vectors": ...,
|
||||
"residual_energy": ...,
|
||||
},
|
||||
"metadata": {"gop": "P"},
|
||||
},
|
||||
}
|
||||
```
|
||||
|
||||
The gate callable receives the current rolling window and returns either a
|
||||
`MageVLGateDecision`, a mapping with `should_respond` / `text` / `event_id`, a bool, or
|
||||
a text string. Repeated `event_id` values are de-duplicated at the session adapter.
|
||||
|
||||
## Verification
|
||||
|
||||
```bash
|
||||
pytest tests/e2e/features/fullduplex/test_mage_vl_adapter.py \
|
||||
tests/examples/offline_inference/test_mage_vl.py
|
||||
```
|
||||
@@ -278,7 +278,7 @@ steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
|
||||
|
||||
### Attention Backends
|
||||
|
||||
On datacenter Blackwell GPUs, MiniMax H3 defaults to dense BF16
|
||||
On supported datacenter Blackwell systems, MiniMax H3 defaults to dense BF16
|
||||
`TRTLLM_ATTN`; no attention backend flag is required. To select it explicitly,
|
||||
use:
|
||||
|
||||
@@ -298,17 +298,33 @@ FA4 remains available by explicitly selecting the `FLASH_ATTN` backend:
|
||||
--diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
|
||||
`Using CuTe FlashAttention-4 on Blackwell` before recording FA4 measurements.
|
||||
With the optional `[fa4]` dependency installed, `FLASH_ATTN` prefers FA4 on
|
||||
Blackwell. Confirm the server log contains `Using CuTe FlashAttention-4 on
|
||||
Blackwell` before recording FA4 measurements.
|
||||
|
||||
`TRTLLM_ATTN` additionally supports two **lossy** optimizations for the long main
|
||||
DiT attention sequence: SAGE attention quantization and Skip-Softmax Sparse
|
||||
Attention. SAGE quantizes Q/K to the configured dtype and V to FP8. This example uses
|
||||
`fp8_e4m3` for Q/K; B200 also supports `int8` Q/K. The TRTLLM SAGE path fixes V
|
||||
to FP8, so vLLM-Omni only exposes the Q/K dtype. The token refiner is a short
|
||||
attention path, so the `per_role` override leaves SAGE and Skip-Softmax disabled
|
||||
for it. The example enables the calibration-free Skip-Softmax path with
|
||||
`threshold=0.05`, after the normalized timestep reaches `0.97`:
|
||||
`TRTLLM_ATTN` additionally offers two **lossy** optimizations for the long main
|
||||
DiT attention sequence: SAGE attention quantization and Skip-Softmax sparse
|
||||
attention. Both work under the pure Ulysses parallelism of the profile above
|
||||
(`--usp 4 --ring 1`). The example below enables both:
|
||||
|
||||
- SAGE with `fp8_e4m3` Q/K; P and V are always FP8 in this kernel. B200
|
||||
additionally supports `int8` Q/K, which preserves accuracy better than FP8.
|
||||
- Skip-Softmax with a direct `threshold=0.05` (the calibrated
|
||||
`target_sparsity` control needs ModelOpt metadata that the official H3
|
||||
checkpoint does not include). Together with the cutoff below this is a
|
||||
**conservative** setting: a low threshold skips only clearly negligible tiles,
|
||||
and a cutoff close to `1.0` leaves a substantial dense prefix. Raise
|
||||
`threshold` or lower `disabled_until_timestep` for more speedup once quality
|
||||
is verified.
|
||||
- `disabled_until_timestep=0.97` keeps the early high-noise steps dense. The
|
||||
gate compares against the video sigma, which H3's default flow shift of 12
|
||||
keeps high for much of the run: at 50 steps, `0.99`, `0.97`, and `0.95`
|
||||
leave the first 6, 14, and 19 of 49 denoiser forwards dense. See the
|
||||
[Skip-Softmax design](https://github.com/vllm-project/vllm-omni/blob/main/docs/design/feature/skip_softmax.md#timestep-gating)
|
||||
for how the cutoff maps to steps.
|
||||
- A `per_role` entry that keeps the token refiner, a short attention path,
|
||||
dense. A per-role spec does not inherit `quant` or `skip_softmax` from
|
||||
`default`.
|
||||
|
||||
```bash
|
||||
--diffusion-attention-config '{
|
||||
@@ -332,10 +348,13 @@ for it. The example enables the calibration-free Skip-Softmax path with
|
||||
}'
|
||||
```
|
||||
|
||||
For configuration details, see
|
||||
[TRTLLM_ATTN Backend and Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-backend-and-skip-softmax)
|
||||
Both optimizations trade fidelity for speed and their effects compound.
|
||||
Compare against dense output on the same prompt and seed before adopting them.
|
||||
For the full key reference, see
|
||||
[Skip-Softmax](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends/trtllm.md#skip-softmax)
|
||||
and
|
||||
[TRTLLM_ATTN SAGE Quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends.md#trtllm_attn-sage-quantization).
|
||||
[SAGE quantization](https://github.com/vllm-project/vllm-omni/blob/main/docs/user_guide/diffusion/attention_backends/trtllm.md#sage-quantization)
|
||||
in the TRTLLM attention guide.
|
||||
|
||||
### Text encoder tensor parallelism
|
||||
|
||||
@@ -962,13 +981,27 @@ mostly base H3 weights on a four-step schedule. Offload is refused for the same
|
||||
reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and
|
||||
`--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast.
|
||||
|
||||
> [!NOTE]
|
||||
> Only the **Dense / Data-Free** variant is supported today. The three VSA
|
||||
> variants carry `.set_weight` compression gates for modules that do not exist
|
||||
> in vLLM-Omni's dense H3 attention, and the release states dense attention is
|
||||
> not a supported substitute for them, so they are refused at startup rather
|
||||
> than silently run wrong. Video Sparse Attention for H3's packed
|
||||
> `[text | cond | audio | video]` sequence is follow-up work.
|
||||
The VSA variants are supported through FastVideo's external kernel. Install a
|
||||
`fastvideo-kernel` build that provides the `fastvideo_kernel` Python module,
|
||||
then add the following flags to the same command:
|
||||
|
||||
```bash
|
||||
--diffusion-attention-backend FASTVIDEO_VSA \
|
||||
--fastvideo-vsa-topk 64
|
||||
```
|
||||
|
||||
FastH3 VSA applies its learned `.set_weight` compression gates to the complete
|
||||
packed `[text | cond | audio | video]` document using the official H3 geometry:
|
||||
text/condition/audio prefix tiles never cross segment boundaries, and target
|
||||
video rows use `(4, 4, 4)` 3-D tiles (64 tokens). Prefix queries remain dense;
|
||||
video queries select all prefix tiles plus the configured top-k video tiles.
|
||||
Pure Ulysses sequence parallelism is supported: the learned gate follows the
|
||||
same sequence-to-head all-to-all as Q/K/V before VSA runs. Ring and all-gather
|
||||
SP remain unsupported because they do not present a complete packed sequence
|
||||
to each block-sparse kernel rank.
|
||||
|
||||
The Dense / Data-Free variant does not require `fastvideo-kernel` and should be
|
||||
served with a dense attention backend.
|
||||
|
||||
Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at
|
||||
1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded:
|
||||
|
||||
@@ -119,6 +119,30 @@ MODEL=Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
|
||||
bash examples/online_serving/text_to_video/run_curl_sana_video.sh
|
||||
```
|
||||
|
||||
##### Parallel native serving
|
||||
|
||||
The native pipelines support tensor parallelism (up to 2 GPUs) and CFG
|
||||
parallelism. Tensor parallelism on 2 GPUs:
|
||||
|
||||
```bash
|
||||
vllm serve Efficient-Large-Model/SANA-Video_2B_480p_diffusers \
|
||||
--omni \
|
||||
--model-class-name SanaVideoPipeline \
|
||||
--tensor-parallel-size 2 \
|
||||
--dtype bfloat16 \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
For CFG parallelism on 2 GPUs, replace `--tensor-parallel-size 2` with
|
||||
`--cfg-parallel-size 2`; it splits the guided and unguided branches across the
|
||||
two GPUs and only helps when `guidance_scale` is above 1. Passing both flags
|
||||
combines the two on 4 GPUs.
|
||||
|
||||
Whether tensor parallelism lowers latency depends on the interconnect. Each
|
||||
transformer block adds an all-reduce, so it speeds up generation only on fast
|
||||
GPU links such as NVLink and can be slower than a single GPU on PCIe-only
|
||||
systems. Measure on your hardware before enabling it.
|
||||
|
||||
To run the black-box compatibility backend for T2V, replace the server script
|
||||
with `run_server_sana_video_diffusers.sh`. The same `/v1/videos` request
|
||||
works; `num_frames` is adapted to Diffusers' `frames` argument. The script
|
||||
@@ -178,6 +202,75 @@ Use the native `SanaVideoPipeline` and `SanaImageToVideoPipeline` for the
|
||||
primary SANA execution paths. The Diffusers adapter is retained as a
|
||||
validated compatibility/reference backend.
|
||||
|
||||
#### Native Cache-DiT and CPU offload
|
||||
|
||||
The native T2V and I2V pipelines support Cache-DiT, model-level CPU offload,
|
||||
and layerwise CPU offload through the common diffusion flags. Layerwise mode
|
||||
keeps non-block transformer modules, the text encoder, and the VAE on the
|
||||
runtime device while prefetching DiT blocks in order.
|
||||
|
||||
Add one of the following flag sets to either native offline command above:
|
||||
|
||||
```bash
|
||||
# Cache-DiT
|
||||
--cache-backend cache_dit
|
||||
|
||||
# Model-level CPU offload
|
||||
--enable-cpu-offload
|
||||
|
||||
# Layerwise CPU offload
|
||||
--enable-layerwise-offload
|
||||
|
||||
# Cache-DiT plus model-level offload
|
||||
--cache-backend cache_dit --enable-cpu-offload
|
||||
|
||||
# Cache-DiT plus layerwise offload
|
||||
--cache-backend cache_dit --enable-layerwise-offload
|
||||
|
||||
# Distributed layerwise offload (weights sharded across the DP group)
|
||||
--enable-distributed-layerwise-offload
|
||||
```
|
||||
|
||||
If both CPU offload flags are supplied, the common offloader keeps its existing
|
||||
layerwise precedence. Cache-DiT refreshes each request from the explicit
|
||||
`num_inference_steps`; when that field is omitted, both native SANA pipelines
|
||||
default to 50 inference steps.
|
||||
|
||||
Cache-DiT and CPU offload require TP1, CFG1, and SP1. Combining them with
|
||||
tensor, CFG, or sequence parallelism raises before checkpoint components are
|
||||
loaded, as do other cache backends such as TeaCache. Cache-DiT cannot be
|
||||
combined with distributed layerwise offload: per-rank cache skips would
|
||||
desynchronize the weight AllGather.
|
||||
|
||||
Measured on one A800-SXM4-80GB (driver 580.126.09, CUDA 13.0, PyTorch
|
||||
2.11.0+cu130, Diffusers 0.38.0) with `DIFFUSION_ATTENTION_BACKEND=CUDNN_ATTN`,
|
||||
81 frames, 50 steps, seed 42, each configuration in its own process:
|
||||
|
||||
| Configuration | 480p T2V latency | 480p T2V generation peak |
|
||||
|---|---:|---:|
|
||||
| Baseline | 120.19 s | 24.07 GiB |
|
||||
| Cache-DiT | 77.11 s (1.56x) | 24.07 GiB |
|
||||
| Model CPU offload | 121.15 s | 16.81 GiB (-7.26 GiB) |
|
||||
| Layerwise offload | 118.63 s | 20.75 GiB (-3.32 GiB) |
|
||||
|
||||
Cache-DiT gives 1.54x on 720p I2V (37.0 s to 24.1 s). Latency is the median of
|
||||
three measured runs after one warmup. Offload trades startup peak too: 14.28
|
||||
GiB baseline against 6.14 GiB for model-level and 6.80 GiB for layerwise.
|
||||
|
||||
Cache-DiT is approximate: it skips block computations, so its output differs
|
||||
from an uncached run. Similarity against a reference run is not used as a gate
|
||||
here because this pipeline does not reproduce its own trajectory on this
|
||||
hardware -- two runs of identical code differ by SSIM 0.9712 / rel_l2 8.7e-2 at
|
||||
50 steps. Cache-DiT output was checked frame by frame against the uncached run
|
||||
on 480p T2V and 720p I2V instead, and matches it in subject, composition,
|
||||
sharpness, and lighting.
|
||||
|
||||
The numbers above use the offline example's default cache configuration.
|
||||
Lowering `residual_diff_threshold` in `cache_config` caches fewer steps,
|
||||
trading speed for a trajectory closer to the uncached run; see the
|
||||
[Cache-DiT guide](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md)
|
||||
for the tuning knobs.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Key flags: select I2V explicitly with `--model-class-name
|
||||
@@ -194,8 +287,14 @@ validated compatibility/reference backend.
|
||||
The denoising loop intentionally retains the checkpoint-compatible
|
||||
Diffusers `DPMSolverMultistepScheduler`.
|
||||
- Known limitations:
|
||||
- Sequence/tensor/CFG parallelism, Cache-DiT, TeaCache, and step execution
|
||||
are not validated for the native pipeline.
|
||||
- Tensor parallelism (up to 2 GPUs) and CFG parallelism are supported for
|
||||
the native pipeline. Sequence parallelism, TeaCache, and step execution
|
||||
are not validated.
|
||||
- Cache-DiT and CPU offload are limited to TP1/CFG1/SP1. Cache-DiT with
|
||||
distributed layerwise offload is not supported by the native pipeline.
|
||||
- Cache-DiT speedup and offload memory numbers above are single-GPU A800
|
||||
measurements; other hardware will differ. Distributed layerwise offload is
|
||||
only exercised single-rank here.
|
||||
- The Diffusers backend is a compatibility path and does not provide native
|
||||
vLLM-Omni parallelism or continuous batching.
|
||||
- Native describes pipeline and Transformer ownership, not a zero-Diffusers
|
||||
|
||||
@@ -166,6 +166,7 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-ty
|
||||
- Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config.
|
||||
- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`.
|
||||
- Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail.
|
||||
- Stored voices: Uploaded and precomputed voices use the Base task and require a Base checkpoint. Built-in preset speakers remain CustomVoice voices.
|
||||
- Base codec termination: Base requests without an explicit `max_new_tokens`
|
||||
use a text-scaled safety ceiling (at least 192 codec frames and no more than
|
||||
the configured model limit). If the Talker reaches that ceiling without
|
||||
|
||||
@@ -78,6 +78,7 @@ recipes/
|
||||
| [`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 |
|
||||
| [`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 |
|
||||
| [`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 |
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
# 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.
|
||||
@@ -13,8 +13,9 @@ transformers >= 5.10.1, < 5.15
|
||||
av>=14.0.0
|
||||
omegaconf>=2.3.0
|
||||
diffusers==0.40.0
|
||||
# Hub kernels, including LTX-2.5 NATTEN
|
||||
kernels==0.14.1
|
||||
# Hub kernels, including LTX-2.5 NATTEN. Keep this pin in sync with the
|
||||
# supported transformers range above.
|
||||
kernels==0.15.2
|
||||
safetensors>=0.8.0
|
||||
accelerate==1.12.0
|
||||
soundfile>=0.13.1
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
首先,我分析所有输入图像:图像1是一个圆形的logo,设计现代且抽象。它由不同色调的蓝色(深蓝、中蓝、浅蓝)和白色构成,这些色块以流畅的曲线相互交织,形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形,背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴,<recaption>这幅图像以产品摄影的精致风格,呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。冰箱贴居中放置在`image_2`的灰色织物背景上,其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感:表面覆盖着一层光滑如镜的珐琅釉面,反射出柔和而清晰的高光;图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒,增强了立体感。柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡,并在其下方投下淡淡的、轮廓模糊的阴影,使其与织物背景无缝融合,营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓,以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构,精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果,以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>
|
||||
@@ -0,0 +1 @@
|
||||
首先,我分析所有输入图像:图像1是一个圆形logo,背景为黑色,主体由深蓝、中蓝、浅蓝和白色组成的抽象曲线构成,整体呈现出一种动态的、类似旋涡的视觉效果。图像2展示了四个并排摆放在灰色织物表面上的冰箱贴。这些冰箱贴的造型分别是长颈鹿、小鸟、蜜蜂和海豚,其关键特征是材质:它们具有光滑、高反光的珐琅质感,边缘有明显的银色金属包边,呈现出立体感。用户的指令是“基于图一的logo,参考图二中冰箱贴的材质,制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计(logo)转化为一个具有特定三维物理属性(冰箱贴材质)的实体物品。核心在于保留logo的设计精髓,同时精确复刻参考图像中冰箱贴的材质和工艺感。为了构建最终的答案图像,我会将图一的圆形logo作为新冰箱贴的主体图案。这个圆形冰箱贴将拥有与图二冰箱贴一致的厚度和立体感。其表面会覆盖一层高光泽的透明釉质,完美再现珐琅的质感。logo内部的蓝色和白色区域将保持其原有的色彩,但会因为表面的光泽而出现明亮的高光反射。整个圆形的边缘会包裹一圈纤细而有光泽的银色金属边框,这不仅增加了产品的精致感,也与图二中的金属包边工艺相呼应。最后,将这个新制作的圆形冰箱贴放置在图二所示的灰色织物背景上,确保其光照、阴影和反射效果与周围环境和其他冰箱贴完全一致,使其看起来就像是和它们一起被拍摄的,从而创造出一个逼真且风格统一的视觉效果。</think><recaption>这幅图像以产品摄影的精准风格,展示了一枚根据`image_1`的logo设计定制的圆形冰箱贴。最终图像使用`image_2`的分辨率。该冰箱贴居中放置于`image_2`的灰色织物背景之上,其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝及白色构成的动态旋涡状抽象图案。整个冰箱贴被赋予了与`image_2`中冰箱贴完全一致的高级珐琅质感,表面覆盖着一层晶莹剔透的高光釉面,在柔和的顶光照射下,呈现出清晰而温润的高光反射。一圈纤薄而精致的银色金属边框紧密包裹着圆形主体,其金属光泽与釉面的质感形成精妙对比。冰箱贴在织物背景上投下了一层极其自然的柔和阴影,证明了其与环境光影的无缝融合,整体视觉效果真实可信,宛如一件精心制作的实体工艺品。<relation_1>最终图像完整保留了`image_1`中logo的全部核心设计元素。这包括其完美的圆形轮廓,以及内部由深蓝、中蓝、浅蓝和白色四种颜色构成的抽象旋涡图案的精确形态、色彩和布局。</relation_1><relation_2>最终图像的分辨率、背景和材质风格完全来自`image_2`。这包括其作为背景的灰色织物表面的独特纹理,以及应用于新冰箱贴上的、与`image_2`中冰箱贴完全相同的珐琅材质(高光釉面质感)和银色金属边框的物理属性。画面的整体光照环境,包括柔和的光源方向和质感,也与`image_2`保持一致。</relation_2></recaption><answer><boi><img_size_1024><img_ratio_36><timestep><guidance><timestep_r>[<img>]{3600}<eoi></answer>
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.7 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.6 MiB |
@@ -0,0 +1,42 @@
|
||||
# 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,4 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
@@ -13,6 +13,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "com
|
||||
from skip_ci import resolve_ci_decision # noqa: E402
|
||||
from upload_pipeline import ( # noqa: E402
|
||||
_expand_mirror_hardwares,
|
||||
_get_mirror_hw_selector,
|
||||
_load_bootstrap_steps,
|
||||
_render_bootstrap_pipeline,
|
||||
_render_test_pipeline,
|
||||
@@ -116,7 +117,8 @@ def test_yaml_gated_l2_still_enables_image_via_ready_base() -> None:
|
||||
assert "if: true" not in rendered
|
||||
|
||||
|
||||
def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins() -> None:
|
||||
def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "")
|
||||
doc = {
|
||||
"steps": [
|
||||
{
|
||||
@@ -129,8 +131,29 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins() -> None:
|
||||
rendered = _render_test_pipeline(doc, changed_files=None)
|
||||
step = rendered["steps"][0]
|
||||
assert "mirror_hardwares" not in step
|
||||
assert step["agents"]["queue"] == "gpu_1_queue"
|
||||
assert step["plugins"][0]["docker#v5.2.0"]["image"].endswith("$BUILDKITE_COMMIT")
|
||||
assert step["agents"]["queue"] == "l4-k8s"
|
||||
assert step["retry"] == {
|
||||
"automatic": [
|
||||
{"exit_status": -1, "limit": 1},
|
||||
{"exit_status": 128, "limit": 1},
|
||||
{"signal_reason": "agent_stop", "limit": 1},
|
||||
{"signal_reason": "agent_refused", "limit": 1},
|
||||
],
|
||||
}
|
||||
container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0]
|
||||
assert container["image"].endswith("$BUILDKITE_COMMIT")
|
||||
assert container["resources"]["limits"]["nvidia.com/gpu"] == 1
|
||||
|
||||
|
||||
def test_mirror_hardwares_l4_preserves_explicit_retry() -> None:
|
||||
step = _expand_mirror_hardwares(
|
||||
{
|
||||
"label": "opt-out",
|
||||
"mirror_hardwares": "l4_1",
|
||||
"retry": {"automatic": [{"exit_status": 255, "limit": 2}]},
|
||||
},
|
||||
)
|
||||
assert step["retry"] == {"automatic": [{"exit_status": 255, "limit": 2}]}
|
||||
|
||||
|
||||
def test_mirror_hardwares_conflicts_with_explicit_agents() -> None:
|
||||
@@ -161,6 +184,156 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
|
||||
]
|
||||
|
||||
|
||||
def _gpu_limit(step: dict) -> int:
|
||||
return step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0]["resources"]["limits"]["nvidia.com/gpu"]
|
||||
|
||||
|
||||
def test_mirror_hardwares_mapping_is_rejected() -> None:
|
||||
with pytest.raises(ValueError, match="must be a preset string"):
|
||||
_expand_mirror_hardwares(
|
||||
{"label": "bad", "mirror_hardwares": {"default": "h100_2", "b200": "b200_2"}},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("hardware", [2, "2", "h100_99", "not_a_preset"])
|
||||
def test_mirror_hardwares_unknown_name_is_rejected(hardware: int | str) -> None:
|
||||
with pytest.raises(ValueError, match="unknown mirror_hardwares"):
|
||||
_expand_mirror_hardwares({"label": "unknown preset", "mirror_hardwares": hardware})
|
||||
|
||||
|
||||
def test_mirror_hardwares_string_ignores_pytest_cards_mark(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "")
|
||||
step = _expand_mirror_hardwares(
|
||||
{
|
||||
"label": "forced",
|
||||
"commands": ['pytest -sv tests/e2e -m "full_model and L4 and cards_4"'],
|
||||
"mirror_hardwares": "h100_1",
|
||||
},
|
||||
)
|
||||
assert step is not None
|
||||
assert step["agents"]["queue"] == "mithril-h100-pool"
|
||||
assert _gpu_limit(step) == 1
|
||||
|
||||
|
||||
def test_mirror_hardwares_b200_omits_cuda_strings_and_remaps_inferred(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "b200")
|
||||
assert _expand_mirror_hardwares({"label": "h100", "mirror_hardwares": "h100_4"}) is None
|
||||
assert _expand_mirror_hardwares({"label": "l4", "mirror_hardwares": "l4_1"}) is None
|
||||
b200 = _expand_mirror_hardwares({"label": "b200", "mirror_hardwares": "b200_2"})
|
||||
assert b200 is not None and b200["agents"]["queue"] == "b200-k8s"
|
||||
npu = _expand_mirror_hardwares({"label": "npu", "mirror_hardwares": "a2b3_npu_4"})
|
||||
assert npu is not None and npu["agents"]["queue"] == "ascend-a2b3"
|
||||
|
||||
rendered = _render_test_pipeline(
|
||||
{
|
||||
"steps": [
|
||||
{
|
||||
"group": ":card_index_dividers: Mixed",
|
||||
"steps": [
|
||||
{"label": "H100 only", "mirror_hardwares": "h100_4"},
|
||||
{
|
||||
"label": "Count remap",
|
||||
"commands": [
|
||||
'pytest -sv tests/e2e -m "full_model and H100 and B200 and omni and cards_2"',
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"group": ":card_index_dividers: H100 only group",
|
||||
"steps": [{"label": "Skip me", "mirror_hardwares": "h100_1"}],
|
||||
},
|
||||
],
|
||||
},
|
||||
changed_files=None,
|
||||
)
|
||||
groups = [step.get("group") for step in rendered["steps"]]
|
||||
assert ":card_index_dividers: H100 only group" not in groups
|
||||
mixed = next(step for step in rendered["steps"] if step.get("group") == ":card_index_dividers: Mixed")
|
||||
assert [child["label"] for child in mixed["steps"]] == ["Count remap"]
|
||||
assert mixed["steps"][0]["agents"]["queue"] == "b200-k8s"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("selector", "expr", "queue", "gpus"),
|
||||
[
|
||||
("", "H100 and B200 and cards_2", "mithril-h100-pool", 2),
|
||||
("b200", "H100 and B200 and cards_2", "b200-k8s", 2),
|
||||
("", "H100 or L4 and cards_4", "mithril-h100-pool", 4),
|
||||
("", "L4 and B200 and cards_4", "l4-k8s", 4),
|
||||
("", "H100 and cards_2 and cards_3", "mithril-h100-pool", 3),
|
||||
("", "H100 and not cards_1", "mithril-h100-pool", 4),
|
||||
("b200", "H100 and B200 and not cards_1", "b200-k8s", 4),
|
||||
],
|
||||
)
|
||||
def test_mirror_hardwares_inferred_from_marks(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
selector: str,
|
||||
expr: str,
|
||||
queue: str,
|
||||
gpus: int,
|
||||
) -> None:
|
||||
monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: selector)
|
||||
step = _expand_mirror_hardwares({"label": "job", "commands": [f'pytest -sv tests/e2e -m "{expr}"']})
|
||||
assert step is not None
|
||||
assert step["agents"]["queue"] == queue
|
||||
assert _gpu_limit(step) == gpus
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("selector", "expr"),
|
||||
[
|
||||
("b200", "H100 and cards_2"),
|
||||
("", "B200 and cards_2"),
|
||||
("", "full_model and cards_2"),
|
||||
],
|
||||
)
|
||||
def test_mirror_hardwares_inferred_skips_unmatched_chip(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
selector: str,
|
||||
expr: str,
|
||||
) -> None:
|
||||
monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: selector)
|
||||
assert _expand_mirror_hardwares({"label": "job", "commands": [f'pytest -sv tests/e2e -m "{expr}"']}) is None
|
||||
|
||||
|
||||
def test_mirror_hardwares_inferred_missing_preset_is_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "")
|
||||
with pytest.raises(ValueError, match="has no preset 'h100_8'"):
|
||||
_expand_mirror_hardwares(
|
||||
{"label": "H100 8-gpu", "commands": ['pytest -sv tests/e2e -m "H100 and cards_8"']},
|
||||
)
|
||||
|
||||
|
||||
def test_cpu_step_without_mirror_hardwares_is_unchanged() -> None:
|
||||
step = {"label": "CPU report", "commands": ["echo ok"], "agents": {"queue": "cpu_queue_premerge"}}
|
||||
assert _expand_mirror_hardwares(step) is step
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("raw", "expected"), [("", ""), (" ", ""), ("b200", "b200"), ("B200", "b200")])
|
||||
def test_mirror_hw_selector_empty_or_b200(monkeypatch: pytest.MonkeyPatch, raw: str, expected: str) -> None:
|
||||
monkeypatch.setenv("MIRROR_HW", raw)
|
||||
assert _get_mirror_hw_selector() == expected
|
||||
|
||||
|
||||
def test_mirror_hw_typo_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("MIRROR_HW", "b20o")
|
||||
with pytest.raises(ValueError, match=r"unsupported MIRROR_HW='b20o'"):
|
||||
_get_mirror_hw_selector()
|
||||
with pytest.raises(ValueError, match=r"unsupported MIRROR_HW='b20o'"):
|
||||
_render_test_pipeline(
|
||||
{
|
||||
"steps": [
|
||||
{"label": "CPU report", "commands": ["echo ok"]},
|
||||
{"label": "H100 string", "mirror_hardwares": "h100_4"},
|
||||
],
|
||||
},
|
||||
changed_files=None,
|
||||
)
|
||||
|
||||
|
||||
def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
|
||||
rendered = _render_test_pipeline(doc, changed_files=changed_files)
|
||||
labels: set[str] = set()
|
||||
|
||||
@@ -41,6 +41,11 @@ from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
SINGLE_STAGE_PIPE_CFG = (StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
|
||||
# Validation strings to match against in post init
|
||||
NO_STAGES_MATCH_STR = "no stages"
|
||||
NO_TERMINAL_STAGE_MATCH_STR = "No terminal stage"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _stable_test_platform(monkeypatch):
|
||||
@@ -618,6 +623,7 @@ class TestPipelineDiscovery:
|
||||
p = PipelineConfig(
|
||||
model_type="custom_collide",
|
||||
hf_architectures=("SomeCollidingArch",),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
assert p.hf_architectures == ("SomeCollidingArch",)
|
||||
|
||||
@@ -639,41 +645,42 @@ class TestStagePipelineConfig:
|
||||
|
||||
|
||||
class TestPipelineConfigNew:
|
||||
def test_simple_init(self):
|
||||
# Ensure pipeline config validates in post init.
|
||||
PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
|
||||
def test_frozen(self):
|
||||
p = PipelineConfig(model_type="t", model_arch="A")
|
||||
p = PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
with pytest.raises(AttributeError):
|
||||
p.model_type = "changed"
|
||||
|
||||
def test_validate_valid(self):
|
||||
p = PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=(
|
||||
StagePipelineConfig(stage_id=0, model_stage="a"),
|
||||
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,), final_output=True),
|
||||
),
|
||||
)
|
||||
assert p.validate() == []
|
||||
|
||||
def test_validate_no_stages(self):
|
||||
p = PipelineConfig(model_type="t", model_arch="A")
|
||||
assert any("no stages" in e.lower() for e in p.validate())
|
||||
with pytest.raises(ValueError, match=NO_STAGES_MATCH_STR):
|
||||
PipelineConfig(model_type="t", model_arch="A")
|
||||
|
||||
def test_validate_no_terminal_stage(self):
|
||||
"""A pipeline with no ``final_output`` stage can never emit a result."""
|
||||
p = PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=(
|
||||
StagePipelineConfig(stage_id=0, model_stage="a"),
|
||||
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
|
||||
),
|
||||
)
|
||||
assert any("no terminal stage" in e.lower() for e in p.validate())
|
||||
with pytest.raises(ValueError, match=NO_TERMINAL_STAGE_MATCH_STR):
|
||||
PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=(
|
||||
StagePipelineConfig(stage_id=0, model_stage="a"),
|
||||
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
|
||||
),
|
||||
)
|
||||
|
||||
def test_validate_terminal_stage_need_not_be_last(self):
|
||||
"""The check is cycle-agnostic: any stage may carry ``final_output``."""
|
||||
p = PipelineConfig(
|
||||
def test_pipeline_can_have_final_output_in_any_stage(self):
|
||||
"""Ensure any stage may carry ``final_output``."""
|
||||
PipelineConfig(
|
||||
model_type="t",
|
||||
model_arch="A",
|
||||
stages=(
|
||||
@@ -681,15 +688,20 @@ class TestPipelineConfigNew:
|
||||
StagePipelineConfig(stage_id=1, model_stage="b", input_sources=(0,)),
|
||||
),
|
||||
)
|
||||
assert p.validate() == []
|
||||
|
||||
|
||||
class TestPipelineRegistration:
|
||||
def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path):
|
||||
deploy_key = "deploy_selected_pipeline"
|
||||
model_type_key = "hf_model_type_pipeline"
|
||||
deploy_pipe = PipelineConfig(model_type=deploy_key)
|
||||
model_type_pipe = PipelineConfig(model_type=model_type_key)
|
||||
deploy_pipe = PipelineConfig(
|
||||
model_type=deploy_key,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
model_type_pipe = PipelineConfig(
|
||||
model_type=model_type_key,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(deploy_pipe)
|
||||
register_pipeline(model_type_pipe)
|
||||
|
||||
@@ -715,7 +727,11 @@ class TestPipelineRegistration:
|
||||
tmp_path,
|
||||
):
|
||||
model_type_key = "registered_model_type_pipeline"
|
||||
register_pipeline(PipelineConfig(model_type=model_type_key))
|
||||
pipeline_cfg = PipelineConfig(
|
||||
model_type=model_type_key,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipeline_cfg)
|
||||
|
||||
deploy_path = tmp_path / "deploy.yaml"
|
||||
deploy_path.write_text("pipeline: missing_pipeline\n", encoding="utf-8")
|
||||
@@ -750,6 +766,7 @@ class TestPipelineRegistration:
|
||||
pipe_cfg = PipelineConfig(
|
||||
model_type=pipeline_key,
|
||||
hf_architectures=("ArchitectureFallbackForTest",),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipe_cfg)
|
||||
|
||||
@@ -780,15 +797,18 @@ class TestPipelineRegistration:
|
||||
model_type="rejected_by_predicate",
|
||||
hf_architectures=(shared_arch,),
|
||||
hf_config_predicate=rejecting_predicate,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
raises_cfg = PipelineConfig(
|
||||
model_type="raises_in_predicate",
|
||||
hf_architectures=(shared_arch,),
|
||||
hf_config_predicate=raising_predicate,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
accept_cfg = PipelineConfig(
|
||||
model_type="accepted_by_predicate",
|
||||
hf_architectures=(shared_arch,),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(reject_cfg)
|
||||
register_pipeline(raises_cfg)
|
||||
@@ -822,6 +842,7 @@ class TestPipelineRegistration:
|
||||
model_type="predicate_without_arch_match",
|
||||
hf_architectures=("DifferentArchitectureForTest",),
|
||||
hf_config_predicate=predicate,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipe_cfg)
|
||||
|
||||
@@ -845,6 +866,7 @@ class TestPipelineRegistration:
|
||||
resolved_cfg = PipelineConfig(
|
||||
model_type="callable_resolved_pipeline",
|
||||
hf_architectures=("CallableArchitectureForTest",),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
seen_hf_configs = []
|
||||
|
||||
@@ -918,7 +940,10 @@ class TestPipelineRegistration:
|
||||
def test_pipeline_registration(self, clean_pipeline_registry):
|
||||
"""Ensure that we can register and create a custom pipeline config."""
|
||||
new_model_type = "new_model_type"
|
||||
pipe_cfg = PipelineConfig(model_type=new_model_type)
|
||||
pipe_cfg = PipelineConfig(
|
||||
model_type=new_model_type,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
|
||||
# Register the new PipelineConfig
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
@@ -961,7 +986,10 @@ class TestPipelineRegistration:
|
||||
def custom_resolver(
|
||||
hf_config: FakeConfig,
|
||||
) -> PipelineConfig:
|
||||
return PipelineConfig(model_type=resolved_type)
|
||||
return PipelineConfig(
|
||||
model_type=resolved_type,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
|
||||
# Register the new PipelineConfig
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
@@ -1018,7 +1046,7 @@ class TestPipelineRegistration:
|
||||
pipe_cfg = PipelineConfig(
|
||||
model_type=pipeline_key,
|
||||
model_arch="DeployOnlyArch",
|
||||
stages=(StagePipelineConfig(stage_id=0, model_stage="diffusion"),),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipe_cfg)
|
||||
|
||||
@@ -1051,7 +1079,11 @@ class TestPipelineRegistration:
|
||||
|
||||
def test_structured_path_loads_explicit_deploy_config_once(self, clean_pipeline_registry, tmp_path):
|
||||
pipeline_key = "single_load_pipeline"
|
||||
register_pipeline(PipelineConfig(model_type=pipeline_key))
|
||||
pipeline_cfg = PipelineConfig(
|
||||
model_type=pipeline_key,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipeline_cfg)
|
||||
deploy_path = tmp_path / "single_load.yaml"
|
||||
deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8")
|
||||
|
||||
@@ -1081,6 +1113,7 @@ class TestPipelineRegistration:
|
||||
pipeline = PipelineConfig(
|
||||
model_type="pipeline_with_default",
|
||||
default_deploy_config_name=default_name,
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
|
||||
with patch(
|
||||
@@ -1100,8 +1133,16 @@ class TestPipelineRegistration:
|
||||
OmniServingCapability.COMPLETIONS,
|
||||
"pipeline_a blocks completions",
|
||||
)
|
||||
pipe_a = PipelineConfig(model_type="detect_type", endpoint_restrictions=(restriction,))
|
||||
pipe_b = PipelineConfig(model_type="override_type", endpoint_restrictions=())
|
||||
pipe_a = PipelineConfig(
|
||||
model_type="detect_type",
|
||||
endpoint_restrictions=(restriction,),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
pipe_b = PipelineConfig(
|
||||
model_type="override_type",
|
||||
endpoint_restrictions=(),
|
||||
stages=SINGLE_STAGE_PIPE_CFG,
|
||||
)
|
||||
register_pipeline(pipe_a)
|
||||
register_pipeline(pipe_b)
|
||||
|
||||
@@ -1641,6 +1682,7 @@ stages:
|
||||
model_stage="ar",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
requires_multimodal_data=True,
|
||||
final_output=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -1661,6 +1703,7 @@ stages:
|
||||
model_stage="ar",
|
||||
execution_type=StageExecutionType.LLM_AR,
|
||||
scheduler_cls=scheduler_cls,
|
||||
final_output=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
@@ -1906,7 +1949,6 @@ class TestQwen3OmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
|
||||
assert len(p.stages) == 3
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker(self):
|
||||
p = resolve_pipeline_config(
|
||||
@@ -1960,7 +2002,6 @@ class TestQwen2_5OmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
|
||||
assert len(p.stages) == 3
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker(self):
|
||||
p = resolve_pipeline_config("qwen2_5_omni")
|
||||
@@ -2009,7 +2050,6 @@ class TestQwen3TTSPipeline:
|
||||
assert p is not None
|
||||
assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_talker_stage(self):
|
||||
p = resolve_pipeline_config("qwen3_tts")
|
||||
@@ -2089,7 +2129,6 @@ class TestMingFlashOmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker_stage(self):
|
||||
p = resolve_pipeline_config("ming_flash_omni")
|
||||
@@ -2146,7 +2185,6 @@ class TestMingFlashOmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
|
||||
assert len(p.stages) == 1
|
||||
assert p.validate() == []
|
||||
|
||||
def test_tts_stage(self):
|
||||
p = resolve_pipeline_config("ming_flash_omni_tts")
|
||||
@@ -2203,7 +2241,6 @@ class TestMingFlashOmniPipeline:
|
||||
assert isinstance(p, PipelineConfig)
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 1
|
||||
assert p.validate() == []
|
||||
|
||||
def test_thinker_only_stage(self):
|
||||
p = resolve_pipeline_config("ming_flash_omni_thinker_only")
|
||||
@@ -2242,7 +2279,6 @@ class TestMingFlashOmniPipeline:
|
||||
assert p is not None
|
||||
assert p.model_arch == "MingFlashOmniForConditionalGeneration"
|
||||
assert len(p.stages) == 2
|
||||
assert p.validate() == []
|
||||
|
||||
def test_image_thinker_stage(self):
|
||||
s = resolve_pipeline_config("ming_flash_omni_image").get_stage(0)
|
||||
@@ -2363,6 +2399,21 @@ class TestBaseConfigInheritance:
|
||||
# CI overrides max_tokens
|
||||
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):
|
||||
"""An overlay with only ``base_config`` inherits everything."""
|
||||
base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
|
||||
@@ -2824,6 +2875,54 @@ stages:
|
||||
assert omega_config.engine_args.max_model_len == 8192
|
||||
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:
|
||||
"""Caller-typed (non-None) values win over YAML; None values fall through
|
||||
|
||||
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 24,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 58,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
|
||||
EnvironmentVariableCategory.INTERNAL: 2,
|
||||
}
|
||||
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
|
||||
)
|
||||
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 33,
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 35,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
|
||||
|
||||
@@ -45,6 +45,7 @@ from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
StageDeployConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
)
|
||||
@@ -1039,7 +1040,10 @@ def test_from_pipeline_config_uses_hf_config_for_callable_resolver():
|
||||
|
||||
|
||||
def test_from_pipeline_config_accepts_pre_resolved_pipeline():
|
||||
resolved_pipeline = PipelineConfig(model_type="callable_resolved_variant")
|
||||
resolved_pipeline = PipelineConfig(
|
||||
model_type="callable_resolved_variant",
|
||||
stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
|
||||
)
|
||||
|
||||
omni_config = VllmOmniConfig.from_pipeline_config(resolved_pipeline)
|
||||
|
||||
@@ -1072,6 +1076,7 @@ def test_from_pipeline_config_default_deploy_name_ignores_cwd(monkeypatch, tmp_p
|
||||
pipeline = PipelineConfig(
|
||||
model_type="pipeline_with_default",
|
||||
default_deploy_config_name=default_name,
|
||||
stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),),
|
||||
)
|
||||
loaded_paths = []
|
||||
|
||||
|
||||
@@ -1,16 +1,22 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Tests for out of tree registration to OMNI_PIPELINES."""
|
||||
|
||||
import pytest
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline
|
||||
from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver
|
||||
from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig, pipeline_cfg_resolver
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def build_fake_pipeline_config(model_type: str) -> PipelineConfig:
|
||||
return PipelineConfig(
|
||||
model_type=model_type, stages=(StagePipelineConfig(stage_id=0, model_stage="a", final_output=True),)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def custom_resolver():
|
||||
"""Build a reusable custom resolver for PipelineConfigs."""
|
||||
@@ -22,7 +28,7 @@ def custom_resolver():
|
||||
def custom_resolver(
|
||||
hf_config: CustomConfigType,
|
||||
) -> PipelineConfig:
|
||||
return PipelineConfig(model_type="resolved_type")
|
||||
return build_fake_pipeline_config("resolved_type")
|
||||
|
||||
return custom_resolver
|
||||
|
||||
@@ -30,7 +36,7 @@ def custom_resolver():
|
||||
def test_register_pipeline_config(clean_pipeline_registry):
|
||||
"""Ensure that we can register a custom pipeline config to OMNI_PIPELINES."""
|
||||
new_model_type = "new_model_type"
|
||||
pipe_cfg = PipelineConfig(model_type=new_model_type)
|
||||
pipe_cfg = build_fake_pipeline_config(new_model_type)
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
register_pipeline(pipe_cfg)
|
||||
assert new_model_type in OMNI_PIPELINES
|
||||
@@ -41,7 +47,7 @@ def test_register_pipeline_config_with_model_type(clean_pipeline_registry):
|
||||
"""Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES."""
|
||||
new_model_type = "new_model_type"
|
||||
unused_model_type = "foo"
|
||||
pipe_cfg = PipelineConfig(model_type=unused_model_type)
|
||||
pipe_cfg = build_fake_pipeline_config(unused_model_type)
|
||||
assert new_model_type not in OMNI_PIPELINES
|
||||
assert unused_model_type not in OMNI_PIPELINES
|
||||
|
||||
@@ -73,20 +79,3 @@ def test_minimax_h3_disaggregation_is_explicit_opt_in():
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
assert pipeline.model_type == "minimax_h3_disaggregated"
|
||||
|
||||
|
||||
def test_in_tree_pipelines_pass_topology_validation():
|
||||
"""Every statically registered in-tree pipeline must have a valid topology.
|
||||
|
||||
``PipelineConfig.validate()`` only runs on ``register_pipeline``, which
|
||||
in-tree pipelines never go through, so a shipped topology error would
|
||||
otherwise stay invisible until serving hangs or misroutes.
|
||||
Resolver entries are skipped: they need an ``hf_config`` to produce a
|
||||
``PipelineConfig``.
|
||||
"""
|
||||
failures = {
|
||||
model_type: errors
|
||||
for model_type, pipeline in OMNI_PIPELINES.items()
|
||||
if isinstance(pipeline, PipelineConfig) and (errors := pipeline.validate())
|
||||
}
|
||||
assert not failures, f"in-tree pipelines failed topology validation: {failures}"
|
||||
@@ -50,9 +50,6 @@ from tests.dfx.conftest import (
|
||||
)
|
||||
from tests.helpers.runtime import get_open_port
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inline field processing
|
||||
@@ -837,7 +834,7 @@ def run_benchmark(
|
||||
"Quantization": _to_quantization_value(server_type, serve_args_dict),
|
||||
"offload": _to_offload_string(server_type, serve_args_dict),
|
||||
"compile": _to_compile_value(server_type, serve_args_dict),
|
||||
"Attn_backend": os.environ.get("DIFFUSION_ATTENTION_BACKEND", ""),
|
||||
"Attn_backend": str(serve_args_dict.get("diffusion-attention-backend") or ""),
|
||||
"num_inference_steps": params.get("num-inference-steps", ""),
|
||||
"completed": completed,
|
||||
"failed": failed,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_boogu_image_edit_single_device",
|
||||
"description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. Text-only edit path (2 predictions/step); a double-guidance (guidance_scale_2) case is a follow-up since the random dataset cannot inject extra_body params. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.",
|
||||
"description": "Single-device baseline (no parallelism), concurrency sweep 1/2/4 at 512x512, 28 steps. This row measures text-only edit guidance (2 predictions/step); dedicated rows below use extra-body to measure double guidance. Baselines recorded on 1x A40 46GB (measured c1/c2/c4: qps 0.0337/0.0336/0.0336, latency_mean 29.64/56.50/101.48 s, peak mem 36754 MB) with ~10% margin.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Boogu/Boogu-Image-0.1-Edit",
|
||||
@@ -33,5 +33,184 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_boogu_image_edit_double_guidance_single_device",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device sequential three-branch double-guidance control row for direct A/B comparison with cfg_parallel_size=2 and 3 on identical H100 hardware.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Boogu/Boogu-Image-0.1-Edit",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps28_i2i_double_cfg1",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 28,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"warmup-requests": 4,
|
||||
"warmup-concurrency": 1,
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"extra-body": {
|
||||
"guidance_scale": 5.0,
|
||||
"guidance_scale_2": 2.0,
|
||||
"seed": 42
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_boogu_image_edit_text_cfg_parallel_2",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-device CFG-parallel Edit text-only guidance measurement row. Compare with the single-device row on identical H100 hardware and software.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Boogu/Boogu-Image-0.1-Edit",
|
||||
"serve_args": {
|
||||
"cfg-parallel-size": 2,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps28_i2i_text_cfg2",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 28,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"warmup-requests": 4,
|
||||
"warmup-concurrency": 1,
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"extra-body": {
|
||||
"guidance_scale": 5.0,
|
||||
"guidance_scale_2": 1.0,
|
||||
"seed": 42
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_boogu_image_edit_double_cfg_parallel_2",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-device round-robin execution of the three Edit double-guidance branches. Populate A/B latency, throughput, and per-GPU memory after GPU validation.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Boogu/Boogu-Image-0.1-Edit",
|
||||
"serve_args": {
|
||||
"cfg-parallel-size": 2,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps28_i2i_double_cfg2",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 28,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"warmup-requests": 4,
|
||||
"warmup-concurrency": 1,
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"extra-body": {
|
||||
"guidance_scale": 5.0,
|
||||
"guidance_scale_2": 2.0,
|
||||
"seed": 42
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_boogu_image_edit_double_cfg_parallel_3",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 3
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Full three-device execution of the three Edit double-guidance branches. Compare directly with sequential and cfg_parallel_size=2 on identical H100 hardware.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Boogu/Boogu-Image-0.1-Edit",
|
||||
"serve_args": {
|
||||
"cfg-parallel-size": 3,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps28_i2i_double_cfg3",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 28,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"warmup-requests": 4,
|
||||
"warmup-concurrency": 1,
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"extra-body": {
|
||||
"guidance_scale": 5.0,
|
||||
"guidance_scale_2": 2.0,
|
||||
"seed": 42
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -33,5 +33,49 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_boogu_image_cfg_parallel_2",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-device CFG-parallel Base T2I measurement row. Run beside test_boogu_image_single_device on the same H100 revision and populate latency, throughput, speedup, and per-GPU peak-memory baselines from the resulting report.",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Boogu/Boogu-Image-0.1-Base",
|
||||
"serve_args": {
|
||||
"cfg-parallel-size": 2,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps28_cfg2",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 28,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"warmup-requests": 4,
|
||||
"warmup-concurrency": 1,
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"extra-body": {
|
||||
"guidance_scale": 4.0,
|
||||
"seed": 42
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -30,7 +30,6 @@
|
||||
"vae-patch-parallel-size": 4,
|
||||
"vae-parallel-mode": "tile",
|
||||
"vae-use-tiling": true,
|
||||
"diffusion-attention-backend": "FLASH_ATTN",
|
||||
"stage-init-timeout": 1800,
|
||||
"init-timeout": 1800,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
@@ -102,7 +101,6 @@
|
||||
"vae-patch-parallel-size": 4,
|
||||
"vae-parallel-mode": "tile",
|
||||
"vae-use-tiling": true,
|
||||
"diffusion-attention-backend": "FLASH_ATTN",
|
||||
"stage-init-timeout": 1800,
|
||||
"init-timeout": 1800,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
@@ -175,7 +173,6 @@
|
||||
"vae-patch-parallel-size": 4,
|
||||
"vae-parallel-mode": "tile",
|
||||
"vae-use-tiling": true,
|
||||
"diffusion-attention-backend": "FLASH_ATTN",
|
||||
"stage-init-timeout": 1800,
|
||||
"init-timeout": 1800,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
@@ -247,7 +244,6 @@
|
||||
"vae-patch-parallel-size": 4,
|
||||
"vae-parallel-mode": "tile",
|
||||
"vae-use-tiling": true,
|
||||
"diffusion-attention-backend": "FLASH_ATTN",
|
||||
"stage-init-timeout": 1800,
|
||||
"init-timeout": 1800,
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
|
||||
@@ -197,6 +197,7 @@ def test_create_unique_server_pytest_params_applies_marks(tmp_path):
|
||||
assert len(by_id["test_with_mark"].values) == 1
|
||||
assert isinstance(by_id["test_with_mark"].values[0], tuple)
|
||||
assert any(m.name == "H100" for m in by_id["test_with_mark"].marks)
|
||||
assert not any(m.name == "B200" for m in by_id["test_with_mark"].marks)
|
||||
assert not any(m.name == "H100" for m in by_id["test_without_mark"].marks)
|
||||
|
||||
|
||||
|
||||
@@ -18,6 +18,7 @@ import torch
|
||||
|
||||
import vllm_omni.diffusion.attention.layer as layer_mod
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.config import (
|
||||
get_current_diffusion_config,
|
||||
@@ -73,6 +74,12 @@ class TestAttentionSpec:
|
||||
assert bk["dtype_vo"] == "fp8_e4m3"
|
||||
assert bk["flashinfer_backend"] == "trtllm-gen"
|
||||
|
||||
def test_quant_flashinfer_backend_serialized_without_dtype(self):
|
||||
spec = AttentionSpec(backend="FLASHINFER_ATTN", quant={"flashinfer_backend": "fa2"})
|
||||
assert spec.quant is not None
|
||||
assert spec.quant.enabled is True
|
||||
assert spec.backend_kwargs() == {"quant": {"flashinfer_backend": "fa2"}}
|
||||
|
||||
def test_quant_and_skip_softmax_coexist(self):
|
||||
bk = AttentionSpec(
|
||||
backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5}
|
||||
@@ -150,8 +157,8 @@ class TestAttentionSpec:
|
||||
def test_block_sparse_precision_enum_input_normalized(self):
|
||||
"""RainFusionPrecision enum members must normalize to their string value.
|
||||
|
||||
Regression: the str,Enum members passed the value check but str() later
|
||||
produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it).
|
||||
Regression: ``str(str, Enum)`` later produced ``RainFusionPrecision.MIX``
|
||||
instead of ``mix``. ``BlockSparseSpec`` stores ``member.value``.
|
||||
"""
|
||||
from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig
|
||||
from vllm_omni.diffusion.data import RainFusionPrecision
|
||||
@@ -509,6 +516,49 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
|
||||
dense_flash.assert_not_called()
|
||||
sdpa.assert_called_once()
|
||||
|
||||
@pytest.mark.parametrize("fa4_available", [False, True])
|
||||
def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch, fa4_available):
|
||||
attention = Attention.__new__(Attention)
|
||||
attention._scheduler_paged_kv = True
|
||||
attention.paged_kv_cache_role = "primary"
|
||||
attention.backend_pref = "FLASH_ATTN"
|
||||
attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
|
||||
flash_output = torch.ones(1)
|
||||
sdpa_output = torch.zeros(1)
|
||||
dense_flash = Mock(return_value=flash_output)
|
||||
sdpa = Mock(return_value=sdpa_output)
|
||||
attention.attention = SimpleNamespace(forward=dense_flash)
|
||||
attention.sdpa_fallback = SimpleNamespace(forward=sdpa)
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
layer_mod,
|
||||
"get_forward_context",
|
||||
lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True),
|
||||
)
|
||||
monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: False)
|
||||
monkeypatch.setattr(layer_mod.current_omni_platform, "is_cuda", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
layer_mod.current_omni_platform,
|
||||
"supports_diffusion_dense_flash_attention",
|
||||
lambda: fa4_available,
|
||||
)
|
||||
|
||||
result = attention._run_local_attention(
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
AttentionMetadata(),
|
||||
)
|
||||
|
||||
if fa4_available:
|
||||
assert result is flash_output
|
||||
dense_flash.assert_called_once()
|
||||
sdpa.assert_not_called()
|
||||
else:
|
||||
assert result is sdpa_output
|
||||
dense_flash.assert_not_called()
|
||||
sdpa.assert_called_once()
|
||||
|
||||
def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
|
||||
attention = Attention.__new__(Attention)
|
||||
attention._scheduler_paged_kv = False
|
||||
@@ -579,7 +629,106 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
|
||||
|
||||
assert len(calls) == 1
|
||||
assert attention.attn_backend is _DenseBackend
|
||||
assert attention.backend_pref is None
|
||||
# Auto path records the resolved backend name for Ring/SP; it is not explicit.
|
||||
assert attention.backend_pref == _DenseBackend.get_name()
|
||||
assert attention.backend_explicit is False
|
||||
|
||||
def test_internal_paged_selector_does_not_require_dense_fa4(self, monkeypatch):
|
||||
"""Blackwell CUDNN default + marked paged layer must not act as explicit FLASH_ATTN."""
|
||||
|
||||
class _FakeAttentionImpl:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
def forward(self, query, key, value, attn_metadata=None):
|
||||
return query
|
||||
|
||||
class _CuDNNBackend:
|
||||
supports_paged_kv = False
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "CUDNN_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls():
|
||||
return _FakeAttentionImpl
|
||||
|
||||
def _fake_get_attn_backend_for_role(*args, attention_config=None, **kwargs):
|
||||
if attention_config is not None:
|
||||
spec, _ = attention_config.resolve_with_source(role="self", role_category=None)
|
||||
if spec is not None and spec.backend.upper() == "FLASH_ATTN":
|
||||
raise ValueError(
|
||||
"FLASH_ATTN was explicitly selected but is unsupported "
|
||||
"(Blackwell requires CuTe FlashAttention-4 "
|
||||
"(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only). "
|
||||
"Select a compatible backend."
|
||||
)
|
||||
return _CuDNNBackend, None
|
||||
|
||||
monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
|
||||
monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
|
||||
monkeypatch.setattr(FlashAttentionBackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
|
||||
monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(),
|
||||
diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
parallel_config=SimpleNamespace(ring_degree=1),
|
||||
diffusion_kv_cache_dtype=None,
|
||||
diffusion_kv_cache_skip_step_indices=None,
|
||||
diffusion_kv_cache_skip_layer_indices=None,
|
||||
)
|
||||
|
||||
with set_current_diffusion_config(od_config):
|
||||
attention = Attention(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
causal=False,
|
||||
softmax_scale=1.0,
|
||||
paged_kv_cache_role="primary",
|
||||
)
|
||||
|
||||
assert attention.attn_backend is FlashAttentionBackend
|
||||
assert attention.backend_pref == "FLASH_ATTN"
|
||||
assert attention.backend_explicit is False
|
||||
assert attention.attn_spec is None
|
||||
|
||||
def test_user_explicit_flash_attn_still_requires_dense_fa4(self, monkeypatch):
|
||||
def _fake_get_attn_backend_for_role(*args, attention_config=None, **kwargs):
|
||||
if attention_config is not None:
|
||||
spec, _ = attention_config.resolve_with_source(role="self", role_category=None)
|
||||
if spec is not None and spec.backend.upper() == "FLASH_ATTN":
|
||||
raise ValueError(
|
||||
"FLASH_ATTN was explicitly selected but is unsupported "
|
||||
"(Blackwell requires CuTe FlashAttention-4 "
|
||||
"(flash_attn.cute / vllm-omni[fa4]); FA2/FA3 kernels are Hopper-only). "
|
||||
"Select a compatible backend."
|
||||
)
|
||||
raise AssertionError("explicit FLASH_ATTN must not fall through to the platform default")
|
||||
|
||||
monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
|
||||
monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")),
|
||||
diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
parallel_config=SimpleNamespace(ring_degree=1),
|
||||
diffusion_kv_cache_dtype=None,
|
||||
diffusion_kv_cache_skip_step_indices=None,
|
||||
diffusion_kv_cache_skip_layer_indices=None,
|
||||
)
|
||||
|
||||
with set_current_diffusion_config(od_config), pytest.raises(ValueError, match="explicitly selected"):
|
||||
Attention(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
causal=False,
|
||||
softmax_scale=1.0,
|
||||
paged_kv_cache_role="primary",
|
||||
)
|
||||
|
||||
def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
|
||||
class _FakeAttentionImpl:
|
||||
@@ -614,9 +763,10 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
|
||||
return _FakeBackend, AttentionSpec(backend="TRTLLM_ATTN", skip_softmax={"target_sparsity": 0.5})
|
||||
|
||||
class _FakeRingParallelAttention:
|
||||
def __init__(self, sp_group, attn_backend_pref=None):
|
||||
def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False):
|
||||
self.sp_group = sp_group
|
||||
self.attn_backend_pref = attn_backend_pref
|
||||
self.attn_backend_explicit = attn_backend_explicit
|
||||
|
||||
monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
|
||||
monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
|
||||
@@ -662,11 +812,134 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
|
||||
assert attn.use_ring is True
|
||||
assert attn.ring_runner is not None
|
||||
assert attn.ring_runner.attn_backend_pref == "TRTLLM_ATTN"
|
||||
assert attn.ring_runner.attn_backend_explicit is True
|
||||
|
||||
def test_attention_init_forwards_flashinfer_backend_without_dtype(self, monkeypatch):
|
||||
class _FakeAttentionImpl:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
class _FakeBackend:
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "FLASHINFER_ATTN"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls():
|
||||
return _FakeAttentionImpl
|
||||
|
||||
spec = AttentionSpec(backend="FLASHINFER_ATTN", quant={"flashinfer_backend": "fa2"})
|
||||
|
||||
monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", lambda **kwargs: (_FakeBackend, spec))
|
||||
monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
|
||||
monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
|
||||
monkeypatch.setattr(layer_mod, "get_sp_group", lambda: None)
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(default=spec),
|
||||
parallel_config=SimpleNamespace(ring_degree=1),
|
||||
diffusion_kv_cache_dtype=None,
|
||||
diffusion_kv_cache_skip_step_indices=None,
|
||||
diffusion_kv_cache_skip_layer_indices=None,
|
||||
)
|
||||
|
||||
with set_current_diffusion_config(od_config):
|
||||
attn = Attention(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
causal=False,
|
||||
softmax_scale=1.0,
|
||||
role="self",
|
||||
)
|
||||
|
||||
assert attn.backend_pref == "FLASHINFER_ATTN"
|
||||
assert attn.attention.kwargs["backend_kwargs"] == {"quant": {"flashinfer_backend": "fa2"}}
|
||||
assert attn.attention.kwargs["backend_explicit"] is True
|
||||
|
||||
def test_attention_init_propagates_ring_setup_failure(self, monkeypatch):
|
||||
class _FakeAttentionImpl:
|
||||
def __init__(self, **kwargs):
|
||||
pass
|
||||
|
||||
class _FakeBackend:
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "TORCH_SDPA"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls():
|
||||
return _FakeAttentionImpl
|
||||
|
||||
monkeypatch.setattr(
|
||||
layer_mod,
|
||||
"get_attn_backend_for_role",
|
||||
lambda **kwargs: (_FakeBackend, AttentionSpec(backend="TORCH_SDPA")),
|
||||
)
|
||||
monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
|
||||
monkeypatch.setattr(
|
||||
layer_mod,
|
||||
"get_sp_group",
|
||||
lambda: (_ for _ in ()).throw(RuntimeError("SP group is not initialized")),
|
||||
)
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
|
||||
parallel_config=SimpleNamespace(ring_degree=2),
|
||||
)
|
||||
with set_current_diffusion_config(od_config):
|
||||
with pytest.raises(RuntimeError, match="SP group is not initialized"):
|
||||
Attention(num_heads=4, head_size=64, causal=False, softmax_scale=1.0)
|
||||
|
||||
def test_float32_main_dispatch_does_not_override_selected_backend(self):
|
||||
sentinel = torch.randn(1, 2, 4, 8)
|
||||
selected_calls = []
|
||||
|
||||
def _selected_forward(*args):
|
||||
selected_calls.append(args)
|
||||
return sentinel
|
||||
|
||||
fake_attention = SimpleNamespace(
|
||||
attention=SimpleNamespace(
|
||||
forward=_selected_forward,
|
||||
),
|
||||
sdpa_fallback=SimpleNamespace(
|
||||
forward=lambda *args: pytest.fail("unexpected SDPA fallback"),
|
||||
),
|
||||
_assert_metadata_compatible=lambda metadata: None,
|
||||
_scheduler_paged_kv=False,
|
||||
paged_kv_cache_role=None,
|
||||
)
|
||||
query = torch.randn(1, 2, 4, 8, dtype=torch.float32)
|
||||
|
||||
output = Attention._run_local_attention(fake_attention, query, query, query, None)
|
||||
|
||||
assert output is sentinel
|
||||
assert len(selected_calls) == 1
|
||||
|
||||
def test_unsupported_attention_mask_is_rejected_before_backend_call(self):
|
||||
fake_backend = SimpleNamespace(
|
||||
get_name=lambda: "NO_MASK_BACKEND",
|
||||
supports_attention_mask=lambda *args, **kwargs: False,
|
||||
)
|
||||
fake_attention = SimpleNamespace(attn_backend=fake_backend)
|
||||
metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool))
|
||||
|
||||
with pytest.raises(ValueError, match="does not support attn_mask"):
|
||||
Attention._assert_metadata_compatible(fake_attention, metadata)
|
||||
|
||||
def test_ring_attention_rejects_mask_instead_of_ignoring_it(self):
|
||||
fake_attention = SimpleNamespace(attention=SimpleNamespace(skip=None), ring_runner=None)
|
||||
query = torch.randn(1, 2, 4, 8)
|
||||
metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool))
|
||||
|
||||
with pytest.raises(ValueError, match="Ring attention does not support attn_mask"):
|
||||
Attention._run_ring_attention(fake_attention, query, query, query, metadata)
|
||||
|
||||
|
||||
class TestDiffusionKvCacheQuantization:
|
||||
@staticmethod
|
||||
def _install_attention_init_stubs(monkeypatch):
|
||||
def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True):
|
||||
class _FakeAttentionImpl:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
@@ -674,6 +947,10 @@ class TestDiffusionKvCacheQuantization:
|
||||
def forward(self, query, key, value, attn_metadata=None):
|
||||
return query
|
||||
|
||||
@staticmethod
|
||||
def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
|
||||
return supports_kv_cache_dtype
|
||||
|
||||
class _FakeBackend:
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
@@ -683,14 +960,11 @@ class TestDiffusionKvCacheQuantization:
|
||||
def get_impl_cls():
|
||||
return _FakeAttentionImpl
|
||||
|
||||
@staticmethod
|
||||
def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
|
||||
return True
|
||||
|
||||
class _FakeRingParallelAttention:
|
||||
def __init__(self, sp_group, attn_backend_pref=None):
|
||||
def __init__(self, sp_group, attn_backend_pref=None, attn_backend_explicit=False):
|
||||
self.sp_group = sp_group
|
||||
self.attn_backend_pref = attn_backend_pref
|
||||
self.attn_backend_explicit = attn_backend_explicit
|
||||
|
||||
monkeypatch.setattr(
|
||||
layer_mod,
|
||||
@@ -764,3 +1038,22 @@ class TestDiffusionKvCacheQuantization:
|
||||
causal=False,
|
||||
softmax_scale=1.0,
|
||||
)
|
||||
|
||||
def test_explicit_unsupported_kv_cache_dtype_raises(self, monkeypatch):
|
||||
self._install_attention_init_stubs(monkeypatch, supports_kv_cache_dtype=False)
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(),
|
||||
parallel_config=SimpleNamespace(ring_degree=1),
|
||||
diffusion_kv_cache_dtype="fp8",
|
||||
diffusion_kv_cache_skip_step_indices=None,
|
||||
diffusion_kv_cache_skip_layer_indices=None,
|
||||
)
|
||||
|
||||
with set_current_diffusion_config(od_config):
|
||||
with pytest.raises(ValueError, match="does not support kv_cache_dtype='fp8'"):
|
||||
Attention(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
causal=False,
|
||||
softmax_scale=1.0,
|
||||
)
|
||||
@@ -0,0 +1,265 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import sys
|
||||
import types
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.attention.selector import _cached_get_backend_cls
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
from vllm_omni.diffusion.data import AttentionConfig
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
||||
|
||||
# Importing CudaOmniPlatform pulls in vllm.platforms.cuda, which imports the
|
||||
# CUDA-only ``vllm._C_stable_libtorch`` extension at module top level. That
|
||||
# extension is absent on non-CUDA builds (e.g. Intel/XPU), so skip the whole
|
||||
# module there before the import can crash collection. ``is_cuda()`` resolves
|
||||
# the platform without importing cuda.py.
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("CUDA-only diffusion backend selection tests", allow_module_level=True)
|
||||
|
||||
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform # noqa: E402
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _blackwell_sm120(monkeypatch: pytest.MonkeyPatch, *, cudnn_version: int = 90500) -> None:
|
||||
monkeypatch.setattr(
|
||||
CudaOmniPlatform,
|
||||
"get_device_capability",
|
||||
classmethod(lambda cls, device_id=0: DeviceCapability(12, 0)),
|
||||
)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
|
||||
monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: False))
|
||||
monkeypatch.setattr(torch.backends.cudnn, "version", lambda: cudnn_version)
|
||||
|
||||
|
||||
def _install_dummy_flashinfer(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Top-level ``flashinfer`` only — no ``prefill`` wrapper (partial install)."""
|
||||
monkeypatch.setitem(sys.modules, "flashinfer", types.ModuleType("flashinfer"))
|
||||
# A previously imported real submodule must not satisfy the wrapper probe.
|
||||
monkeypatch.delitem(sys.modules, "flashinfer.prefill", raising=False)
|
||||
|
||||
|
||||
def _install_dummy_flashinfer_prefill(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""FlashInfer stub that exposes the prefill wrapper FLASHINFER_ATTN needs."""
|
||||
prefill = types.ModuleType("flashinfer.prefill")
|
||||
setattr(prefill, "BatchPrefillWithRaggedKVCacheWrapper", type("BatchPrefillWithRaggedKVCacheWrapper", (), {}))
|
||||
pkg = types.ModuleType("flashinfer")
|
||||
setattr(pkg, "prefill", prefill)
|
||||
monkeypatch.setitem(sys.modules, "flashinfer", pkg)
|
||||
monkeypatch.setitem(sys.modules, "flashinfer.prefill", prefill)
|
||||
|
||||
|
||||
def test_auto_selects_cudnn_for_supported_blackwell_head_size(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=128)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
|
||||
|
||||
|
||||
def test_auto_routes_incompatible_head_size_to_sdpa(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
_install_dummy_flashinfer(monkeypatch)
|
||||
|
||||
# 320 is outside cuDNN FMHA (<=256, multiple of 8) and FlashInfer {64,128,256}.
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=320)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
|
||||
|
||||
|
||||
def test_auto_routes_non_multiple_head_size_to_sdpa_without_flashinfer(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
monkeypatch.setitem(sys.modules, "flashinfer", None)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=12)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
|
||||
|
||||
|
||||
def test_auto_selects_flashinfer_when_cudnn_too_old(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch, cudnn_version=90499)
|
||||
_install_dummy_flashinfer_prefill(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=128)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
|
||||
|
||||
|
||||
def test_auto_skips_flashinfer_when_only_toplevel_package(monkeypatch: pytest.MonkeyPatch):
|
||||
"""A stub that only imports ``flashinfer`` must not be auto-selected."""
|
||||
_blackwell_sm120(monkeypatch, cudnn_version=90499)
|
||||
_install_dummy_flashinfer(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=128)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
|
||||
|
||||
|
||||
def test_explicit_flashinfer_raises_without_prefill_wrapper(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
_install_dummy_flashinfer(monkeypatch)
|
||||
|
||||
with pytest.raises(ValueError, match="BatchPrefillWithRaggedKVCacheWrapper"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASHINFER_ATTN", head_size=128)
|
||||
|
||||
|
||||
def test_auto_skips_flashinfer_for_unsupported_head_size_when_cudnn_too_old(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch, cudnn_version=90499)
|
||||
_install_dummy_flashinfer(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=72)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
|
||||
|
||||
|
||||
def test_explicit_cudnn_raises_for_incompatible_head_size(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
|
||||
with pytest.raises(ValueError, match="head_size=12 is unsupported"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("CUDNN_ATTN", head_size=12)
|
||||
|
||||
|
||||
def test_explicit_cudnn_accepts_supported_head_size(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls("CUDNN_ATTN", head_size=72)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
|
||||
|
||||
|
||||
def test_explicit_cudnn_skips_head_size_check_for_capability_probe(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls("CUDNN_ATTN", head_size=-1)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_size", [32, 320])
|
||||
def test_explicit_flashinfer_raises_for_unsupported_head_size(monkeypatch: pytest.MonkeyPatch, head_size: int):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
_install_dummy_flashinfer_prefill(monkeypatch)
|
||||
|
||||
with pytest.raises(ValueError, match=f"head_size={head_size} is unsupported"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASHINFER_ATTN", head_size=head_size)
|
||||
|
||||
|
||||
def test_explicit_flashinfer_accepts_supported_head_size(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
_install_dummy_flashinfer_prefill(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASHINFER_ATTN", head_size=128)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
|
||||
|
||||
|
||||
def test_explicit_flashinfer_skips_head_size_check_for_capability_probe(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
_install_dummy_flashinfer_prefill(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASHINFER_ATTN", head_size=-1)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.FLASHINFER_ATTN.get_path()
|
||||
|
||||
|
||||
def test_auto_selects_cudnn_for_unknown_head_size_probe(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
_install_dummy_flashinfer_prefill(monkeypatch)
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=-1)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.CUDNN_ATTN.get_path()
|
||||
|
||||
|
||||
def test_explicit_flash_attn_rejects_blackwell_without_fa4(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
|
||||
|
||||
with pytest.raises(ValueError, match="requires CuTe FlashAttention-4"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
|
||||
|
||||
|
||||
def test_explicit_flash_attn_accepts_blackwell_when_fa4_available(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: True))
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
|
||||
|
||||
|
||||
def test_blackwell_dense_flash_capability_tracks_fa4(monkeypatch: pytest.MonkeyPatch):
|
||||
_blackwell_sm120(monkeypatch)
|
||||
|
||||
assert CudaOmniPlatform.supports_diffusion_dense_flash_attention() is False
|
||||
|
||||
monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: True))
|
||||
|
||||
assert CudaOmniPlatform.supports_diffusion_dense_flash_attention() is True
|
||||
|
||||
|
||||
def test_hopper_dense_flash_capability_tracks_fa_package(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(
|
||||
CudaOmniPlatform,
|
||||
"get_device_capability",
|
||||
classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)),
|
||||
)
|
||||
monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_package", classmethod(lambda cls: False))
|
||||
|
||||
assert CudaOmniPlatform.supports_diffusion_dense_flash_attention() is False
|
||||
|
||||
monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_package", classmethod(lambda cls: True))
|
||||
|
||||
assert CudaOmniPlatform.supports_diffusion_dense_flash_attention() is True
|
||||
|
||||
|
||||
def test_auto_selects_flash_attn_on_blackwell_when_cudnn_old_and_fa4_available(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
_blackwell_sm120(monkeypatch, cudnn_version=90499)
|
||||
monkeypatch.setitem(sys.modules, "flashinfer", None)
|
||||
monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: True))
|
||||
|
||||
path = CudaOmniPlatform.get_diffusion_attn_backend_cls(None, head_size=128)
|
||||
|
||||
assert path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
|
||||
|
||||
|
||||
def test_marked_paged_layer_constructs_on_blackwell_cudnn_without_fa4(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Scheduler-paged remap must not require dense FA4 when CUDNN is the default."""
|
||||
_blackwell_sm120(monkeypatch)
|
||||
_cached_get_backend_cls.cache_clear()
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(),
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
parallel_config=SimpleNamespace(ring_degree=1, allgather_degree=1),
|
||||
diffusion_kv_cache_dtype=None,
|
||||
diffusion_kv_cache_skip_step_indices=None,
|
||||
diffusion_kv_cache_skip_layer_indices=None,
|
||||
model_class_name=None,
|
||||
)
|
||||
|
||||
with set_current_diffusion_config(od_config):
|
||||
attention = Attention(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
causal=False,
|
||||
softmax_scale=1.0,
|
||||
paged_kv_cache_role="primary",
|
||||
)
|
||||
|
||||
assert attention.attn_backend.get_name() == "FLASH_ATTN"
|
||||
assert attention.backend_pref == "FLASH_ATTN"
|
||||
assert attention.backend_explicit is False
|
||||
assert attention.attn_spec is None
|
||||
@@ -1,62 +1,84 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from contextlib import nullcontext
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.nn.attention import SDPBackend
|
||||
|
||||
import vllm_omni.diffusion.attention.backends.cudnn_attn as cudnn_backend
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
|
||||
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionBackend, CuDNNAttentionImpl
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("query_sequence_length", "key_sequence_length", "expected_backend"),
|
||||
[
|
||||
(1, 1, SDPBackend.MATH),
|
||||
(8, 1, SDPBackend.MATH),
|
||||
(1, 8, SDPBackend.MATH),
|
||||
(8, 8, SDPBackend.CUDNN_ATTENTION),
|
||||
],
|
||||
)
|
||||
def test_cudnn_routes_single_token_attention_to_math(
|
||||
monkeypatch, query_sequence_length, key_sequence_length, expected_backend
|
||||
):
|
||||
def test_cudnn_backend_uses_math_for_kv_seq_len_one(monkeypatch):
|
||||
"""Automatic CUDNN_ATTN (platform default) may use MATH for singleton K/V."""
|
||||
selected_backends = []
|
||||
|
||||
@contextmanager
|
||||
def fake_sdpa_kernel(backends):
|
||||
selected_backends.append(backends)
|
||||
return nullcontext()
|
||||
selected_backends.append(tuple(backends))
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
|
||||
fake_sdpa_kernel,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
torch.nn.functional,
|
||||
"scaled_dot_product_attention",
|
||||
lambda query, _key, _value, **_kwargs: query,
|
||||
)
|
||||
impl = CuDNNAttentionImpl(
|
||||
num_heads=2,
|
||||
head_size=4,
|
||||
softmax_scale=0.5,
|
||||
)
|
||||
query = torch.randn(1, query_sequence_length, 2, 4)
|
||||
key = torch.randn(1, key_sequence_length, 2, 4)
|
||||
value = torch.randn_like(key)
|
||||
def fake_sdpa(query, key, value, **kwargs):
|
||||
return query
|
||||
|
||||
output = impl.forward_cuda(query, key, value)
|
||||
monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel)
|
||||
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
|
||||
|
||||
impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
|
||||
query = torch.randn(1, 2, 2, 8)
|
||||
singleton_kv = torch.randn(1, 1, 2, 8)
|
||||
|
||||
output = impl.forward_cuda(query, singleton_kv, singleton_kv)
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert selected_backends == [[expected_backend]]
|
||||
assert selected_backends == [(SDPBackend.MATH,)]
|
||||
|
||||
|
||||
def test_explicit_cudnn_rejects_kv_seq_len_one():
|
||||
impl = CuDNNAttentionImpl(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=0.5,
|
||||
backend_explicit=True,
|
||||
)
|
||||
query = torch.randn(1, 2, 2, 8)
|
||||
singleton_kv = torch.randn(1, 1, 2, 8)
|
||||
|
||||
with pytest.raises(ValueError, match="explicitly selected.*sequence length 1"):
|
||||
impl.forward_cuda(query, singleton_kv, singleton_kv)
|
||||
|
||||
|
||||
def test_cudnn_backend_pins_cudnn_only_when_kv_seq_len_gt_one(monkeypatch):
|
||||
selected_backends = []
|
||||
|
||||
@contextmanager
|
||||
def fake_sdpa_kernel(backends):
|
||||
selected_backends.append(tuple(backends))
|
||||
yield
|
||||
|
||||
def reject_shape(*args, **kwargs):
|
||||
raise RuntimeError("No available kernel. Aborting execution.")
|
||||
|
||||
monkeypatch.setattr(cudnn_backend, "sdpa_kernel", fake_sdpa_kernel)
|
||||
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", reject_shape)
|
||||
|
||||
impl = CuDNNAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5)
|
||||
tensors = torch.randn(1, 2, 2, 8)
|
||||
|
||||
with pytest.raises(RuntimeError, match="No available kernel"):
|
||||
impl.forward_cuda(tensors, tensors, tensors)
|
||||
|
||||
assert selected_backends == [(SDPBackend.CUDNN_ATTENTION,)]
|
||||
|
||||
|
||||
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
|
||||
observed = {}
|
||||
observed: dict[str, Any] = {}
|
||||
|
||||
def fake_sdpa(query, key, value, **kwargs):
|
||||
observed.update(query=query, key=key, value=value, kwargs=kwargs)
|
||||
@@ -105,3 +127,15 @@ def test_cudnn_rejects_invalid_valid_kv_length():
|
||||
query,
|
||||
AttentionMetadata(extra={"valid_kv_length": 9}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_size", [8, 64, 128, 256])
|
||||
def test_cudnn_backend_accepts_blackwell_fmha_head_sizes(head_size):
|
||||
assert CuDNNAttentionBackend.supports_head_size(head_size)
|
||||
assert head_size in CuDNNAttentionBackend.get_supported_head_sizes()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("head_size", [0, 7, 12, 320])
|
||||
def test_cudnn_backend_rejects_incompatible_head_sizes(head_size):
|
||||
assert not CuDNNAttentionBackend.supports_head_size(head_size)
|
||||
assert head_size not in CuDNNAttentionBackend.get_supported_head_sizes()
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import math
|
||||
import sys
|
||||
import types
|
||||
from typing import Any
|
||||
@@ -8,9 +9,17 @@ from typing import Any
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
AttentionMetadata,
|
||||
PackedPaddingMetadata,
|
||||
VideoTokenLayout,
|
||||
VideoTokenSpan,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
|
||||
FastVideoVSABackend,
|
||||
FastVideoVSAImpl,
|
||||
_build_h3_block_map,
|
||||
_get_h3_tile_metadata,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.backends.registry import (
|
||||
DiffusionAttentionBackendEnum,
|
||||
@@ -23,6 +32,12 @@ def test_fastvideo_vsa_backend_is_registered():
|
||||
assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
|
||||
|
||||
|
||||
def test_fastvideo_vsa_reports_missing_optional_kernel(monkeypatch):
|
||||
monkeypatch.setattr("importlib.util.find_spec", lambda name: None)
|
||||
with pytest.raises(ImportError, match="fastvideo-kernel"):
|
||||
FastVideoVSABackend.validate_available()
|
||||
|
||||
|
||||
def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
fake_module = types.ModuleType("fastvideo_kernel")
|
||||
@@ -201,3 +216,190 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
|
||||
# CPU/float32 is rejected later, but k=N itself must not trigger fallback.
|
||||
assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
|
||||
assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
|
||||
|
||||
|
||||
def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
|
||||
extra: dict[str, Any] = {
|
||||
"vsa_h3_prefix_segments": prefix_segments,
|
||||
}
|
||||
if gate is not None:
|
||||
extra["gate_compress"] = gate
|
||||
return AttentionMetadata(
|
||||
packed_padding=packed_padding,
|
||||
video_layout=VideoTokenLayout(
|
||||
used_len=sum(prefix_segments) + math.prod(video_shape),
|
||||
video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
|
||||
),
|
||||
extra=extra,
|
||||
)
|
||||
|
||||
|
||||
def _fake_block_sparse_kernel(monkeypatch, calls):
|
||||
"""Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
|
||||
|
||||
def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
|
||||
calls["block_map"] = block_map.clone()
|
||||
calls["variable_block_sizes"] = variable_block_sizes.tolist()
|
||||
calls["q_shape"] = tuple(q.shape)
|
||||
blocks = block_map.shape[-1]
|
||||
assert blocks * 64 == q.shape[2]
|
||||
# Expand the per-block map to a per-token mask so the stand-in honours
|
||||
# the sparsity contract the real kernel implements, and drop the rows
|
||||
# each partial block pads with, which the real kernel skips through
|
||||
# variable_block_sizes.
|
||||
token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
|
||||
within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
|
||||
token_mask = token_mask & within.reshape(1, 1, 1, -1)
|
||||
out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
|
||||
return out.to(q.dtype), None
|
||||
|
||||
module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
|
||||
setattr(module, "block_sparse_attn", block_sparse_attn)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
|
||||
|
||||
|
||||
def _h3_impl(**backend_kwargs):
|
||||
kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
|
||||
kwargs.update(backend_kwargs)
|
||||
return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
|
||||
|
||||
|
||||
def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
_fake_block_sparse_kernel(monkeypatch, calls)
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5, 70), (2, 4, 4)
|
||||
seq_len = sum(prefix_segments) + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
|
||||
assert output.shape == query.shape
|
||||
# topk covers every video block, so the block map is dense and the tiled
|
||||
# kernel must reproduce plain attention over the same rows.
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
*(x.float().transpose(1, 2) for x in (query, query, query))
|
||||
).transpose(1, 2)
|
||||
torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
|
||||
# 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
|
||||
assert calls["variable_block_sizes"] == [5, 64, 6, 32]
|
||||
|
||||
|
||||
def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
_fake_block_sparse_kernel(monkeypatch, calls)
|
||||
impl = _h3_impl(topk=1)
|
||||
prefix_segments, video_shape = (5,), (8, 4, 4)
|
||||
seq_len = 5 + 8 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
|
||||
# 1 prefix block + 2 video blocks. The odd count is padded to an even one
|
||||
# for the paired-CTA kernel contract, and the Triton route slices that
|
||||
# transport-only partner off again before launching.
|
||||
block_map = calls["block_map"]
|
||||
assert block_map.shape[-1] == 3
|
||||
assert calls["variable_block_sizes"] == [5, 64, 64]
|
||||
assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
|
||||
assert block_map[..., :1].all(), "prefix keys are exempt from selection"
|
||||
assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
|
||||
|
||||
# The same odd geometry must still untile back onto the original rows.
|
||||
dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
*(x.float().transpose(1, 2) for x in (query, query, query))
|
||||
).transpose(1, 2)
|
||||
torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
|
||||
_fake_block_sparse_kernel(monkeypatch, {})
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5,), (2, 4, 4)
|
||||
seq_len = 5 + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
zero_gate = torch.zeros_like(query)
|
||||
with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
|
||||
live_gate = torch.full_like(query, 0.5)
|
||||
with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
|
||||
|
||||
torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
|
||||
assert not torch.allclose(with_live_gate.float(), without_gate.float())
|
||||
|
||||
|
||||
def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
|
||||
_fake_block_sparse_kernel(monkeypatch, {})
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5,), (2, 4, 4)
|
||||
valid = 5 + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
|
||||
padding = PackedPaddingMetadata(
|
||||
q_length=valid,
|
||||
kv_length=valid,
|
||||
cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
|
||||
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
|
||||
)
|
||||
|
||||
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert torch.count_nonzero(output[:, valid:]) == 0
|
||||
|
||||
|
||||
def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
|
||||
# The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
|
||||
# building the padding mask. Every fallback out of the VSA path must honour
|
||||
# that contract itself; SDPA reads attn_mask and nothing else.
|
||||
impl = FastVideoVSAImpl(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
backend_kwargs={"topk": 1, "min_seq_len": 4096},
|
||||
)
|
||||
valid = 40
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, 64, 2, 8)
|
||||
padding = PackedPaddingMetadata(
|
||||
q_length=valid,
|
||||
kv_length=valid,
|
||||
cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
|
||||
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
|
||||
)
|
||||
|
||||
baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
|
||||
perturbed_input = query.clone()
|
||||
perturbed_input[:, valid:] += 100.0
|
||||
perturbed = impl.forward_cuda(
|
||||
perturbed_input, perturbed_input, perturbed_input, AttentionMetadata(packed_padding=padding, extra={})
|
||||
)
|
||||
|
||||
assert baseline.shape == query.shape
|
||||
torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
|
||||
assert torch.count_nonzero(baseline[:, valid:]) == 0
|
||||
|
||||
|
||||
def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
|
||||
partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
|
||||
(5, 70, 9), (5, 6, 6), torch.device("cpu")
|
||||
)
|
||||
assert sizes[:4].tolist() == [5, 64, 6, 9]
|
||||
assert prefix_blocks == 4
|
||||
assert video_blocks == 8
|
||||
assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
|
||||
assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
|
||||
|
||||
|
||||
def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
|
||||
scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
|
||||
block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
|
||||
assert block_map[:, :, :2].all()
|
||||
assert block_map[..., :2].all()
|
||||
assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()
|
||||
@@ -0,0 +1,112 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends import flashinfer_attn
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.flashinfer_attn import (
|
||||
FlashInferAttentionBackend,
|
||||
FlashInferAttentionImpl,
|
||||
)
|
||||
from vllm_omni.diffusion.data import AttentionSpec, AttnQuantSpec
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _impl(*, causal: bool = False, backend_explicit: bool = False):
|
||||
# Avoid CUDA/wrapper init; these tests only cover mask validation.
|
||||
obj = FlashInferAttentionImpl.__new__(FlashInferAttentionImpl)
|
||||
obj.causal = causal
|
||||
obj.softmax_scale = 0.5
|
||||
obj.flashinfer_backend = "fa2"
|
||||
obj.backend_explicit = backend_explicit
|
||||
obj._sdpa_fallback = None
|
||||
return obj
|
||||
|
||||
|
||||
def test_flashinfer_rejects_float_mask_instead_of_falling_back(monkeypatch):
|
||||
monkeypatch.setattr(flashinfer_attn, "HAS_FLASHINFER", True)
|
||||
query = torch.randn(1, 2, 2, 8)
|
||||
metadata = AttentionMetadata(attn_mask=torch.zeros(2, 2))
|
||||
|
||||
with pytest.raises(ValueError, match="boolean-only"):
|
||||
_impl(backend_explicit=True).forward_cuda(query, query, query, metadata)
|
||||
|
||||
|
||||
def test_flashinfer_rejects_causal_custom_mask_instead_of_falling_back(monkeypatch):
|
||||
monkeypatch.setattr(flashinfer_attn, "HAS_FLASHINFER", True)
|
||||
query = torch.randn(1, 2, 2, 8)
|
||||
metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False], [True, True]]))
|
||||
|
||||
with pytest.raises(ValueError, match="causal=True"):
|
||||
_impl(causal=True, backend_explicit=True).forward_cuda(query, query, query, metadata)
|
||||
|
||||
|
||||
def test_explicit_cute_dsl_rejects_custom_mask_instead_of_falling_back(monkeypatch):
|
||||
monkeypatch.setattr(flashinfer_attn, "HAS_FLASHINFER", True)
|
||||
impl = _impl(backend_explicit=True)
|
||||
impl.flashinfer_backend = "cute-dsl"
|
||||
query = torch.randn(1, 2, 2, 8)
|
||||
metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False], [True, True]]))
|
||||
|
||||
with pytest.raises(ValueError, match="cute-dsl"):
|
||||
impl.forward_cuda(query, query, query, metadata)
|
||||
|
||||
|
||||
def test_auto_cute_dsl_falls_back_to_sdpa_for_custom_mask(monkeypatch):
|
||||
monkeypatch.setattr(flashinfer_attn, "HAS_FLASHINFER", True)
|
||||
impl = _impl(backend_explicit=False)
|
||||
impl.flashinfer_backend = "cute-dsl"
|
||||
impl._run_batch_prefill = lambda *args, **kwargs: pytest.fail("unexpected cute-dsl prefill")
|
||||
query = torch.randn(1, 2, 2, 8)
|
||||
metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False], [True, True]]))
|
||||
|
||||
output = impl.forward_cuda(query, query, query, metadata)
|
||||
|
||||
assert output.shape == query.shape
|
||||
|
||||
|
||||
def test_explicit_cute_dsl_is_not_mask_capable(monkeypatch):
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
|
||||
spec = AttentionSpec(backend="FLASHINFER_ATTN")
|
||||
|
||||
assert FlashInferAttentionBackend.supports_attention_mask() is True
|
||||
assert FlashInferAttentionBackend.supports_attention_mask(spec) is False
|
||||
|
||||
|
||||
def test_explicit_fa2_flashinfer_is_mask_capable(monkeypatch):
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
|
||||
spec = AttentionSpec(backend="FLASHINFER_ATTN", quant=AttnQuantSpec(flashinfer_backend="fa2"))
|
||||
|
||||
assert FlashInferAttentionBackend.supports_attention_mask(spec) is True
|
||||
|
||||
|
||||
def test_explicit_fa2_backend_kwargs_select_fa2_on_blackwell(monkeypatch):
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
|
||||
spec = AttentionSpec(backend="FLASHINFER_ATTN", quant=AttnQuantSpec(flashinfer_backend="fa2"))
|
||||
kwargs = spec.backend_kwargs() or {}
|
||||
requested = (kwargs.get("quant") or {}).get("flashinfer_backend", "auto")
|
||||
|
||||
assert kwargs == {"quant": {"flashinfer_backend": "fa2"}}
|
||||
assert FlashInferAttentionImpl._select_backend(requested) == "fa2"
|
||||
assert FlashInferAttentionBackend.supports_attention_mask(spec) is True
|
||||
|
||||
|
||||
def test_auto_flashinfer_backend_kwargs_select_cute_dsl_on_blackwell(monkeypatch):
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
|
||||
spec = AttentionSpec(backend="FLASHINFER_ATTN")
|
||||
kwargs = spec.backend_kwargs() or {}
|
||||
requested = (kwargs.get("quant") or {}).get("flashinfer_backend", "auto")
|
||||
|
||||
assert "quant" not in kwargs
|
||||
assert FlashInferAttentionImpl._select_backend(requested) == "cute-dsl"
|
||||
assert FlashInferAttentionBackend.supports_attention_mask(spec) is False
|
||||
|
||||
|
||||
def test_hopper_explicit_flashinfer_is_mask_capable(monkeypatch):
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (9, 0))
|
||||
spec = AttentionSpec(backend="FLASHINFER_ATTN")
|
||||
|
||||
assert FlashInferAttentionBackend.supports_attention_mask(spec) is True
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import sys
|
||||
import types
|
||||
@@ -9,18 +9,13 @@ import torch
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
|
||||
from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test that when kernels package is not available, the platform fallback logic
|
||||
|
||||
routes to local FLASH_ATTN.
|
||||
"""
|
||||
def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
|
||||
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
||||
|
||||
# Temporarily hide/remove kernels module if present
|
||||
@@ -38,19 +33,73 @@ def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch):
|
||||
)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
|
||||
|
||||
# Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable
|
||||
backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
|
||||
assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
|
||||
with pytest.raises(ImportError, match="explicitly selected"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
|
||||
with pytest.raises(ImportError, match="explicitly selected"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
||||
|
||||
# Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable
|
||||
backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
||||
assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
|
||||
|
||||
# Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs
|
||||
kernels_module = types.ModuleType("kernels")
|
||||
monkeypatch.setitem(sys.modules, "kernels", kernels_module)
|
||||
backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
||||
assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path()
|
||||
with pytest.raises(ValueError, match="require.*Hopper GPU"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
||||
|
||||
monkeypatch.setattr(
|
||||
CudaOmniPlatform,
|
||||
"get_device_capability",
|
||||
classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
|
||||
)
|
||||
with pytest.raises(ValueError, match="current FA2 kernels require"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64)
|
||||
with pytest.raises(ValueError, match="require.*Hopper GPU"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_explicit_flash_attention_unavailable_does_not_fallback(monkeypatch: pytest.MonkeyPatch):
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
||||
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
||||
|
||||
monkeypatch.setattr(
|
||||
CudaOmniPlatform,
|
||||
"get_device_capability",
|
||||
classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)),
|
||||
)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
|
||||
|
||||
with pytest.raises(ValueError, match="explicitly selected but is unsupported"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
|
||||
|
||||
monkeypatch.setattr(
|
||||
CudaOmniPlatform,
|
||||
"get_device_capability",
|
||||
classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
|
||||
)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True})
|
||||
monkeypatch.setattr(CudaOmniPlatform, "has_flash_attn_4", classmethod(lambda cls: False))
|
||||
with pytest.raises(ValueError, match="requires CuTe FlashAttention-4"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_explicit_sage_attention_unsupported_arch_does_not_reach_kernel(monkeypatch: pytest.MonkeyPatch):
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
||||
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
||||
|
||||
monkeypatch.setattr(
|
||||
CudaOmniPlatform,
|
||||
"get_device_capability",
|
||||
classmethod(lambda cls, device_id=0: DeviceCapability(10, 3)),
|
||||
)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
|
||||
|
||||
with pytest.raises(ValueError, match="does not provide a kernel for sm_103"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN", head_size=64)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
|
||||
@@ -0,0 +1,199 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.ring import ring_globals, ring_kernels
|
||||
from vllm_omni.diffusion.attention.parallel import factory, ring
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_source_build_fa3_is_hopper_only(monkeypatch):
|
||||
fake_interface = SimpleNamespace(
|
||||
_flash_attn_forward=lambda *args, **kwargs: None,
|
||||
flash_attn_func=lambda *args, **kwargs: None,
|
||||
)
|
||||
try:
|
||||
with monkeypatch.context() as context:
|
||||
context.setitem(sys.modules, "flash_attn_interface", fake_interface)
|
||||
reloaded = importlib.reload(ring_globals)
|
||||
assert reloaded.HAS_FA3
|
||||
assert reloaded.FA3_SUPPORTED_CUDA_MAJORS == frozenset({9})
|
||||
finally:
|
||||
importlib.reload(ring_globals)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("capability", "expected"),
|
||||
[
|
||||
((8, 0), True),
|
||||
((8, 9), True),
|
||||
((9, 0), True),
|
||||
((10, 0), False),
|
||||
((10, 3), False),
|
||||
((12, 0), False),
|
||||
],
|
||||
)
|
||||
def test_can_use_fa3_checks_device_arch(monkeypatch, capability, expected):
|
||||
monkeypatch.setattr(ring, "HAS_FA3", True)
|
||||
monkeypatch.setattr(ring, "FA3_SUPPORTED_CUDA_MAJORS", frozenset({8, 9}))
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: capability)
|
||||
|
||||
assert ring._can_use_fa3(torch.device("cuda")) is expected
|
||||
|
||||
|
||||
def test_can_use_fa3_preserves_unknown_extension_contract(monkeypatch):
|
||||
monkeypatch.setattr(ring, "HAS_FA3", True)
|
||||
monkeypatch.setattr(ring, "FA3_SUPPORTED_CUDA_MAJORS", None)
|
||||
|
||||
assert ring._can_use_fa3(torch.device("cuda"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("capability", "expected"),
|
||||
[
|
||||
((8, 0), True),
|
||||
((9, 0), True),
|
||||
((10, 0), False),
|
||||
((10, 3), False),
|
||||
((12, 0), False),
|
||||
],
|
||||
)
|
||||
def test_can_use_fa2_checks_device_arch(monkeypatch, capability, expected):
|
||||
monkeypatch.setattr(ring, "HAS_FLASH_ATTN", True)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: capability)
|
||||
|
||||
assert ring._can_use_fa2(torch.device("cuda")) is expected
|
||||
|
||||
|
||||
def test_ring_falls_back_to_sdpa_when_fa3_extension_has_no_device_kernel(monkeypatch):
|
||||
monkeypatch.setattr(ring, "HAS_FA3", True)
|
||||
monkeypatch.setattr(ring, "HAS_FA4", False)
|
||||
monkeypatch.setattr(ring, "HAS_FLASH_ATTN", False)
|
||||
monkeypatch.setattr(ring, "HAS_AITER", False)
|
||||
monkeypatch.setattr(ring, "_can_use_fa3", lambda _device: False)
|
||||
|
||||
expected = torch.randn(1, 2, 1, 8)
|
||||
|
||||
def fake_ring_sdpa(*args, **kwargs):
|
||||
return expected
|
||||
|
||||
module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
|
||||
monkeypatch.setitem(sys.modules, module_name, SimpleNamespace(ring_pytorch_attn_func=fake_ring_sdpa))
|
||||
|
||||
strategy = ring.RingParallelAttention(SimpleNamespace(ring_group=object()))
|
||||
query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
|
||||
actual = strategy.run_attention(query, query, query, attn_metadata=None)
|
||||
|
||||
assert actual is expected
|
||||
|
||||
|
||||
def test_explicit_ring_backend_does_not_fallback_to_sdpa(monkeypatch):
|
||||
monkeypatch.setattr(ring, "HAS_FA3", True)
|
||||
monkeypatch.setattr(ring, "HAS_FA4", False)
|
||||
monkeypatch.setattr(ring, "HAS_FLASH_ATTN", False)
|
||||
monkeypatch.setattr(ring, "HAS_AITER", False)
|
||||
monkeypatch.setattr(ring, "_can_use_fa3", lambda _device: False)
|
||||
|
||||
strategy = ring.RingParallelAttention(
|
||||
SimpleNamespace(ring_group=object()),
|
||||
attn_backend_pref="FLASH_ATTN",
|
||||
attn_backend_explicit=True,
|
||||
)
|
||||
query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
|
||||
|
||||
with pytest.raises(RuntimeError, match="explicitly selected"):
|
||||
strategy.run_attention(query, query, query, attn_metadata=None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"backend_pref",
|
||||
[
|
||||
"CUDNN_ATTN",
|
||||
"FLASH_ATTN_HUB",
|
||||
"FLASH_ATTN_3_HUB",
|
||||
],
|
||||
)
|
||||
def test_explicit_non_ring_backend_is_rejected(monkeypatch, backend_pref):
|
||||
# Local FA is available so a missing rejection would silently run FA4/FA3/FA2
|
||||
# instead of the requested Hub/cuDNN backend.
|
||||
monkeypatch.setattr(ring, "HAS_FA4", True)
|
||||
monkeypatch.setattr(ring, "_can_use_fa4", lambda _device: True)
|
||||
module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
module_name,
|
||||
SimpleNamespace(ring_pytorch_attn_func=lambda *args, **kwargs: pytest.fail("unexpected SDPA fallback")),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"vllm_omni.diffusion.attention.backends.ring_flash_attn",
|
||||
SimpleNamespace(ring_flash_attn_func=lambda *args, **kwargs: pytest.fail("unexpected local FA ring")),
|
||||
)
|
||||
strategy = ring.RingParallelAttention(
|
||||
SimpleNamespace(ring_group=object()),
|
||||
attn_backend_pref=backend_pref,
|
||||
attn_backend_explicit=True,
|
||||
)
|
||||
query = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16)
|
||||
|
||||
with pytest.raises(ValueError, match="ring sequence parallelism has no implementation"):
|
||||
strategy.run_attention(query, query, query, attn_metadata=None)
|
||||
|
||||
|
||||
def test_explicit_flash_ring_rejects_float32_instead_of_falling_back(monkeypatch):
|
||||
module_name = "vllm_omni.diffusion.attention.backends.ring_pytorch_attn"
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
module_name,
|
||||
SimpleNamespace(ring_pytorch_attn_func=lambda *args, **kwargs: pytest.fail("unexpected SDPA fallback")),
|
||||
)
|
||||
strategy = ring.RingParallelAttention(
|
||||
SimpleNamespace(ring_group=object()),
|
||||
attn_backend_pref="FLASH_ATTN",
|
||||
attn_backend_explicit=True,
|
||||
)
|
||||
query = torch.randn(1, 2, 1, 8, dtype=torch.float32)
|
||||
|
||||
with pytest.raises(ValueError, match="does not support float32"):
|
||||
strategy.run_attention(query, query, query, attn_metadata=None)
|
||||
|
||||
|
||||
def test_pytorch_ring_flash_op_rejects_float32():
|
||||
query = torch.randn(1, 2, 1, 8, dtype=torch.float32)
|
||||
|
||||
with pytest.raises(ValueError, match="does not support float32"):
|
||||
ring_kernels.pytorch_attn_forward(query, query, query, op_type="flash")
|
||||
|
||||
|
||||
def test_configured_sp_does_not_fallback_when_group_is_unavailable(monkeypatch):
|
||||
parallel_config = SimpleNamespace(ulysses_degree=2, ring_degree=1, allgather_degree=1)
|
||||
forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(parallel_config=parallel_config))
|
||||
monkeypatch.setattr(factory, "is_forward_context_available", lambda: True)
|
||||
monkeypatch.setattr(factory, "get_forward_context", lambda: forward_context)
|
||||
monkeypatch.setattr(
|
||||
factory,
|
||||
"get_sp_group",
|
||||
lambda: (_ for _ in ()).throw(RuntimeError("SP group is not initialized")),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="SP is configured"):
|
||||
factory.build_parallel_attention_strategy(scatter_idx=2, gather_idx=1, use_sync=False)
|
||||
|
||||
|
||||
def test_configured_sp_rejects_single_rank_group(monkeypatch):
|
||||
parallel_config = SimpleNamespace(ulysses_degree=2, ring_degree=1, allgather_degree=1)
|
||||
forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(parallel_config=parallel_config))
|
||||
monkeypatch.setattr(factory, "is_forward_context_available", lambda: True)
|
||||
monkeypatch.setattr(factory, "get_forward_context", lambda: forward_context)
|
||||
monkeypatch.setattr(factory, "get_sp_group", lambda: SimpleNamespace())
|
||||
monkeypatch.setattr(factory, "get_sequence_parallel_world_size", lambda: 1)
|
||||
|
||||
with pytest.raises(RuntimeError, match="world size is not greater than one"):
|
||||
factory.build_parallel_attention_strategy(scatter_idx=2, gather_idx=1, use_sync=False)
|
||||
@@ -72,3 +72,25 @@ def test_ring_prefers_fa4_on_blackwell(monkeypatch):
|
||||
|
||||
assert strategy.run_attention(query, "key", "value", None) == "output"
|
||||
assert captured["attn_type"] is AttnType.FA4
|
||||
|
||||
|
||||
def test_flashinfer_wrapper_canonicalizes_unbatched_lse_to_bhs(monkeypatch):
|
||||
seq_len = num_heads = 4
|
||||
query = torch.randn(1, seq_len, num_heads, 8)
|
||||
unbatched_out = torch.randn(seq_len, num_heads, 8)
|
||||
unbatched_lse = torch.arange(num_heads * seq_len, dtype=torch.float32).reshape(num_heads, seq_len)
|
||||
assert not torch.equal(unbatched_lse, unbatched_lse.transpose(0, 1))
|
||||
|
||||
def fake_prefill(_q, _k, _v, **_kwargs):
|
||||
return unbatched_out.clone(), unbatched_lse.clone()
|
||||
|
||||
monkeypatch.setattr(ring_kernels, "HAS_FLASHINFER", True)
|
||||
# FlashInfer symbols are only bound when HAS_FLASHINFER was true at import.
|
||||
monkeypatch.setattr(ring_kernels, "single_prefill_with_kv_cache", fake_prefill, raising=False)
|
||||
monkeypatch.setattr(ring_kernels, "_LOG2_E", 1.0, raising=False)
|
||||
|
||||
out, lse = ring_kernels.flashinfer_attn_forward(query, query, query, softmax_scale=0.1)
|
||||
|
||||
assert out.shape == query.shape
|
||||
assert lse.shape == (1, num_heads, seq_len)
|
||||
assert torch.equal(lse[0], unbatched_lse)
|
||||
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.ring.ring_utils import update_out_and_lse
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("layout", ["bhs", "bsh", "bhs_padded", "bhs1"])
|
||||
def test_ring_lse_known_layouts_are_normalized(layout):
|
||||
block_out = torch.randn(2, 5, 3, 4)
|
||||
canonical = torch.randn(2, 5, 3)
|
||||
if layout == "bhs":
|
||||
block_lse = canonical.transpose(1, 2)
|
||||
lse_layout = "bhs"
|
||||
elif layout == "bsh":
|
||||
block_lse = canonical
|
||||
lse_layout = "bsh"
|
||||
elif layout == "bhs_padded":
|
||||
block_lse = torch.cat([canonical.transpose(1, 2), torch.randn(2, 3, 2)], dim=2)
|
||||
lse_layout = "bhs"
|
||||
else:
|
||||
block_lse = canonical.transpose(1, 2).unsqueeze(-1)
|
||||
lse_layout = "bhs"
|
||||
|
||||
out, lse = update_out_and_lse(None, None, block_out, block_lse, lse_layout=lse_layout)
|
||||
|
||||
assert out.shape == block_out.shape
|
||||
assert out.dtype == torch.float32
|
||||
assert lse.shape == (2, 5, 3, 1)
|
||||
assert torch.equal(lse.squeeze(-1), canonical)
|
||||
|
||||
|
||||
def test_ring_lse_seq_len_equals_num_heads_keeps_asymmetric_values():
|
||||
batch, seq_len, num_heads, head_dim = 2, 4, 4, 8
|
||||
block_out = torch.randn(batch, seq_len, num_heads, head_dim)
|
||||
canonical = torch.arange(batch * seq_len * num_heads, dtype=torch.float32).reshape(batch, seq_len, num_heads)
|
||||
bhs = canonical.transpose(1, 2)
|
||||
|
||||
assert bhs.shape == canonical.shape
|
||||
assert not torch.equal(bhs, canonical)
|
||||
|
||||
_, lse_from_bhs = update_out_and_lse(None, None, block_out, bhs, lse_layout="bhs")
|
||||
_, lse_from_bsh = update_out_and_lse(None, None, block_out, canonical, lse_layout="bsh")
|
||||
|
||||
assert torch.equal(lse_from_bhs.squeeze(-1), canonical)
|
||||
assert torch.equal(lse_from_bsh.squeeze(-1), canonical)
|
||||
|
||||
_, lse_if_bsh_read_as_bhs = update_out_and_lse(None, None, block_out, canonical, lse_layout="bhs")
|
||||
assert not torch.equal(lse_if_bsh_read_as_bhs.squeeze(-1), canonical)
|
||||
assert torch.equal(lse_if_bsh_read_as_bhs.squeeze(-1), bhs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"invalid_lse",
|
||||
[
|
||||
torch.randn(2, 5),
|
||||
torch.randn(2, 4, 4),
|
||||
torch.randn(2, 3, 5, 2),
|
||||
],
|
||||
)
|
||||
def test_ring_lse_unknown_layout_is_rejected(invalid_lse):
|
||||
block_out = torch.randn(2, 5, 3, 4)
|
||||
|
||||
with pytest.raises(ValueError, match="Ring LSE"):
|
||||
update_out_and_lse(None, None, block_out, invalid_lse, lse_layout="bhs")
|
||||
|
||||
|
||||
def test_ring_accumulation_rejects_mismatched_block_output():
|
||||
out = torch.randn(2, 5, 3, 4)
|
||||
lse = torch.randn(2, 5, 3, 1)
|
||||
block_out = torch.randn(2, 6, 3, 4)
|
||||
block_lse = torch.randn(2, 3, 6)
|
||||
|
||||
with pytest.raises(ValueError, match="block output shape"):
|
||||
update_out_and_lse(out, lse, block_out, block_lse, lse_layout="bhs")
|
||||
|
||||
|
||||
def test_ring_accumulation_matches_logsumexp_weighting():
|
||||
first_out = torch.zeros(1, 2, 1, 1)
|
||||
second_out = torch.ones(1, 2, 1, 1)
|
||||
first_lse = torch.zeros(1, 1, 2)
|
||||
second_lse = torch.zeros(1, 1, 2)
|
||||
|
||||
out, lse = update_out_and_lse(None, None, first_out, first_lse, lse_layout="bhs")
|
||||
out, lse = update_out_and_lse(out, lse, second_out, second_lse, lse_layout="bhs")
|
||||
|
||||
assert torch.allclose(out, torch.full_like(out, 0.5))
|
||||
assert torch.allclose(lse, torch.full_like(lse, torch.log(torch.tensor(2.0))))
|
||||
@@ -0,0 +1,36 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_sage_attention_rejects_mask_instead_of_ignoring_it(monkeypatch):
|
||||
fake_package = types.ModuleType("sageattention")
|
||||
fake_package.sageattn = lambda *args, **kwargs: pytest.fail("unexpected Sage kernel call")
|
||||
monkeypatch.setitem(sys.modules, "sageattention", fake_package)
|
||||
module_name = "vllm_omni.diffusion.attention.backends.sage_attn"
|
||||
sys.modules.pop(module_name, None)
|
||||
try:
|
||||
backend_module = importlib.import_module(module_name)
|
||||
impl = backend_module.SageAttentionImpl(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
softmax_scale=1.0 / 8.0,
|
||||
causal=False,
|
||||
)
|
||||
query = torch.randn(1, 2, 4, 64)
|
||||
metadata = AttentionMetadata(attn_mask=torch.ones(1, 2, dtype=torch.bool))
|
||||
|
||||
with pytest.raises(ValueError, match="does not support attn_mask"):
|
||||
impl.forward_cuda(query, query, query, metadata)
|
||||
finally:
|
||||
sys.modules.pop(module_name, None)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
@@ -9,6 +9,8 @@ import pytest
|
||||
import torch
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@@ -17,7 +19,7 @@ SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"
|
||||
|
||||
def load_sage_attn3_module(monkeypatch: pytest.MonkeyPatch, kernel_impl):
|
||||
fake_module = types.ModuleType("sageattn3")
|
||||
fake_module.sageattn3_blackwell = kernel_impl
|
||||
setattr(fake_module, "sageattn3_blackwell", kernel_impl)
|
||||
monkeypatch.setitem(sys.modules, "sageattn3", fake_module)
|
||||
sys.modules.pop(SAGE_ATTN3_MODULE, None)
|
||||
return importlib.import_module(SAGE_ATTN3_MODULE)
|
||||
@@ -51,21 +53,11 @@ def test_sage_attn3_forward_uses_blackwell_layout(monkeypatch: pytest.MonkeyPatc
|
||||
assert torch.allclose(output, expected)
|
||||
|
||||
|
||||
def test_sage_attn3_falls_back_to_sdpa_for_gqa(monkeypatch: pytest.MonkeyPatch):
|
||||
def test_sage_attn3_rejects_gqa_instead_of_falling_back(monkeypatch: pytest.MonkeyPatch):
|
||||
def fake_kernel(*args, **kwargs):
|
||||
raise AssertionError("sageattn3_blackwell should not be used for GQA")
|
||||
|
||||
backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
|
||||
sdpa_calls = {}
|
||||
|
||||
def fake_sdpa(query, key, value, **kwargs):
|
||||
sdpa_calls["query_shape"] = query.shape
|
||||
sdpa_calls["key_shape"] = key.shape
|
||||
sdpa_calls["enable_gqa"] = kwargs["enable_gqa"]
|
||||
return query + 1
|
||||
|
||||
monkeypatch.setattr(backend_module.F, "scaled_dot_product_attention", fake_sdpa)
|
||||
|
||||
impl = backend_module.SageAttention3Impl(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
@@ -77,13 +69,38 @@ def test_sage_attn3_falls_back_to_sdpa_for_gqa(monkeypatch: pytest.MonkeyPatch):
|
||||
key = torch.randn(2, 8, 2, 64)
|
||||
value = torch.randn(2, 8, 2, 64)
|
||||
|
||||
output = impl.forward_cuda(query, key, value)
|
||||
with pytest.raises(NotImplementedError, match="does not support GQA/MQA"):
|
||||
impl.forward_cuda(query, key, value)
|
||||
|
||||
assert sdpa_calls["query_shape"] == (2, 4, 8, 64)
|
||||
assert sdpa_calls["key_shape"] == (2, 2, 8, 64)
|
||||
assert sdpa_calls["enable_gqa"] is True
|
||||
expected = (query.permute(0, 2, 1, 3) + 1).permute(0, 2, 1, 3)
|
||||
assert torch.allclose(output, expected)
|
||||
|
||||
def test_sage_attn3_rejects_mask_instead_of_ignoring_it(monkeypatch: pytest.MonkeyPatch):
|
||||
def fake_kernel(*args, **kwargs):
|
||||
raise AssertionError("sageattn3_blackwell should not run with an unsupported mask")
|
||||
|
||||
backend_module = load_sage_attn3_module(monkeypatch, fake_kernel)
|
||||
impl = backend_module.SageAttention3Impl(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
softmax_scale=1.0 / 8.0,
|
||||
causal=False,
|
||||
)
|
||||
query = torch.randn(2, 8, 4, 64)
|
||||
metadata = AttentionMetadata(attn_mask=torch.ones(2, 8, dtype=torch.bool))
|
||||
|
||||
with pytest.raises(ValueError, match="does not support attn_mask"):
|
||||
impl.forward_cuda(query, query, query, metadata)
|
||||
|
||||
|
||||
def test_sage_attn3_rejects_custom_softmax_scale(monkeypatch: pytest.MonkeyPatch):
|
||||
backend_module = load_sage_attn3_module(monkeypatch, lambda *args, **kwargs: None)
|
||||
|
||||
with pytest.raises(ValueError, match="does not expose a custom softmax scale"):
|
||||
backend_module.SageAttention3Impl(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
softmax_scale=1.0,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
|
||||
@@ -115,10 +132,39 @@ def test_cuda_platform_selects_sage_attn3_alias(monkeypatch: pytest.MonkeyPatch)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
|
||||
def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch: pytest.MonkeyPatch):
|
||||
@pytest.mark.parametrize("head_size", [32, 320])
|
||||
def test_cuda_platform_rejects_explicit_sage_attn3_unsupported_head_size(
|
||||
monkeypatch: pytest.MonkeyPatch, head_size: int
|
||||
):
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
||||
from vllm_omni.platforms.cuda import platform as cuda_platform_module
|
||||
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
||||
|
||||
original_import_module = importlib.import_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
CudaOmniPlatform,
|
||||
"get_device_capability",
|
||||
classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)),
|
||||
)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
|
||||
monkeypatch.setattr(
|
||||
cuda_platform_module.importlib,
|
||||
"import_module",
|
||||
lambda module_name: object() if module_name == "sageattn3" else original_import_module(module_name),
|
||||
)
|
||||
load_sage_attn3_module(monkeypatch, lambda *args, **kwargs: None)
|
||||
|
||||
with pytest.raises(ValueError, match=f"head_size={head_size} is unsupported"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=head_size)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
|
||||
def test_cuda_platform_rejects_explicit_sage_attn3_on_unsupported_gpu(monkeypatch: pytest.MonkeyPatch):
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
|
||||
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
||||
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
||||
|
||||
@@ -129,6 +175,33 @@ def test_cuda_platform_falls_back_when_sage_attn3_gpu_is_unsupported(monkeypatch
|
||||
)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
|
||||
|
||||
backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
|
||||
with pytest.raises(ValueError, match="explicitly selected"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
|
||||
|
||||
assert backend_path == DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_cuda(), reason="sage_attn3 tests require CUDA platform")
|
||||
def test_cuda_platform_rejects_missing_explicit_sage_attn3(monkeypatch: pytest.MonkeyPatch):
|
||||
from vllm.platforms.interface import DeviceCapability
|
||||
|
||||
from vllm_omni.diffusion.envs import PACKAGES_CHECKER
|
||||
from vllm_omni.platforms.cuda import platform as cuda_platform_module
|
||||
from vllm_omni.platforms.cuda.platform import CudaOmniPlatform
|
||||
|
||||
original_import_module = importlib.import_module
|
||||
|
||||
monkeypatch.setattr(
|
||||
CudaOmniPlatform,
|
||||
"get_device_capability",
|
||||
classmethod(lambda cls, device_id=0: DeviceCapability(10, 0)),
|
||||
)
|
||||
monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": False})
|
||||
|
||||
def _missing_sageattn3(module_name):
|
||||
if module_name == "sageattn3":
|
||||
raise ImportError("sageattn3 not installed")
|
||||
return original_import_module(module_name)
|
||||
|
||||
monkeypatch.setattr(cuda_platform_module.importlib, "import_module", _missing_sageattn3)
|
||||
|
||||
with pytest.raises(ImportError, match="explicitly selected"):
|
||||
CudaOmniPlatform.get_diffusion_attn_backend_cls("SAGE_ATTN_3", head_size=64)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit tests for diffusion attention backend selection."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
@@ -100,9 +100,12 @@ def test_platform_default_used_without_resolved_spec(monkeypatch, attention_conf
|
||||
def test_backend_class_resolution_is_cached(monkeypatch):
|
||||
platform_calls = []
|
||||
load_calls = []
|
||||
fake_platform = SimpleNamespace(
|
||||
get_diffusion_attn_backend_cls=lambda **kwargs: (platform_calls.append(kwargs) or "fake.module.Backend")
|
||||
)
|
||||
|
||||
def fake_get_cls(**kwargs):
|
||||
platform_calls.append(kwargs)
|
||||
return "fake.module.Backend"
|
||||
|
||||
fake_platform = SimpleNamespace(get_diffusion_attn_backend_cls=fake_get_cls)
|
||||
|
||||
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform)
|
||||
|
||||
@@ -126,6 +129,39 @@ def test_backend_class_resolution_is_cached(monkeypatch):
|
||||
assert load_calls == ["fake.module.Backend"]
|
||||
|
||||
|
||||
def test_capability_query_loads_explicit_cudnn_without_platform(monkeypatch):
|
||||
platform_calls = []
|
||||
|
||||
def fake_get_cls(**kwargs):
|
||||
platform_calls.append(kwargs)
|
||||
return "fake.module.MustNotBeCalled"
|
||||
|
||||
fake_platform = SimpleNamespace(get_diffusion_attn_backend_cls=fake_get_cls)
|
||||
monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform)
|
||||
|
||||
config = AttentionConfig(default=AttentionSpec(backend="CUDNN_ATTN"))
|
||||
backend = selector.get_attn_backend_for_capability(role="self", attention_config=config)
|
||||
|
||||
assert backend.get_name() == "CUDNN_ATTN"
|
||||
assert backend.supports_attention_mask()
|
||||
assert platform_calls == []
|
||||
|
||||
|
||||
def test_capability_query_uses_unknown_head_size_for_platform_default(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_get_backend(backend_name, head_size, allow_trtllm_default=True):
|
||||
calls.append((backend_name, head_size, allow_trtllm_default))
|
||||
return _PlatformBackend
|
||||
|
||||
monkeypatch.setattr(selector, "_cached_get_backend_cls", fake_get_backend)
|
||||
|
||||
backend = selector.get_attn_backend_for_capability(role="self", attention_config=AttentionConfig())
|
||||
|
||||
assert backend is _PlatformBackend
|
||||
assert calls == [(None, selector.HEAD_SIZE_UNKNOWN, True)]
|
||||
|
||||
|
||||
def test_load_backend_cls_reports_missing_module():
|
||||
with pytest.raises(ImportError, match="Failed to import module missing_attention_backend"):
|
||||
selector._load_backend_cls("missing_attention_backend.Backend")
|
||||
|
||||
@@ -29,31 +29,343 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
|
||||
lambda default: "advanced_uaa",
|
||||
def test_uaa_gqa_head_padding_preserves_the_query_to_kv_ratio(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Q must be padded to a multiple of the *KV* padding, not of world_size.
|
||||
|
||||
Rounding each up independently is the obvious implementation and is wrong:
|
||||
28Q/7KV at SP2 would become 28/8, i.e. 14/4 per rank, whose 3.5 ratio is not
|
||||
a valid GQA shape. Deriving Q from the padded KV count gives 32/8 -> 16/4.
|
||||
"""
|
||||
from vllm_omni.diffusion.attention.parallel import ulysses
|
||||
|
||||
class FakeGroup:
|
||||
ulysses_group = object()
|
||||
ulysses_world_size = 2
|
||||
ulysses_rank = 0
|
||||
ring_world_size = 1
|
||||
|
||||
observed = []
|
||||
|
||||
def fake_all_to_all(pg, tensor, **kwargs):
|
||||
observed.append(kwargs["padded_head_cnt"])
|
||||
return tensor, tensor.shape[2]
|
||||
|
||||
monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
|
||||
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
|
||||
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
|
||||
strategy = ulysses.UlyssesParallelAttention(
|
||||
FakeGroup(),
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
|
||||
with set_forward_context():
|
||||
strategy.pre_attention(
|
||||
torch.zeros(1, 3, 28, 4),
|
||||
torch.zeros(1, 3, 7, 4),
|
||||
torch.zeros(1, 3, 7, 4),
|
||||
None,
|
||||
)
|
||||
|
||||
q_padded, k_padded, v_padded = observed
|
||||
assert (q_padded, k_padded, v_padded) == (32, 8, 8)
|
||||
# The per-rank shape must still be a valid GQA shape.
|
||||
assert (q_padded // 2) % (k_padded // 2) == 0
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_uaa_rejects_head_counts_that_are_not_a_gqa_shape(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.attention.parallel import ulysses
|
||||
|
||||
class FakeGroup:
|
||||
ulysses_group = object()
|
||||
ulysses_world_size = 2
|
||||
ulysses_rank = 0
|
||||
ring_world_size = 1
|
||||
|
||||
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
|
||||
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
|
||||
strategy = ulysses.UlyssesParallelAttention(
|
||||
FakeGroup(),
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
|
||||
with set_forward_context(), pytest.raises(ValueError, match="multiple of KV heads"):
|
||||
strategy.pre_attention(
|
||||
torch.zeros(1, 3, 10, 4),
|
||||
torch.zeros(1, 3, 4, 4),
|
||||
torch.zeros(1, 3, 4, 4),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize(
|
||||
"q_heads, kv_heads, world_size, expected_warn",
|
||||
[
|
||||
# 28Q/7KV @ U=2 -> pad to 32Q, 32/28 = 1.14x, below threshold, no warn.
|
||||
(28, 7, 2, False),
|
||||
# 32Q/1KV @ U=8 -> pad KV to 8 -> pad Q to 256, 256/32 = 8x, warn.
|
||||
(32, 1, 8, True),
|
||||
# 16Q/2KV @ U=4 -> pad KV to 4 -> pad Q to 32, 32/16 = 2x, warn.
|
||||
(16, 2, 4, True),
|
||||
],
|
||||
)
|
||||
def test_uaa_padding_ratio_warning_fires_when_blowup_is_large(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
q_heads: int,
|
||||
kv_heads: int,
|
||||
world_size: int,
|
||||
expected_warn: bool,
|
||||
) -> None:
|
||||
"""advanced_uaa must warn once when GQA padding materially inflates Q heads.
|
||||
|
||||
MQA-style shapes at high ulysses_degree pad Q by U x, which grows attention
|
||||
FLOPs and temporary VRAM by the same factor. We do not fall back (there is
|
||||
no cheap alternative -- Ulysses needs KV divisible by U for the head-dim
|
||||
split), but we surface a one-shot warning so the user can pick a
|
||||
ulysses_degree that divides KV_heads if the overhead is unacceptable.
|
||||
"""
|
||||
from vllm_omni.diffusion.attention.parallel import ulysses
|
||||
|
||||
# Reset the warn-once dedup set so parametrized cases do not shadow each other.
|
||||
ulysses._uaa_pad_ratio_warned.clear()
|
||||
|
||||
class FakeGroup:
|
||||
ulysses_group = object()
|
||||
ulysses_world_size = world_size
|
||||
ulysses_rank = 0
|
||||
ring_world_size = 1
|
||||
|
||||
def fake_all_to_all(pg, tensor, **kwargs):
|
||||
return tensor, tensor.shape[2]
|
||||
|
||||
monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
|
||||
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
|
||||
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *a, **kw: [3] * world_size)
|
||||
strategy = ulysses.UlyssesParallelAttention(
|
||||
FakeGroup(),
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
|
||||
caplog.set_level("WARNING", logger="vllm_omni.diffusion.attention.parallel.ulysses")
|
||||
with set_forward_context():
|
||||
strategy.pre_attention(
|
||||
torch.zeros(1, 3, q_heads, 4),
|
||||
torch.zeros(1, 3, kv_heads, 4),
|
||||
torch.zeros(1, 3, kv_heads, 4),
|
||||
None,
|
||||
)
|
||||
|
||||
warn_msgs = [r.message for r in caplog.records if "GQA padding inflates" in r.message]
|
||||
if expected_warn:
|
||||
assert len(warn_msgs) == 1, warn_msgs
|
||||
# Second call with the same shape must not re-warn (warn-once dedup).
|
||||
with set_forward_context():
|
||||
strategy.pre_attention(
|
||||
torch.zeros(1, 3, q_heads, 4),
|
||||
torch.zeros(1, 3, kv_heads, 4),
|
||||
torch.zeros(1, 3, kv_heads, 4),
|
||||
None,
|
||||
)
|
||||
warn_msgs = [r.message for r in caplog.records if "GQA padding inflates" in r.message]
|
||||
assert len(warn_msgs) == 1, "warn-once dedup should suppress the second call"
|
||||
else:
|
||||
assert warn_msgs == [], f"unexpected warning for {q_heads}Q/{kv_heads}KV @ U={world_size}: {warn_msgs}"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_uaa_joint_gqa_head_padding_preserves_the_query_to_kv_ratio(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Joint Q/K/V must also preserve the GQA ratio through Ulysses padding.
|
||||
|
||||
MMDiT-style models (SD3, FLUX, Qwen-Image, Hunyuan-Video, ...) pass a
|
||||
separate text stream via ``AttentionMetadata.joint_*``. Before the fix the
|
||||
joint block padded K/V by Q's amount and integer-divided K/V by world_size,
|
||||
which for 28Q/7KV at SP=2 produces mismatched joint vs main head counts and
|
||||
the downstream ring-attention concat blows up. Post-fix the joint block
|
||||
mirrors the main-path rule: pad KV to a world_size multiple, derive Q as
|
||||
``padded_kv * (Q // KV)``.
|
||||
"""
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.parallel import ulysses
|
||||
|
||||
class FakeGroup:
|
||||
ulysses_group = object()
|
||||
ulysses_world_size = 2
|
||||
ulysses_rank = 0
|
||||
ring_world_size = 1
|
||||
|
||||
def fake_all_to_all(pg, tensor, **kwargs):
|
||||
# Simulate the head-dim sharding so the post-a2a joint concat has
|
||||
# matching per-rank head counts on Q and joint_Q.
|
||||
heads_per_rank = kwargs["padded_head_cnt"] // 2
|
||||
return tensor[:, :, :heads_per_rank, :].contiguous(), tensor.shape[2]
|
||||
|
||||
monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all)
|
||||
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
|
||||
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
|
||||
strategy = ulysses.UlyssesParallelAttention(
|
||||
FakeGroup(),
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
|
||||
attn_metadata = AttentionMetadata(
|
||||
joint_query=torch.zeros(1, 5, 28, 4),
|
||||
joint_key=torch.zeros(1, 5, 7, 4),
|
||||
joint_value=torch.zeros(1, 5, 7, 4),
|
||||
)
|
||||
|
||||
with set_forward_context():
|
||||
strategy.pre_attention(
|
||||
torch.zeros(1, 3, 28, 4),
|
||||
torch.zeros(1, 3, 7, 4),
|
||||
torch.zeros(1, 3, 7, 4),
|
||||
attn_metadata,
|
||||
)
|
||||
|
||||
# 28Q/7KV at SP=2 -> pad to 32/8 -> per-rank slice is 16Q/4KV.
|
||||
assert attn_metadata.joint_key.shape[-2] == 4
|
||||
assert attn_metadata.joint_value.shape[-2] == 4
|
||||
# Per-rank joint shape must remain a valid GQA shape (Q multiple of KV).
|
||||
per_rank_joint_q = 32 // 2
|
||||
per_rank_joint_kv = attn_metadata.joint_key.shape[-2]
|
||||
assert per_rank_joint_q % per_rank_joint_kv == 0
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_uaa_rejects_joint_head_counts_that_are_not_a_gqa_shape(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.parallel import ulysses
|
||||
|
||||
class FakeGroup:
|
||||
ulysses_group = object()
|
||||
ulysses_world_size = 2
|
||||
ulysses_rank = 0
|
||||
ring_world_size = 1
|
||||
|
||||
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
|
||||
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
|
||||
strategy = ulysses.UlyssesParallelAttention(
|
||||
FakeGroup(),
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
|
||||
attn_metadata = AttentionMetadata(
|
||||
joint_query=torch.zeros(1, 5, 10, 4),
|
||||
joint_key=torch.zeros(1, 5, 4, 4),
|
||||
joint_value=torch.zeros(1, 5, 4, 4),
|
||||
)
|
||||
|
||||
with set_forward_context(), pytest.raises(ValueError, match="multiple of joint KV heads"):
|
||||
strategy.pre_attention(
|
||||
torch.zeros(1, 3, 8, 4),
|
||||
torch.zeros(1, 3, 4, 4),
|
||||
torch.zeros(1, 3, 4, 4),
|
||||
attn_metadata,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_advanced_uaa_hybrid_rejects_non_gqa_shapes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Hybrid Ulysses+Ring still rejects shapes that are not a GQA shape.
|
||||
|
||||
Valid GQA shapes are safe under advanced_uaa head padding (the derived-Q
|
||||
rule keeps real Q heads paired with real K/V heads after the split, see
|
||||
pre_attention), but a query count that is not a multiple of the K/V count
|
||||
cannot be padded while preserving the ratio, so it must be rejected.
|
||||
"""
|
||||
from vllm_omni.diffusion.attention.parallel import ulysses
|
||||
|
||||
sp_group = SimpleNamespace(
|
||||
ulysses_group=None,
|
||||
ulysses_world_size=2,
|
||||
ulysses_rank=0,
|
||||
ring_world_size=2,
|
||||
ring_group=object(),
|
||||
)
|
||||
monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa")
|
||||
monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3])
|
||||
strategy = UlyssesParallelAttention(
|
||||
sp_group=sp_group,
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
query = torch.zeros(1, 2, 8, 4)
|
||||
key = torch.zeros(1, 2, 1, 4)
|
||||
query = torch.zeros(1, 2, 10, 4)
|
||||
key = torch.zeros(1, 2, 4, 4)
|
||||
value = torch.zeros_like(key)
|
||||
|
||||
with pytest.raises(ValueError, match="does not support K/V head padding"):
|
||||
with pytest.raises(ValueError, match="multiple of KV heads"):
|
||||
strategy.pre_attention(query, key, value, None)
|
||||
|
||||
|
||||
def test_strict_ulysses_reshards_vsa_gate_with_qkv(monkeypatch) -> None:
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
|
||||
lambda default: "strict",
|
||||
)
|
||||
calls: list[torch.Tensor] = []
|
||||
|
||||
def fake_all_to_all(pg, tensor, scatter_idx, gather_idx, use_sync):
|
||||
del pg, scatter_idx, gather_idx, use_sync
|
||||
calls.append(tensor)
|
||||
return tensor.reshape(tensor.shape[0], tensor.shape[1] * 2, tensor.shape[2] // 2, tensor.shape[3])
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.attention.parallel.ulysses.SeqAllToAll4D.apply",
|
||||
fake_all_to_all,
|
||||
)
|
||||
strategy = UlyssesParallelAttention(
|
||||
sp_group=SimpleNamespace(
|
||||
ulysses_group=object(),
|
||||
ulysses_world_size=2,
|
||||
ulysses_rank=0,
|
||||
ring_world_size=1,
|
||||
),
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
query = torch.arange(24).reshape(1, 3, 4, 2)
|
||||
key = query + 10
|
||||
value = query + 20
|
||||
gate = query + 100
|
||||
metadata = AttentionMetadata(extra={"gate_compress": gate})
|
||||
|
||||
query_out, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
|
||||
|
||||
assert all(actual is expected for actual, expected in zip(calls, (query, key, value, gate), strict=True))
|
||||
assert query_out.shape == (1, 6, 2, 2)
|
||||
assert metadata_out is metadata
|
||||
assert metadata.extra["gate_compress"].shape == query_out.shape
|
||||
torch.testing.assert_close(metadata.extra["gate_compress"], gate.reshape_as(query_out))
|
||||
|
||||
|
||||
def _run_attention_case(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
@@ -67,6 +379,9 @@ def _run_attention_case(
|
||||
ring_degree: int = 1,
|
||||
split_sizes: list[int] | None = None,
|
||||
sdp_kernel_mode: str = "math",
|
||||
num_kv_heads: int | None = None,
|
||||
has_joint: bool = False,
|
||||
joint_strategy: str = "front",
|
||||
) -> None:
|
||||
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
|
||||
current_omni_platform.set_device(device)
|
||||
@@ -102,15 +417,33 @@ def _run_attention_case(
|
||||
head_size=head_size,
|
||||
causal=False,
|
||||
softmax_scale=1.0 / (head_size**0.5),
|
||||
num_kv_heads=num_kv_heads,
|
||||
).to(device=device, dtype=torch.float32)
|
||||
|
||||
with np.load(input_file, allow_pickle=False) as payload:
|
||||
q_full = torch.from_numpy(payload["q"]).to(device=device)
|
||||
k_full = torch.from_numpy(payload["k"]).to(device=device)
|
||||
v_full = torch.from_numpy(payload["v"]).to(device=device)
|
||||
if has_joint:
|
||||
joint_q = torch.from_numpy(payload["joint_q"]).to(device=device)
|
||||
joint_k = torch.from_numpy(payload["joint_k"]).to(device=device)
|
||||
joint_v = torch.from_numpy(payload["joint_v"]).to(device=device)
|
||||
|
||||
attn_metadata = None
|
||||
if world_size == 1:
|
||||
q, k, v = q_full, k_full, v_full
|
||||
if has_joint:
|
||||
# Baseline: SP is inactive so pre_attention won't concat joint for us.
|
||||
# Build the full sequence directly and skip the joint metadata.
|
||||
if joint_strategy == "front":
|
||||
q = torch.cat([joint_q, q_full], dim=1)
|
||||
k = torch.cat([joint_k, k_full], dim=1)
|
||||
v = torch.cat([joint_v, v_full], dim=1)
|
||||
else:
|
||||
q = torch.cat([q_full, joint_q], dim=1)
|
||||
k = torch.cat([k_full, joint_k], dim=1)
|
||||
v = torch.cat([v_full, joint_v], dim=1)
|
||||
else:
|
||||
q, k, v = q_full, k_full, v_full
|
||||
else:
|
||||
if split_sizes is None:
|
||||
# NOTE: torch.chunk may return fewer than `world_size` chunks for some
|
||||
@@ -135,6 +468,16 @@ def _run_attention_case(
|
||||
# shard manually for testing.
|
||||
get_forward_context()._sp_shard_depth = 1
|
||||
|
||||
if has_joint:
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
|
||||
attn_metadata = AttentionMetadata(
|
||||
joint_query=joint_q,
|
||||
joint_key=joint_k,
|
||||
joint_value=joint_v,
|
||||
joint_strategy=joint_strategy,
|
||||
)
|
||||
|
||||
if sdp_kernel_mode == "math":
|
||||
sdp_ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
|
||||
elif sdp_kernel_mode == "mem_efficient":
|
||||
@@ -143,25 +486,47 @@ def _run_attention_case(
|
||||
raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}")
|
||||
|
||||
with sdp_ctx:
|
||||
out_local = attn(q, k, v).contiguous()
|
||||
out_local = attn(q, k, v, attn_metadata=attn_metadata).contiguous()
|
||||
|
||||
if world_size == 1:
|
||||
out_full = out_local
|
||||
else:
|
||||
local_len = torch.tensor([out_local.shape[1]], device=device, dtype=torch.int64)
|
||||
if has_joint:
|
||||
joint_len = int(joint_q.shape[1])
|
||||
if joint_strategy == "front":
|
||||
joint_out_local = out_local[:, :joint_len]
|
||||
img_out_local = out_local[:, joint_len:]
|
||||
else:
|
||||
img_out_local = out_local[:, : out_local.shape[1] - joint_len]
|
||||
joint_out_local = out_local[:, -joint_len:]
|
||||
else:
|
||||
joint_out_local = None
|
||||
img_out_local = out_local
|
||||
|
||||
local_len = torch.tensor([img_out_local.shape[1]], device=device, dtype=torch.int64)
|
||||
gathered_lens = [torch.empty_like(local_len) for _ in range(world_size)]
|
||||
torch.distributed.all_gather(gathered_lens, local_len)
|
||||
lens = [int(t.item()) for t in gathered_lens]
|
||||
max_len = max(lens)
|
||||
|
||||
if out_local.shape[1] < max_len:
|
||||
pad = max_len - out_local.shape[1]
|
||||
out_local = torch.nn.functional.pad(out_local, (0, 0, 0, 0, 0, pad)).contiguous()
|
||||
if img_out_local.shape[1] < max_len:
|
||||
pad = max_len - img_out_local.shape[1]
|
||||
img_out_local = torch.nn.functional.pad(img_out_local, (0, 0, 0, 0, 0, pad)).contiguous()
|
||||
else:
|
||||
img_out_local = img_out_local.contiguous()
|
||||
|
||||
gathered = [torch.empty_like(out_local) for _ in range(world_size)]
|
||||
torch.distributed.all_gather(gathered, out_local)
|
||||
gathered = [torch.empty_like(img_out_local) for _ in range(world_size)]
|
||||
torch.distributed.all_gather(gathered, img_out_local)
|
||||
if local_rank == 0:
|
||||
out_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1)
|
||||
img_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1)
|
||||
if has_joint:
|
||||
out_full = (
|
||||
torch.cat([joint_out_local, img_full], dim=1)
|
||||
if joint_strategy == "front"
|
||||
else torch.cat([img_full, joint_out_local], dim=1)
|
||||
)
|
||||
else:
|
||||
out_full = img_full
|
||||
else:
|
||||
out_full = None
|
||||
|
||||
@@ -172,20 +537,43 @@ def _run_attention_case(
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sp_world_size,seq_len,num_heads",
|
||||
"sp_world_size,seq_len,joint_len,num_heads,num_kv_heads",
|
||||
[
|
||||
(2, 6, 3), # head_cnt not divisible by P=2
|
||||
(2, 5, 4), # seq_len not divisible by P=2
|
||||
(4, 9, 30), # Z-Image-like: head_cnt not divisible by P=4
|
||||
(4, 10, 8), # seq_len not divisible by P=4
|
||||
# MHA (num_kv_heads=None, joint_len=None)
|
||||
(2, 6, None, 3, None), # head_cnt not divisible by P=2
|
||||
(2, 5, None, 4, None), # seq_len not divisible by P=2
|
||||
(4, 9, None, 30, None), # Z-Image-like: head_cnt not divisible by P=4
|
||||
(4, 10, None, 8, None), # seq_len not divisible by P=4
|
||||
# GQA (joint_len=None)
|
||||
(2, 6, None, 28, 7), # BOOGU-like GQA: neither Q nor KV divisible by P=2
|
||||
(2, 5, None, 6, 3), # KV divisible by P, Q is not
|
||||
(4, 8, None, 12, 3), # KV not divisible by P=4
|
||||
# Joint GQA (MMDiT-style text stream via AttentionMetadata.joint_*)
|
||||
(2, 6, 5, 28, 7), # BOOGU-like GQA joint: KV=7 not divisible by P=2
|
||||
(2, 5, 4, 6, 3), # KV divisible by P, Q is not
|
||||
],
|
||||
)
|
||||
def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_heads: int) -> None:
|
||||
def test_ulysses_uaa_matches_baseline(
|
||||
sp_world_size: int,
|
||||
seq_len: int,
|
||||
joint_len: int | None,
|
||||
num_heads: int,
|
||||
num_kv_heads: int | None,
|
||||
) -> None:
|
||||
"""End-to-end SP-vs-single-GPU parity under UAA.
|
||||
|
||||
Covers MHA, GQA and joint (MMDiT-style) shapes. The joint variant catches
|
||||
the pre-fix bug where the joint block padded K/V by Q's head count and
|
||||
integer-divided K/V by world_size, giving per-rank shapes that mismatched
|
||||
the main-tensor slice (e.g. 14Q/3KV vs. 16Q/4KV for 28Q/7KV at SP=2).
|
||||
"""
|
||||
if current_omni_platform.get_device_count() < sp_world_size:
|
||||
pytest.skip(f"Test requires {sp_world_size} GPUs")
|
||||
|
||||
batch_size = 2
|
||||
head_size = 8
|
||||
kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
|
||||
has_joint = joint_len is not None
|
||||
|
||||
base_init_method = get_distributed_init_method()
|
||||
sp_init_method = get_distributed_init_method()
|
||||
@@ -200,17 +588,23 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
|
||||
try:
|
||||
torch.manual_seed(0)
|
||||
q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy())
|
||||
k = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32)
|
||||
v = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32)
|
||||
payload = {"q": q.numpy(), "k": k.numpy(), "v": v.numpy()}
|
||||
if has_joint:
|
||||
payload["joint_q"] = torch.randn(batch_size, joint_len, num_heads, head_size, dtype=torch.float32).numpy()
|
||||
payload["joint_k"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
|
||||
payload["joint_v"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
|
||||
np.savez(input_file, **payload)
|
||||
|
||||
common_tail = (1, None, "math", num_kv_heads, has_joint)
|
||||
|
||||
# Baseline (no SP)
|
||||
torch.multiprocessing.spawn(
|
||||
_run_attention_case,
|
||||
args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"),
|
||||
args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict", *common_tail),
|
||||
nprocs=1,
|
||||
)
|
||||
|
||||
# SP (Ulysses-P with UAA)
|
||||
torch.multiprocessing.spawn(
|
||||
_run_attention_case,
|
||||
@@ -223,15 +617,13 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
|
||||
head_size,
|
||||
sp_world_size,
|
||||
"advanced_uaa",
|
||||
*common_tail,
|
||||
),
|
||||
nprocs=sp_world_size,
|
||||
)
|
||||
|
||||
baseline = np.load(baseline_file, allow_pickle=False)
|
||||
sp = np.load(sp_file, allow_pickle=False)
|
||||
|
||||
baseline_t = torch.from_numpy(baseline)
|
||||
sp_t = torch.from_numpy(sp)
|
||||
baseline_t = torch.from_numpy(np.load(baseline_file, allow_pickle=False))
|
||||
sp_t = torch.from_numpy(np.load(sp_file, allow_pickle=False))
|
||||
assert baseline_t.shape == sp_t.shape
|
||||
torch.testing.assert_close(sp_t, baseline_t, atol=1e-5, rtol=1e-5)
|
||||
finally:
|
||||
@@ -242,7 +634,17 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
|
||||
pass
|
||||
|
||||
|
||||
def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
@pytest.mark.parametrize(
|
||||
"num_heads,num_kv_heads,joint_len",
|
||||
[
|
||||
(3, None, None), # MHA baseline (head_cnt not divisible by ulysses_degree=2)
|
||||
(28, 7, None), # GQA: exercises Ring pytorch_attn_forward KV expansion (Q!=K heads)
|
||||
(28, 7, 5), # GQA + MMDiT joint stream: joint K/V ride the padded Ulysses layout and Ring SDPA
|
||||
],
|
||||
)
|
||||
def test_ulysses_uaa_hybrid_ring_matches_baseline(
|
||||
num_heads: int, num_kv_heads: int | None, joint_len: int | None
|
||||
) -> None:
|
||||
sp_world_size = 4
|
||||
ulysses_degree = 2
|
||||
ring_degree = 2
|
||||
@@ -253,9 +655,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
batch_size = 2
|
||||
head_size = 8
|
||||
seq_len = 10
|
||||
# This is the positive Hybrid case: K/V heads must not require padding,
|
||||
# which is unsupported by advanced_uaa when Ring is also enabled.
|
||||
num_heads = 4
|
||||
kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
|
||||
has_joint = joint_len is not None
|
||||
|
||||
# Ensure ring ranks see equal post-Ulysses seq_len:
|
||||
# rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
|
||||
@@ -274,9 +675,14 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
try:
|
||||
torch.manual_seed(0)
|
||||
q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32)
|
||||
np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy())
|
||||
k = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32)
|
||||
v = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32)
|
||||
payload = {"q": q.numpy(), "k": k.numpy(), "v": v.numpy()}
|
||||
if has_joint:
|
||||
payload["joint_q"] = torch.randn(batch_size, joint_len, num_heads, head_size, dtype=torch.float32).numpy()
|
||||
payload["joint_k"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
|
||||
payload["joint_v"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy()
|
||||
np.savez(input_file, **payload)
|
||||
|
||||
# Baseline (no SP)
|
||||
torch.multiprocessing.spawn(
|
||||
@@ -293,6 +699,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
1,
|
||||
None,
|
||||
"mem_efficient",
|
||||
num_kv_heads,
|
||||
has_joint,
|
||||
),
|
||||
nprocs=1,
|
||||
)
|
||||
@@ -312,6 +720,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
ring_degree,
|
||||
split_sizes,
|
||||
"mem_efficient",
|
||||
num_kv_heads,
|
||||
has_joint,
|
||||
),
|
||||
nprocs=sp_world_size,
|
||||
)
|
||||
|
||||
@@ -7,7 +7,7 @@ End-to-end test for diffusion batching via AsyncOmni.
|
||||
This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a
|
||||
diffusion model and validates that every caller receives its correct
|
||||
individual result. When the underlying diffusion stage is configured with
|
||||
``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the
|
||||
``max_num_seqs > 1``, the
|
||||
requests will be batched internally.
|
||||
|
||||
Even without explicit batching config this test is useful for verifying
|
||||
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
|
||||
) -> None:
|
||||
"""Run the same prompts sequentially then in parallel and print a comparison."""
|
||||
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size)
|
||||
try:
|
||||
await warmup(omni, WARMUP_PROMPTS)
|
||||
single_time = await run_single(omni, prompts)
|
||||
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
|
||||
await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
|
||||
return
|
||||
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size)
|
||||
try:
|
||||
await warmup(omni, WARMUP_PROMPTS)
|
||||
|
||||
@@ -389,7 +389,7 @@ def test_diffusion_batching_async_concurrent(model_name: str):
|
||||
fired via asyncio.gather. Each request_id must appear in the results."""
|
||||
|
||||
async def _inner():
|
||||
omni = AsyncOmni(model=model_name, diffusion_batch_size=1)
|
||||
omni = AsyncOmni(model=model_name, max_num_seqs=1)
|
||||
try:
|
||||
prompts = TEST_PROMPTS[:4]
|
||||
sp = _default_sampling_params()
|
||||
@@ -430,7 +430,7 @@ def test_diffusion_batching_list_prompt_rejected(model_name: str):
|
||||
"""
|
||||
|
||||
async def _inner():
|
||||
omni = AsyncOmni(model=model_name, diffusion_batch_size=4)
|
||||
omni = AsyncOmni(model=model_name, max_num_seqs=4)
|
||||
try:
|
||||
prompts = TEST_PROMPTS[:4]
|
||||
sp = _default_sampling_params()
|
||||
@@ -528,7 +528,7 @@ if __name__ == "__main__":
|
||||
"--batch-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
|
||||
help="Scheduler in-flight width (--max-num-seqs)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
|
||||
+32
-17
@@ -6,7 +6,9 @@ Model specific tests for CacheDiT enablement.
|
||||
"""
|
||||
|
||||
import ast
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
@@ -17,7 +19,8 @@ from vllm.distributed import parallel_state
|
||||
import vllm_omni.diffusion.cache.cachedit as cd_backend
|
||||
import vllm_omni.diffusion.cache.cachedit.model_specific as cd_model_specific
|
||||
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary
|
||||
from vllm_omni.diffusion.data import DiffusionCacheConfig
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig
|
||||
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
||||
from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel
|
||||
from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
|
||||
@@ -35,6 +38,17 @@ SEPARATE_CFG_TRANSFORMERS = [
|
||||
SAMPLE_CACHE_CONFIG = DiffusionCacheConfig()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _force_torch_sdpa():
|
||||
"""Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3)."""
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
|
||||
parallel_config=SimpleNamespace(ring_degree=1),
|
||||
)
|
||||
with set_current_diffusion_config(od_config):
|
||||
yield
|
||||
|
||||
|
||||
def test_custom_cache_dit_enablers_are_registered_explicitly():
|
||||
expected_enablers = {
|
||||
"Wan22Pipeline": cd_model_specific.enable_cache_for_wan22,
|
||||
@@ -213,22 +227,23 @@ def test_ltx2_cache_dit_receives_audio_as_encoder(init_fake_tp_group):
|
||||
text_in = torch.full((1, seq_len, 16), 3.0)
|
||||
audio_text_in = torch.full((1, seq_len, 16), 4.0)
|
||||
|
||||
model = LTX2VideoTransformer3DModel(
|
||||
in_channels=16,
|
||||
out_channels=16,
|
||||
patch_size=1,
|
||||
patch_size_t=1,
|
||||
num_attention_heads=2,
|
||||
attention_head_dim=8,
|
||||
cross_attention_dim=16,
|
||||
audio_in_channels=16,
|
||||
audio_out_channels=16,
|
||||
audio_num_attention_heads=2,
|
||||
audio_attention_head_dim=8,
|
||||
audio_cross_attention_dim=16,
|
||||
num_layers=2,
|
||||
caption_channels=16,
|
||||
)
|
||||
with _force_torch_sdpa():
|
||||
model = LTX2VideoTransformer3DModel(
|
||||
in_channels=16,
|
||||
out_channels=16,
|
||||
patch_size=1,
|
||||
patch_size_t=1,
|
||||
num_attention_heads=2,
|
||||
attention_head_dim=8,
|
||||
cross_attention_dim=16,
|
||||
audio_in_channels=16,
|
||||
audio_out_channels=16,
|
||||
audio_num_attention_heads=2,
|
||||
audio_attention_head_dim=8,
|
||||
audio_cross_attention_dim=16,
|
||||
num_layers=2,
|
||||
caption_channels=16,
|
||||
)
|
||||
|
||||
# NOTE: This is currently using the LTX2 custom enabler, but the custom
|
||||
# enablers will be consolidated after
|
||||
|
||||
@@ -13,6 +13,8 @@ import vllm_omni.diffusion.data as diffusion_data
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient, dummy_messages_from_mix_data
|
||||
from vllm_omni.diffusion.data import (
|
||||
AttentionConfig,
|
||||
AttentionSpec,
|
||||
DiffusionOutput,
|
||||
DiffusionParallelConfig,
|
||||
OmniDiffusionConfig,
|
||||
@@ -141,9 +143,96 @@ class TestPipelineArgumentsHandling:
|
||||
assert isinstance(output.output, MockPipelineOutput)
|
||||
assert output.output.image is stub_image
|
||||
|
||||
@staticmethod
|
||||
def _make_adapter(backend: str, set_attention_backend):
|
||||
adapter = DiffusersAdapterPipeline(
|
||||
od_config=_make_od_config(
|
||||
diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend=backend)),
|
||||
)
|
||||
)
|
||||
adapter._pipeline = SimpleNamespace(
|
||||
transformer=SimpleNamespace(set_attention_backend=set_attention_backend),
|
||||
)
|
||||
return adapter
|
||||
|
||||
def test_explicit_unbound_backend_is_rejected(self):
|
||||
calls: list[str] = []
|
||||
adapter = self._make_adapter("CUDNN_ATTN", calls.append)
|
||||
|
||||
with pytest.raises(ValueError, match="do not provide a binding"):
|
||||
adapter._set_attention_backend()
|
||||
|
||||
assert calls == []
|
||||
|
||||
def test_failed_explicit_backend_does_not_fallback_to_native(self):
|
||||
calls: list[str] = []
|
||||
|
||||
def reject_backend(backend):
|
||||
calls.append(backend)
|
||||
raise RuntimeError(f"{backend} unavailable")
|
||||
|
||||
adapter = self._make_adapter("SAGE_ATTN", reject_backend)
|
||||
|
||||
with pytest.raises(RuntimeError, match="explicitly selected attention backend 'SAGE_ATTN'"):
|
||||
adapter._set_attention_backend()
|
||||
|
||||
assert calls == ["sage_hub", "sage", "sage_varlen"]
|
||||
assert "native" not in calls
|
||||
|
||||
def test_explicit_flash_attn_3_hub_does_not_fallback_to_fa2_or_local(self, monkeypatch):
|
||||
monkeypatch.setattr(DiffusersAdapterPipeline, "_is_blackwell_gpu", lambda self: False)
|
||||
calls: list[str] = []
|
||||
|
||||
def accept_later_backends(backend):
|
||||
calls.append(backend)
|
||||
if backend in ("flash_hub", "flash_varlen_hub", "flash", "flash_varlen", "_native_flash"):
|
||||
return
|
||||
raise RuntimeError(f"{backend} unavailable")
|
||||
|
||||
adapter = self._make_adapter("FLASH_ATTN_3_HUB", accept_later_backends)
|
||||
|
||||
with pytest.raises(RuntimeError, match="explicitly selected attention backend 'FLASH_ATTN_3_HUB'"):
|
||||
adapter._set_attention_backend()
|
||||
|
||||
assert calls == ["_flash_3_hub", "_flash_3_varlen_hub"]
|
||||
|
||||
def test_explicit_flash_attn_hub_does_not_fallback_to_local(self, monkeypatch):
|
||||
monkeypatch.setattr(DiffusersAdapterPipeline, "_is_blackwell_gpu", lambda self: False)
|
||||
calls: list[str] = []
|
||||
|
||||
def accept_local_flash(backend):
|
||||
calls.append(backend)
|
||||
if backend in ("flash", "flash_varlen", "_native_flash"):
|
||||
return
|
||||
raise RuntimeError(f"{backend} unavailable")
|
||||
|
||||
adapter = self._make_adapter("FLASH_ATTN_HUB", accept_local_flash)
|
||||
|
||||
with pytest.raises(RuntimeError, match="explicitly selected attention backend 'FLASH_ATTN_HUB'"):
|
||||
adapter._set_attention_backend()
|
||||
|
||||
assert calls == ["flash_hub", "flash_varlen_hub"]
|
||||
|
||||
def test_explicit_torch_sdpa_maps_to_diffusers_native(self):
|
||||
calls: list[str] = []
|
||||
adapter = self._make_adapter("TORCH_SDPA", calls.append)
|
||||
|
||||
adapter._set_attention_backend()
|
||||
|
||||
assert calls == ["native"]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"feature_id",
|
||||
["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"],
|
||||
[
|
||||
"cfg_parallel",
|
||||
"tensor_parallel",
|
||||
"ulysses",
|
||||
"ring",
|
||||
"teacache",
|
||||
"cache_dit",
|
||||
"enforce_eager",
|
||||
"unsupported_quantization",
|
||||
],
|
||||
)
|
||||
def test_adapter_guard_unsupported_feature(self, feature_id):
|
||||
if feature_id == "cfg_parallel":
|
||||
@@ -151,6 +240,11 @@ class TestPipelineArgumentsHandling:
|
||||
parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1),
|
||||
cache_backend="none",
|
||||
)
|
||||
elif feature_id == "tensor_parallel":
|
||||
od_config = _make_od_config(
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=2, sequence_parallel_size=1),
|
||||
cache_backend="none",
|
||||
)
|
||||
elif feature_id == "ulysses":
|
||||
od_config = _make_od_config(
|
||||
parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ulysses_degree=2),
|
||||
|
||||
@@ -17,7 +17,12 @@ import torch
|
||||
|
||||
from tests.helpers.mark import hardware_marks, hardware_test
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _dispatch_branches, _wrap
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import (
|
||||
CFGParallelMixin,
|
||||
_dispatch_branches,
|
||||
_get_cfg_world_size_or_one,
|
||||
_wrap,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
destroy_distributed_env,
|
||||
get_classifier_free_guidance_rank,
|
||||
@@ -393,6 +398,68 @@ def _build_pipeline_on_device(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.cpu
|
||||
def test_cfg_world_size_defaults_only_when_group_is_uninitialized(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setattr(parallel_state, "_CFG", None)
|
||||
|
||||
assert not parallel_state.is_cfg_group_initialized()
|
||||
assert _get_cfg_world_size_or_one() == 1
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.cpu
|
||||
def test_cfg_world_size_does_not_swallow_initialized_group_errors(monkeypatch: pytest.MonkeyPatch):
|
||||
class BrokenCfgGroup:
|
||||
@property
|
||||
def world_size(self) -> int:
|
||||
raise AssertionError("CFG group is broken")
|
||||
|
||||
monkeypatch.setattr(parallel_state, "_CFG", BrokenCfgGroup())
|
||||
|
||||
assert parallel_state.is_cfg_group_initialized()
|
||||
with pytest.raises(AssertionError, match="CFG group is broken"):
|
||||
_get_cfg_world_size_or_one()
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.cpu
|
||||
def test_cfg_defaults_to_sequential_without_initialized_group():
|
||||
pipeline = _build_pipeline_on_device(
|
||||
TestCFGPipeline,
|
||||
hidden_dim=32,
|
||||
model_seed=42,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
positive_kwargs, negative_kwargs = _make_two_branch_inputs(
|
||||
batch_size=1,
|
||||
channels=4,
|
||||
height=4,
|
||||
width=4,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
input_seed=123,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
positive = pipeline.predict_noise(**positive_kwargs)
|
||||
negative = pipeline.predict_noise(**negative_kwargs)
|
||||
actual = pipeline.predict_noise_maybe_with_cfg(
|
||||
do_true_cfg=True,
|
||||
true_cfg_scale=5.0,
|
||||
positive_kwargs=positive_kwargs,
|
||||
negative_kwargs=negative_kwargs,
|
||||
cfg_normalize=False,
|
||||
)
|
||||
|
||||
expected = negative + 5.0 * (positive - negative)
|
||||
torch.testing.assert_close(actual, expected)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.cpu
|
||||
|
||||
@@ -989,6 +989,89 @@ class TestDimensionValidation:
|
||||
# Since expected_dims=3 but tensor has 4 dims, should return original
|
||||
assert result.shape == tensor_4d.shape
|
||||
|
||||
def test_auto_pad_calls_attention_mask_capability(self, monkeypatch):
|
||||
from vllm_omni.diffusion.attention import selector
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
|
||||
from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook
|
||||
|
||||
class NoMaskBackend:
|
||||
@classmethod
|
||||
def supports_attention_mask(cls, attention_spec=None):
|
||||
del attention_spec
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def get_name():
|
||||
return "NO_MASK_BACKEND"
|
||||
|
||||
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
|
||||
monkeypatch.setattr(selector, "get_attn_backend_for_capability", lambda **kwargs: NoMaskBackend)
|
||||
|
||||
hook = SequenceParallelSplitHook(
|
||||
metadata={},
|
||||
config=SequenceParallelConfig(ulysses_degree=2, ring_degree=1),
|
||||
)
|
||||
tensor = torch.randn(1, 3, 8)
|
||||
|
||||
with pytest.raises(ValueError, match="does not support attention_mask"):
|
||||
hook._shard_with_auto_pad(tensor, dim=1)
|
||||
|
||||
def test_auto_pad_allows_explicit_cudnn(self, monkeypatch):
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context
|
||||
from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook
|
||||
|
||||
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
|
||||
monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 0)
|
||||
monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
|
||||
|
||||
hook = SequenceParallelSplitHook(
|
||||
metadata={},
|
||||
config=SequenceParallelConfig(ulysses_degree=2, ring_degree=1),
|
||||
)
|
||||
tensor = torch.randn(1, 3, 8)
|
||||
od_config = OmniDiffusionConfig(
|
||||
model="test",
|
||||
dtype=torch.float32,
|
||||
diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="CUDNN_ATTN")),
|
||||
)
|
||||
|
||||
with set_forward_context(omni_diffusion_config=od_config):
|
||||
result = hook._shard_with_auto_pad(tensor, dim=1)
|
||||
|
||||
# seq_len=3 padded to 4, then sharded by world_size=2 → rank 0 gets 2.
|
||||
assert result.shape == (1, 2, 8)
|
||||
|
||||
def test_auto_pad_rejects_explicit_flashinfer_cute_on_blackwell(self, monkeypatch):
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig
|
||||
from vllm_omni.diffusion.forward_context import set_forward_context
|
||||
from vllm_omni.diffusion.hooks.sequence_parallel import SequenceParallelSplitHook
|
||||
|
||||
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 2)
|
||||
monkeypatch.setattr(parallel_state, "get_sequence_parallel_rank", lambda: 0)
|
||||
monkeypatch.setattr(parallel_state, "get_ring_parallel_world_size", lambda: 1)
|
||||
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: (10, 0))
|
||||
|
||||
hook = SequenceParallelSplitHook(
|
||||
metadata={},
|
||||
config=SequenceParallelConfig(ulysses_degree=2, ring_degree=1),
|
||||
)
|
||||
tensor = torch.randn(1, 3, 8)
|
||||
od_config = OmniDiffusionConfig(
|
||||
model="test",
|
||||
dtype=torch.float32,
|
||||
diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="FLASHINFER_ATTN")),
|
||||
)
|
||||
|
||||
with set_forward_context(omni_diffusion_config=od_config):
|
||||
with pytest.raises(ValueError, match="does not support attention_mask"):
|
||||
hook._shard_with_auto_pad(tensor, dim=1)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
class TestStrictModeSplitValidation:
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for the CustomOp platform-dispatch base class."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.layers.custom_op import CustomOp
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _NativeOnlyOp(CustomOp):
|
||||
"""An op that implements only the PyTorch-native path."""
|
||||
|
||||
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return x * 2
|
||||
|
||||
|
||||
def test_forward_xpu_falls_back_to_native():
|
||||
"""XPU has no bespoke kernels, so the base class must reuse forward_native.
|
||||
|
||||
Without this fallback every CustomOp subclass has to define an identical
|
||||
forward_xpu, and any that forgets crashes with NotImplementedError on XPU.
|
||||
"""
|
||||
op = _NativeOnlyOp()
|
||||
x = torch.randn(4, 8)
|
||||
|
||||
assert torch.equal(op.forward_xpu(x), op.forward_native(x))
|
||||
|
||||
|
||||
def test_forward_cuda_still_requires_an_implementation():
|
||||
"""The native fallback must not mask a missing CUDA kernel."""
|
||||
op = _NativeOnlyOp()
|
||||
|
||||
with pytest.raises(NotImplementedError):
|
||||
op.forward_cuda(torch.randn(4, 8))
|
||||
|
||||
|
||||
def test_mot_rmsnorm_dispatches_on_xpu_without_an_override():
|
||||
"""MoTRMSNorm relies on the inherited fallback (regression test)."""
|
||||
from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
|
||||
|
||||
assert "forward_xpu" not in MoTRMSNorm.__dict__
|
||||
|
||||
norm = MoTRMSNorm(16)
|
||||
x = torch.randn(8, 16)
|
||||
text_indices = torch.arange(0, 4)
|
||||
vae_indices = torch.arange(4, 8)
|
||||
|
||||
assert torch.equal(norm.forward_xpu(x), norm.forward_native(x))
|
||||
assert torch.equal(
|
||||
norm.forward_xpu(x, text_indices, vae_indices),
|
||||
norm.forward_native(x, text_indices, vae_indices),
|
||||
)
|
||||
Loaded 100 of 288 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user