Compare commits

...
Author SHA1 Message Date
hsliuustc0106 0527aadf6b [Bugfix] Align Hub kernel dependency guidance
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-03 23:02:19 +00:00
hsliuustc0106 6b95de36a6 Update Hub kernel dependency floor
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-03 17:04:25 +00:00
WeiQing Chen dff28e6867 [Doc] Add production diffusion model skill (#6097)
Signed-off-by: David Chen <530634352@qq.com>
2026-09-03 20:09:28 +08:00
Mu GuanLin 1dae537313 [Bugfix][LTX] Restore distilled two-stage video serving (#6847)
Signed-off-by: mglyn <1203789601@qq.com>
2026-09-03 18:48:43 +08:00
Eugene Zhao ec22c7ac06 [Bugfix] Restore Omni health route ownership (#6723)
Signed-off-by: qujing226 <qujing226@gmail.com>
Signed-off-by: Eugene Zhao <113281545+qujing226@users.noreply.github.com>
2026-09-03 00:54:36 -07:00
Jim Ban 711daacabf [Bugfix][TTS] Make CosyVoice3 TRT plan publication concurrency-safe (#6955)
Signed-off-by: BANANASJIM <bananasjim1@gmail.com>
2026-09-03 06:44:42 +00:00
Chiquan MoandMciG-ggg 2b47faf601 [Bugfix] Validate --stage-overrides shape in parse_stage_overrides (drop field allowlist) (#6230)
Signed-off-by: Chiquan Mo <mcigggg399026@gmail.com>
Signed-off-by: MciG-ggg <MciG-ggg@users.noreply.github.com>
Co-authored-by: MciG-ggg <MciG-ggg@users.noreply.github.com>
2026-09-02 23:19:08 -07:00
Tai An d1a5a7c178 [Bugfix][Qwen2.5-Omni] Fix two stale call sites in the speech helpers (#6886)
Signed-off-by: Anai-Guo <antai12232931@outlook.com>
2026-09-02 21:41:05 -07:00
Joshna-MedisettyandCursor 24c49cf392 [Bugfix][XPU] Add MoTRMSNorm.forward_xpu delegating to forward_native (#5568)
Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Signed-off-by: Joshna-Medisetty <joshna.medisetty@intel.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-02 23:06:17 -05:00
ZhanshenzhuiandSy03 48ae412a2f [model]Add experimental Mage-VL full-duplex adapter and offline integration (#6537)
Signed-off-by: Zhanshenzhui <99891002+zyforsure@users.noreply.github.com>
Co-authored-by: Sy03 <1370724210@qq.com>
2026-09-02 20:25:32 -07:00
Zeng Chuang b5e898f295 [Feat] Add support and testcase for HunyuanImage-3.0-Instruct-Distil (#4048)
Signed-off-by: zengchuang <zengchuang3@huawei.com>
Signed-off-by: zengchuang-hw <zengchuang1005@gmail.com>
2026-09-03 02:43:13 +00:00
0z5aand0z5a 4927fd5412 [Core] Optimize diffusion LoRA expand accumulation (#6268)
Signed-off-by: 0z5a <0z5a@users.noreply.github.com>
Signed-off-by: 0z5a <dinadegttyarenko6@gmail.com>
Co-authored-by: 0z5a <0z5a@users.noreply.github.com>
2026-09-03 10:37:35 +08:00
hurukawa 8f284b343a [Feature][Diffusion]: CFG parallelism support for Boogu-Image (#6786)
Signed-off-by: nagisa <1434936049@qq.com>
Signed-off-by: nagisa-kun <1434936049@qq.com>
2026-09-03 10:37:21 +08:00
Rahul SteigerandRahul Steiger 376b28aced [Bugfix][Cosmos3] Support multi-chunk transfer with distributed VAE (#6920)
Signed-off-by: Rahul Steiger <rsteiger@nvl72d147-T16.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d147-T16.cm.cluster>
2026-09-03 09:01:53 +08:00
maithilijoshi20andmaithilijoshi20 63ed0fdef4 [Bugfix] Fix vLLM serve error response imports (#6707)
Signed-off-by: maithilijoshi20 <maithilijoshi2003@gmail.com>
Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com>
Co-authored-by: maithilijoshi20 <maithilijoshi2003@gmail.com>
2026-09-02 15:48:01 -07:00
DanaerLee ca6b4e2112 [Perf][Bugfix][OmniVoice] Restore float16 serving, and fuse the generator hot loop (#6317)
Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
2026-09-02 19:54:55 +00:00
Jim Ban b9a88bf51e [Bugfix][TTS] Isolate concurrent StepAudio2 TRT builds (#6957)
Signed-off-by: BANANASJIM <bananasjim1@gmail.com>
2026-09-02 19:53:02 +00:00
GXINandHongsheng Liu 3ab582c993 [Bugfix][Qwen3-TTS] Reject task mismatches before EngineCore dispatch (#6113)
Signed-off-by: XIN GAO <1037396230@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-02 11:16:55 -07:00
3b6ce98421 [Model] Add Cache-DiT and CPU offload support for SANA-Video 2B (#5882)
Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-09-02 23:40:48 +08:00
95a9311696 [CI] Select mirror_hardwares presets via MIRROR_HW (H100/B200) (#5543)
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: [Your Name] <your.email@example.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Bo Li <22713281+bobboli@users.noreply.github.com>
2026-09-02 15:39:49 +00:00
wangyu d5fd9eb6bb [Bugfix][Diffusion] Register Sana I2V pipeline as video output (#6953)
Signed-off-by: wangyu <410167048@qq.com>
2026-09-02 23:16:58 +08:00
Bo Li 5d20f6b900 [Doc] Refine TRTLLM attention guidance (#6724)
Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
2026-09-02 19:42:37 +08:00
汪志鹏 c59b48b47e [Feature][MiniMax-H3] Complete VSA and Ulysses support (#6909)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
2026-09-02 07:19:21 -04:00
vOv 57c08d34c3 [Model] Add TP and CFG parallelism to SANA-Video 2B (#5861)
Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>
2026-09-02 17:44:41 +08:00
wangyuandamy-why-3459 b6e72f3879 [CI] Migrate L4 jobs to l4-k8s and split by GPU count (#6890)
Signed-off-by: wangyu <410167048@qq.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
2026-09-02 17:09:59 +08:00
jingchengtian 67ef0b64b9 [Hardware][Ascend] Use npu_rotary_mul fused kernel for MOSS-TTS codec RoPE (#6908)
Signed-off-by: jingchengtian <tjc1995@126.com>
2026-09-02 08:58:23 +00:00
XenorynandClaude e51fe6ec1b [Bugfix][Diffusion] Keep the GQA ratio when padding Ulysses heads (#5716)
Signed-off-by: Xenoryn <timberh2000@outlook.com>
Co-authored-by: Claude <noreply@anthropic.com>
2026-09-02 08:24:51 +08:00
leaandHongsheng Liu 5f49848536 [Bugfix][Ascend NPU][Boogu] Replace complex RoPE with real-valued operations (#6571)
Signed-off-by: lea_lei <JiaWen_Lei_7@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-02 07:11:01 +08:00
Arifuzzaman Joy 3cddda9781 [Bugfix] Prevent build_engine_args_dict from mutating stage_config.engine_args (#6783)
Signed-off-by: Arifuzzamanjoy <s1710374103@ru.ac.bd>
2026-09-01 16:45:28 +00:00
DanaerLeeandHongsheng Liu 0288c3f56e [Model] Support SenseNova-U1.5-8B-MoT and its distilled 8-step LoRA (#6516)
Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-01 22:53:53 +08:00
andyluo7andHongsheng Liu 477083a103 [CI][ROCm] Fix deterministic AMD L3 blockers (#6884)
Signed-off-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-01 22:53:00 +08:00
Rahul SteigerandRahul Steiger e1b2f589c1 [Bugfix][Cosmos3] Avoid logging prompts at INFO level (#6913)
Signed-off-by: Rahul Steiger <rsteiger@nvl72d233-T16.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d233-T16.cm.cluster>
2026-09-01 22:52:15 +08:00
Hongsheng Liuandhsliu_ustc 44c8d1c449 docs: add concise API server endpoint guide (#6114)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
2026-09-01 20:34:09 +08:00
Andy Zhou 2298359f67 [Bugfix][Diffusion] Preserve stage attention backend (#6645)
Signed-off-by: AndyZhou952 <jzhoubc@connect.ust.hk>
2026-09-01 20:18:44 +08:00
Hongsheng Liu d1c8beeb18 [Doc] Update README and docs for v0.28.0 (#6858)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-09-01 19:50:26 +08:00
WeiQing Chen e2b83db683 [Test] Add MiniMax-H3 DLO DP2 variants and Turbo LoRA L3 test (#6556)
Signed-off-by: david6666666 <530634352@qq.com>
2026-09-01 10:48:52 +00:00
AuFlowandHongsheng Liu 49c6fc1540 [Frontend] Add configurable output sample rate to Speech API (#6553)
Signed-off-by: 郑倩倩 <2458983186@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-01 17:25:09 +08:00
Alex Brooks a65e89cbd4 [Core] Validate PipelineConfig in __post_init__ (#6829)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-09-01 11:17:57 +08:00
Yuxuan RhoandHongsheng Liu 7992c2213e [Config] Remove diffusion_batch_size and use max_num_seqs instead. (#6484)
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Signed-off-by: Yuxuan Rho <104440472+LyxWxj@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-09-01 01:02:54 +00:00
XRay2016 8e26dce1f4 [Diffusion] Add Boogu Edit-Turbo support docs and tests (#6701)
Signed-off-by: xRay2016 <1150722393@qq.com>
2026-09-01 08:41:23 +08:00
288 changed files with 21369 additions and 2685 deletions

No files matched your search

+4 -7
View File
@@ -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
+434 -32
View File
@@ -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"
+3 -36
View File
@@ -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
+275 -31
View File
@@ -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):
+38
View File
@@ -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
+45 -16
View File
@@ -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"
+16 -10
View File
@@ -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"
+25 -6
View File
@@ -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.
+4
View File
@@ -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
+5 -3
View File
@@ -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()
+11
View File
@@ -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
View File
@@ -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
View File
@@ -57,9 +57,9 @@ vLLM-Omni is flexible and easy to use with:
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
- 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:
+1 -1
View File
@@ -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.
+3 -3
View File
@@ -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. |
+26 -10
View File
@@ -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
+10 -1
View File
@@ -136,7 +136,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<a href="test_writing_guide.md#l4-level-testing-full-functionality-performance-and-documentation-testing">L4</a><br>
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)).
+162 -55
View File
@@ -1,90 +1,197 @@
# Skip-Softmax
Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. Usage — the config keys
and how to pick an operating point — is in
[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax).
The shared selector contract is documented in
[Diffusion Attention Backend Selection](attention_backend_selection.md). This
page explains the algorithm.
Skip-Softmax is the sparse-attention mode of the `TRTLLM_ATTN` backend. It implements
[BLASST](https://arxiv.org/abs/2512.12087) (Dynamic BLocked Attention Sparsity via Softmax
Thresholding), which adds a per-tile skip test to the FlashAttention main loop. This page describes the algorithm, how the user configuration is resolved
into the value the kernel consumes, and how checkpoint calibration and timestep gating participate
in that resolution. Configuration keys and operating-point guidance are covered in
[TRTLLM Attention](../../user_guide/diffusion/attention_backends/trtllm.md#skip-softmax); the
backend selection contract is covered in
[Diffusion Attention Backend Selection](attention_backend_selection.md).
## Motivation
In a long attention row (a video DiT can have tens of thousands of keys), the softmax weight
concentrates on a small fraction of the keys; the rest receive near-zero weight and barely move the
output. Computing softmax and the value-weighted sum over those keys is wasted work. Skip-Softmax
detects, per block of keys, when a block cannot matter and skips its softmax and its value
multiply.
detects, per tile of keys, when a tile cannot matter and skips its softmax and its value multiply.
It is approximate: a skipped block still carries a small non-zero contribution, so the mode is
It is approximate: a skipped tile still carries a small non-zero contribution, so the mode is
opt-in and off by default.
## The online-softmax pass
## Algorithm
Attention is computed in a single streaming pass over the keys, in tiles of 128. Per query row the
![The BLASST algorithm: FlashAttention with a per-tile skip test on the local maximum](../figures/skip_softmax/blasst_algorithm.jpg)
FlashAttention computes attention in a single streaming pass over KV tiles. Per query row the
kernel maintains three running values:
- `m` — the largest score seen so far,
- `l` — the running denominator `Σ exp(sⱼ − m)`,
- `O` — the running numerator `Σ exp(sⱼ − m)·vⱼ`,
and returns `O / l` at the end. For each key tile it computes the tile's scores `Q · K_jᵀ · scale`,
updates `m`, and accumulates that tile's contribution into `l` and `O`. Rescaling `l` and `O` when
`m` grows keeps the single pass numerically exact.
and returns `O / l` at the end. For each KV tile it computes the scores `QK_j^T`, updates `m`, and
accumulates the tile's contribution into `l` and `O`. Rescaling `l` and `O` when `m` grows keeps the
dense pass numerically exact.
## The skip test
Once a tile's scores are known, its largest score `tile_max` is compared against the running
maximum:
A tile is a block of query rows against a block of keys, and the kernel processes it as a unit.
BLASST inserts one test between computing a tile's scores and accumulating them. For each query
row `i` in the tile, let `tile_max[i]` be the row's largest score within the tile, `running_max[i]`
the row's running maximum, and `λ` the effective threshold:
```text
if exp(tile_max − running_max) < threshold:
skip this tile # do not compute its softmax weights or its P·V accumulation
if all(exp(tile_max[i] - running_max[i]) < λ for every row i in the tile):
skip this tile # do not compute its Softmax or PV contribution
```
`exp(tile_max − running_max)` is an upper bound on the softmax weight any key in the tile can
receive: if even the tile's best key is far below the current maximum, every key in the tile is
negligible, and both the softmax (the exponentials) and the `P·V` accumulation for that tile can be
dropped. The tile's contribution to `l` and `O` is simply skipped.
`exp(tile_max[i] − running_max[i])` is an upper bound on the softmax weight any key in the tile can
receive in row `i`: if even the row's best key in this tile is far below the row's current maximum,
the tile is unimportant for that row. Only when this holds for every row in the tile can the
Softmax and `PV` work for the tile be skipped, in which case its contribution to `l` and `O` is
simply omitted. A larger `λ` makes the test more aggressive. The figure states the same test in
log space, `m̃ − m < ln(λ)`, which is how the kernel evaluates it.
## What this bounds
The test needs the tile's scores, so `QK_j^T` always runs and only the Softmax and `PV` work can
be skipped. Since the two matmuls have the same FLOP count, skipping every eligible tile removes at
most roughly half of the attention arithmetic, and the kernel-level speedup stays well under 2×.
Two properties of the test shape the achievable speedup:
## From configuration to the kernel threshold
- **`Q · K_jᵀ` always runs.** The test needs `tile_max`, which comes from the tile's scores, so the
score matmul is never skipped — only the softmax and the `P·V` accumulation are. The score matmul
and the value matmul are comparable in cost, so skipping every eligible tile removes at most
roughly half the attention work; the kernel-level speedup is bounded well under 2×, not unbounded.
- **The decision is per tile, not per key.** A tile is skipped only when *all* of its keys are
collectively negligible; a single important key keeps the whole tile. How many tiles actually
qualify depends on the data and rounds down to tile granularity, so `target_sparsity` selects an
operating point on a calibrated curve — it is not a promise that a fixed fraction of tiles is
skipped.
## The threshold
The per-tile threshold is normalized by sequence length:
The FlashInfer kernel does not take `λ` directly. It takes a `threshold_scale_factor` and divides
it by the KV sequence length of the call:
```text
threshold = factor / seqlen
λ = threshold_scale_factor / max_kv_len
```
`factor` comes from one of two sources:
vLLM-Omni resolves the factor from one of two mutually exclusive user controls:
- **Calibrated** — `factor = a · exp(b · target_sparsity)`, where `a` and `b` are fit per model (and
per expert, for a multi-expert model) so that a given `target_sparsity` lands near that fraction of
skipped tiles on the calibration data.
- **Direct** — `factor = skip_softmax.threshold · seqlen`, the calibration-free path; the user sets
the threshold themselves.
| Control | Factor passed to the kernel | Resulting `λ` |
| --- | --- | --- |
| `threshold` | `threshold * max_kv_len` | `threshold` |
| `target_sparsity=s` | `a * exp(b * s)` | `a * exp(b * s) / max_kv_len` |
`threshold` is therefore `λ` itself and is independent of sequence length: the same value yields
the same per-tile test at any resolution or frame count. `threshold=0` skips no tiles. Because the
left-hand side of the skip test lies in `(0, 1]`, values in `(0, 1)` are the meaningful range; the
schema accepts any finite non-negative value.
`target_sparsity` selects a point on a curve fitted per model by
[NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/diffusers/sparsity),
so that `s` lands near that fraction of skipped tiles on the calibration data. The coefficients
`a` and `b` come from the checkpoint; the achieved sparsity on another prompt, shape, or layer can
differ.
### Skip-Softmax calibration config
ModelOpt stores the calibration in each transformer component's `config.json` under
`sparse_attention_config`. The following is taken from the `transformer/config.json` of
[`nvidia/Wan2.2-T2V-A14B-Diffusers-FP8`](https://huggingface.co/nvidia/Wan2.2-T2V-A14B-Diffusers-FP8),
with the `ignore` list abridged:
```json
{
"sparse_attention_config": {
"config_groups": {
"group_0": {
"algorithm": "skip_softmax",
"targets": ["WanAttention"],
"ignore": ["blocks.0.attn1", "blocks.0.attn2", "blocks.1.attn1", "..."],
"threshold_scale_factor": {
"formula": "a * exp(b * target_sparsity)",
"coefficients": {"a": 2822.6209952221557, "b": 4.388250623520793}
},
"target_sparsity": 0.7,
"disabled_until_timestep": 0.93
}
}
}
}
```
A checkpoint with several transformer components carries one such config per component, each
with its own coefficients; in this example, `transformer_2/config.json` holds the low-noise
expert's curve.
vLLM-Omni consumes `formula`, `coefficients`, and `ignore`, and ignores `targets`,
`target_sparsity`, and `disabled_until_timestep`:
- Only the formula `a * exp(b * target_sparsity)` is accepted; a checkpoint with any other formula
is rejected at startup. If a spec requests `target_sparsity` and the checkpoint carries no
calibration, startup fails and the error points to `threshold` as the calibration-free
alternative.
- Each component's curve is applied to that component's attention layers. If the config for a
secondary component such as `transformer_2` cannot be read, that component stays dense and a
warning is logged.
- `ignore` holds fnmatch patterns matched against both the full module name and the name relative
to the component, so `blocks.0.attn1` matches `transformer.blocks.0.attn1` as well as
`transformer_2.blocks.0.attn1`. Matching layers receive no coefficients and therefore stay dense
regardless of the user configuration.
At each attention call, a layer with a `threshold` uses it directly; a layer with
`target_sparsity` and stamped coefficients evaluates the curve; a layer with `target_sparsity` but
no coefficients runs dense. The timestep gate below is applied to the result before it reaches the
kernel.
## Timestep gating
`disabled_until_timestep = D` keeps the mode off during the early, high-noise denoise steps and
turns it on once the normalized timestep `t` drops to `t ≤ D` (`t` runs `1.0` → `0.0` over the
schedule). The early steps set the global structure of the output and their errors propagate through
every later step, so keeping them dense costs a few skipped-tile opportunities but protects fidelity.
The early, high-noise denoising steps set the global structure of the output, and their errors
propagate through every later step. Keeping just these steps dense markedly improves the quality
of the generated video at a small cost in skipped work. `disabled_until_timestep` implements this:
on each attention call the backend compares the current normalized timestep `t` against the
configured cutoff and passes no skip factor to the kernel while `t` is above it. Skip-Softmax
becomes active on the first step whose `t` is at or below the cutoff and stays active for the rest
of the run, since `t` only decreases.
`t` is the scheduler's own timestep divided by `num_train_timesteps`, published by the pipeline via
`DenoiseProgressMixin.record_denoise_step`. It is deliberately not derived from the step index,
because schedulers space their steps non-uniformly. A pipeline that does not publish a timestep
stays dense when `disabled_until_timestep` is set.
The default `disabled_until_timestep = 0` is a sentinel rather than a cutoff: the gate is off, the
timestep is never read, and Skip-Softmax runs on every step. Any positive value enables the gate.
`disabled_until_timestep = 1.0` therefore also yields no dense steps on a pipeline that publishes
`t`, but on a pipeline that does not publish it the backend stays dense and logs a warning once
rather than guessing from the step index.
### Where `t` comes from
The pipeline publishes `t` once per denoising step before running the transformer. It is the
scheduler's own position in `[0, 1]`, decreasing from near `1.0` (pure noise) to `0.0` (clean
sample): scheduler-based pipelines publish the scheduler timestep divided by `num_train_timesteps`,
and rectified-flow pipelines publish the current sigma. `t` is deliberately not the step index
divided by the step count. Schedulers place their steps non-uniformly in `t`, so expressing the
cutoff in `t` ties it to the same noise level across step counts and schedulers.
### Mapping the cutoff to denoising steps
For a run whose published sequence is `t[0], ..., t[N-1]`, the number of dense steps is
```text
dense_steps = count(t[i] > disabled_until_timestep)
skip_softmax_steps = N - dense_steps
```
This count **depends on the schedule**, and schedules differ from model to model: the
scheduler family, the number of steps, and any time shift all change where the steps land in `t`.
The same cutoff can therefore gate a very different fraction of the run on two models. Take
MiniMax-H3 as an example. Its video branch uses a flow-shifted rectified-flow schedule that applies
a shift `s` to `N + 1` uniform positions `u` from `1` to `0`,
```text
t = s * u / (1 + (s - 1) * u)
```
and a large `s` pushes most positions toward `t ≈ 1`. With the default `s = 12` and
`num_inference_steps = 50` (50 sigma points, `N = 49` denoiser forwards), the published sequence
starts `1.000, 0.998, 0.996, 0.995, 0.993, ...` and stays above `0.9` for 28 steps, giving:
| `disabled_until_timestep` | Dense steps | Skip-Softmax steps |
| :---: | ---: | ---: |
| `1.00` | 0 | 49 |
| `0.99` | 6 | 43 |
| `0.97` | 14 | 35 |
| `0.95` | 19 | 30 |
| `0.90` | 28 | 21 |
| `0.86` | 33 | 16 |
These numbers hold only for that schedule. A model with a smaller shift, a different step count,
or a scheduler that is not rectified flow produces a different sequence and a different table; a
request that overrides `flow_shift` or the step count changes it for MiniMax-H3 too.
Choose the cutoff by counting against the schedule actually served, not by reusing a value from
another model.
Binary file not shown.

After

Width:  |  Height:  |  Size: 562 KiB

+3 -1
View File
@@ -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.
+1 -1
View File
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
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
+10 -14
View File
@@ -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).
+3 -3
View File
@@ -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` | ✅︎ | ✅︎ | | | — |
+187
View File
@@ -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.
+3 -3
View File
@@ -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 \
+86 -54
View File
@@ -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)
+116
View File
@@ -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.
+109
View File
@@ -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.
+82
View File
@@ -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
View File
@@ -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
+109
View File
@@ -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.
---
+3 -2
View File
@@ -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()
+1 -1
View File
@@ -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
View File
@@ -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.
+114
View File
@@ -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
```
+54 -21
View File
@@ -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:
+101 -2
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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 |
+159
View File
@@ -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.
+3 -2
View File
@@ -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

+42
View File
@@ -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
View File
@@ -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
+176 -3
View File
@@ -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()
+142 -43
View File
@@ -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
+2 -2
View File
@@ -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,
+6 -1
View File
@@ -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 = []
+11 -22
View File
@@ -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
+72 -38
View File
@@ -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()
+203 -1
View File
@@ -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
+66 -17
View File
@@ -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)
+96 -23
View File
@@ -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)
+40 -4
View File
@@ -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")
+449 -39
View File
@@ -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
View File
@@ -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:
+55
View File
@@ -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