anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge origin/main into dev/vllm-align

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouamandCursor committed 2026-02-12 03:41:38 +00:00
2 parents 0099eebb3f 20a8fe1dcc
commit 50f2b1e854
81 files changed
+2504 -1573
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/pipeline.yml
+75 -75
View File
@@ -9,13 +9,13 @@ steps:
agents:
queue: "cpu_queue_premerge"
# - label: "Test on NPU"
# depends_on: ~
# key: npu-test
# commands:
# - ".buildkite/scripts/hardware_ci/run_npu_test.sh"
# agents:
# queue: "ascend"
- label: "Upload Nightly Pipeline"
depends_on: image-build
if: build.env("NIGHTLY") == "1"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly.yaml
agents:
queue: "cpu_queue_premerge"
- label: "Simple Unit Test"
depends_on: image-build
@@ -174,12 +174,13 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Benchmark Test"
- label: "Benchmark&Engine Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/benchmarks/test_serve_cli.py
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
agents:
queue: "mithril-h100-pool"
plugins:
@@ -213,78 +214,77 @@ steps:
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Qwen3-TTS E2E Test"
# timeout_in_minutes: 10
# depends_on: image-build
# commands:
# - export VLLM_LOGGING_LEVEL=DEBUG
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
# agents:
# queue: "gpu_4_queue"
# plugins:
# - docker#v5.2.0:
# image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# always-pull: true
# propagate-environment: true
# shm-size: "8gb"
# environment:
# - "HF_HOME=/fsx/hf_cache"
# volumes:
# - "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Omni Model Test with H100"
# timeout_in_minutes: 30
# depends_on: image-build
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
# - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
# - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
# - pytest -s -v tests/e2e/online_serving/test_async_omni.py
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 2
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
- label: "Qwen3-TTS E2E Test"
timeout_in_minutes: 10
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
timeout_in_minutes: 20
.buildkite/test-amd.yaml
+1 -1
View File
@@ -101,7 +101,7 @@ steps:
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_async_omni.py
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
- label: "Diffusion Image Edit Test"
timeout_in_minutes: 15
.buildkite/test-merge.yml
+329
View File
@@ -0,0 +1,329 @@
steps:
- label: ":docker: Build image"
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
- label: "Simple Unit Test"
depends_on: image-build
commands:
- "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Images API LoRA E2E"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model CPU offloading Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Audio Generation Model Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Sequence Parallelism Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Tensor Parallelism Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/diffusion/test_diffusion_worker.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Benchmark&Engine Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/benchmarks/test_serve_cli.py
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni Model Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 30
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Bagel Text2Img Model Test with H100"
# timeout_in_minutes: 30
# depends_on: image-build
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
.buildkite/test-nightly.yaml
+12 -23
View File
@@ -1,14 +1,4 @@
steps:
- label: ":docker: Build image"
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
- label: "Omni Model Test with H100"
timeout_in_minutes: 60
depends_on: image-build
@@ -46,23 +36,22 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni Model Test"
timeout_in_minutes: 60
depends_on: image-build
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
benchmarks/build_dataset/download_process_data_seedtts.md
+1 -1
View File
@@ -56,7 +56,7 @@ python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
**Options:**
- `-i, --input`: Input metadata file (default: `meta.lst`)
- `-o, --output`: Output prompts file (default: `prompts.txt`)
- `-n, --num_lines`: Number of prompts to extract (required)
- `-n, --num-lines`: Number of prompts to extract (required)
### 7. Clean Up (Optional)
benchmarks/build_dataset/extract_tts_prompts.py
+1 -1
View File
@@ -56,7 +56,7 @@ def main():
"-o", "--output", type=str, default="prompts.txt", help="Output txt file path (default: prompts.txt)"
)
parser.add_argument(
"-n", "--num_lines", type=int, required=True, help="Number of lines to extract from the beginning"
"-n", "--num-lines", type=int, required=True, help="Number of lines to extract from the beginning"
)
args = parser.parse_args()
benchmarks/qwen3-omni/README.md
+1 -1
View File
@@ -36,7 +36,7 @@ bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
```
What it does:
- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num_prompts 100`.
- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num-prompts 100`.
- Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
- `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
- `results.json` — per-prompt outputs and audio paths.
benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
+1 -1
View File
@@ -16,7 +16,7 @@ if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]
else
cd benchmarks/qwen3-omni/transformers
python qwen3_omni_moe_transformers.py --prompts_file ../../build_dataset/top100.txt --num_prompts 100
python qwen3_omni_moe_transformers.py --prompts-file ../../build_dataset/top100.txt --num-prompts 100
echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
echo " - perf_stats.json Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py
+4 -4
View File
@@ -227,17 +227,17 @@ def print_stats(stats: dict):
def main():
parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
parser.add_argument(
"--prompts_file",
"--prompts-file",
type=str,
default="benchmark/build_dataset/top100.txt",
help="Path to the prompts file (one prompt per line)",
)
parser.add_argument(
"--output_dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
"--output-dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
)
parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Path to the model")
parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Path to the model")
parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
parser.add_argument("--num_prompts", type=int, default=None, help="Number of prompts to process (default: all)")
parser.add_argument("--num-prompts", type=int, default=None, help="Number of prompts to process (default: all)")
args = parser.parse_args()
# Load model and processor
docs/.nav.yml
+3
View File
@@ -43,6 +43,9 @@ nav:
- Overview: user_guide/diffusion_acceleration.md
- TeaCache: user_guide/diffusion/teacache.md
- Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
- Quantization:
- Overview: user_guide/diffusion/quantization/overview.md
- FP8: user_guide/diffusion/quantization/fp8.md
- Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
- ComfyUI: features/comfyui.md
docs/contributing/features/cfg_parallel.md
+5 -5
View File
@@ -188,11 +188,11 @@ cd examples/offline_inference/text_to_image
python text_to_image.py \
--model Your-org/your-model \
--prompt "a cup of coffee on the table" \
--negative_prompt "ugly, unclear" \
--cfg_scale 4.0 \
--num_inference_steps 50 \
--negative-prompt "ugly, unclear" \
--cfg-scale 4.0 \
--num-inference-steps 50 \
--output "cfg_enabled.png" \
--cfg_parallel_size 2
--cfg-parallel-size 2
```
**Verify:**
@@ -246,4 +246,4 @@ Adding CFG-Parallel support:
1. ✅ **Create mixin** - Inherit from `CFGParallelMixin` and implement `diffuse()` method
2. ✅ **(Optional) Customize** - Override `predict_noise()` or `cfg_normalize_function()` for custom behavior
3. ✅ **Test** - Verify with `--cfg_parallel_size 2` and compare performance
3. ✅ **Test** - Verify with `--cfg-parallel-size 2` and compare performance
docs/contributing/features/sequence_parallel.md
+6 -6
View File
@@ -277,9 +277,9 @@ cd examples/offline_inference/text_to_image
python text_to_image.py \
--model Your-org/your-model \
--prompt "a cup of coffee on the table" \
--num_inference_steps 50 \
--ulysses_degree 2 \
--ring_degree 2 \
--num-inference-steps 50 \
--ulysses-degree 2 \
--ring-degree 2 \
--output sp_test_image_ulysses=2_ring=2.png
```
@@ -298,9 +298,9 @@ cd examples/offline_inference/text_to_image
python text_to_image.py \
--model Your-org/your-model \
--prompt "a cup of coffee on the table" \
--num_inference_steps 50 \
--ulysses_degree 2 \
--tensor_parallel_size 2 \
--num-inference-steps 50 \
--ulysses-degree 2 \
--tensor-parallel-size 2 \
--output sp_test_image_ulysses=2_tp=2.png
```
docs/contributing/features/tensor_parallel.md
+4 -4
View File
@@ -181,11 +181,11 @@ cd examples/offline_inference/text_to_image
python text_to_image.py \
--model Your-org/your-model \
--prompt "a cup of coffee on the table" \
--negative_prompt "ugly, unclear" \
--cfg_scale 4.0 \
--num_inference_steps 50 \
--negative-prompt "ugly, unclear" \
--cfg-scale 4.0 \
--num-inference-steps 50 \
--output "tp_enabled.png" \
--tensor_parallel_size 2
--tensor-parallel-size 2
```
**Verify:**
docs/contributing/profiling.md
+6 -6
View File
@@ -106,7 +106,7 @@ python image_to_video.py \
# Reducing frames to the absolute minimum (2) keeps the
# tensor size small, ensuring the trace file doesn't become
# multi-gigabytes in size.
--num_frames 2 \
--num-frames 2 \
\
# Minimize Iteration Loop (Steps):
# This is the most critical setting for profiling.
@@ -114,12 +114,12 @@ python image_to_video.py \
# Profiling 2 steps gives you the exact same performance
# data as 50 steps, but saves minutes of runtime and
# prevents the trace viewer from freezing.
--num_inference_steps 2 \
--num-inference-steps 2 \
\
--guidance_scale 5.0 \
--guidance_scale_high 6.0 \
--boundary_ratio 0.875 \
--flow_shift 12.0 \
--guidance-scale 5.0 \
--guidance-scale-high 6.0 \
--boundary-ratio 0.875 \
--flow-shift 12.0 \
--fps 16 \
--output i2v_output.mp4
docs/user_guide/diffusion/cache_dit_acceleration.md
+2 -2
View File
@@ -188,8 +188,8 @@ cd examples/offline_inference/text_to_image
python text_to_image.py \
--model Qwen/Qwen-Image \
--prompt "a cup of coffee on the table" \
--cache_backend cache_dit \
--num_inference_steps 50
--cache-backend cache_dit \
--num-inference-steps 50
```
docs/user_guide/diffusion/parallelism_acceleration.md
+3 -3
View File
@@ -317,10 +317,10 @@ python image_edit.py \
--model "Qwen/Qwen-Image-Edit" \
--image "qwen_image_output.png" \
--prompt "turn this cat to a dog" \
--negative_prompt "low quality, blurry" \
--cfg_scale 4.0 \
--negative-prompt "low quality, blurry" \
--cfg-scale 4.0 \
--output "edited_image.png" \
--cfg_parallel_size 2
--cfg-parallel-size 2
```
#### Online Serving
docs/user_guide/diffusion/quantization/fp8.md
+77
View File
@@ -0,0 +1,77 @@
# FP8 Quantization
## Overview
FP8 quantization converts BF16/FP16 weights to FP8 at model load time. No calibration or pre-quantized checkpoint needed.
Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16. See the [per-model table](#supported-models) for which case applies.
Common sensitive layers in DiT-based diffusion models include **image-stream MLPs** (`img_mlp`). These are particularly vulnerable to FP8 precision loss because they process denoising latents whose dynamic range shifts significantly across timesteps, and unlike attention projections (which benefit from QK-Norm stabilization), MLPs have no built-in normalization to absorb quantization error. In deep architectures (e.g., 60+ residual blocks), small per-layer errors compound and degrade output quality. Other layers such as **attention projections** (`to_qkv`, `to_out`) and **text-stream MLPs** (`txt_mlp`) are generally more robust due to normalization or more stable input statistics.
## Configuration
1. **Python API**: set `quantization="fp8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# All layers quantized
omni = Omni(model="<your-model>", quantization="fp8")
# Skip sensitive layers
omni = Omni(
model="<your-model>",
quantization_config={
"method": "fp8",
"ignored_layers": ["<layer-name>"],
},
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
2. **CLI**: pass `--quantization fp8` and optionally `--ignored-layers`.
```bash
# All layers
python text_to_image.py --model <your-model> --quantization fp8
# Skip sensitive layers
python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers "img_mlp"
# Online serving
vllm serve <your-model> --omni --quantization fp8
```
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `method` | str | — | Quantization method (`"fp8"`) |
| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) or `"static"` |
| `weight_block_size` | list[int] \| None | `None` | Block size for block-wise weight quantization |
The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
## Supported Models
| Model | HF Models | Recommendation | `ignored_layers` |
|-------|-----------|---------------|------------------|
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
## Combining with Other Features
FP8 quantization can be combined with cache acceleration:
```python
omni = Omni(
model="<your-model>",
quantization="fp8",
cache_backend="tea_cache",
cache_config={"rel_l1_thresh": 0.2},
)
```
docs/user_guide/diffusion/quantization/overview.md
+17
View File
@@ -0,0 +1,17 @@
# Quantization for Diffusion Transformers
vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and accelerate inference.
## Supported Methods
| Method | Guide |
|--------|-------|
| FP8 | [FP8](fp8.md) |
## Device Compatibility
| GPU Generation | Example GPUs | FP8 Mode |
|---------------|-------------------|----------|
| Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
Kernel selection is automatic.
docs/user_guide/diffusion_acceleration.md
+38 -2
View File
@@ -1,6 +1,6 @@
# Diffusion Acceleration Overview
vLLM-Omni supports various cache acceleration methods to speed up diffusion model inference with minimal quality degradation. These methods include **cache methods** that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, and **parallelism methods** that distribute the computation across multiple devices.
vLLM-Omni supports various acceleration methods to speed up diffusion model inference with minimal quality degradation. These include **cache methods** that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, **parallelism methods** that distribute the computation across multiple devices, and **quantization methods** that reduce memory footprint while preserving accuracy.
## Supported Acceleration Methods
@@ -14,6 +14,10 @@ vLLM-Omni currently supports two main cache acceleration backends:
Both methods can provide significant speedups (typically **1.5x-2.0x**) while maintaining high output quality.
vLLM-Omni also supports quantization methods:
3. **[FP8 Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
vLLM-Omni also supports parallelism methods for diffusion models, including:
1. [Ulysses-SP](diffusion/parallelism_acceleration.md#ulysses-sp) - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
@@ -35,6 +39,12 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
| **TeaCache** | `cache_backend="tea_cache"` | Simple, adaptive caching with minimal configuration | Quick setup, balanced speed/quality |
| **Cache-DiT** | `cache_backend="cache_dit"` | Advanced caching with multiple techniques (DBCache, TaylorSeer, SCM) | Maximum acceleration, fine-grained control |
### Quantization Methods
| Method | Configuration | Description | Best For |
|--------|--------------|-------------|----------|
| **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
## Supported Models
The following table shows which models are currently supported by each acceleration method:
@@ -58,10 +68,18 @@ The following table shows which models are currently supported by each accelerat
### VideoGen
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |CFG-Parallel |
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
### Quantization
| Model | Model Identifier | FP8 |
|-------|------------------|:---:|
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ |
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ |
## Performance Benchmarks
@@ -272,12 +290,30 @@ outputs = omni.generate(
)
```
### Using FP8 Quantization
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="<your-model>",
quantization="fp8",
)
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
## Documentation
For detailed information on each acceleration method:
- **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
- **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
- **[FP8 Quantization Guide](diffusion/quantization/overview.md)** - FP8 quantization for DiT models with per-layer control
- **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models.
- **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
- **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
docs/user_guide/examples/offline_inference/image_to_image.md
+13 -13
View File
@@ -22,8 +22,8 @@ python image_edit.py \
--image qwen-bear.png \
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0
--num-inference-steps 50 \
--cfg-scale 4.0
```
### Multiple Image Editing (Qwen-Image-Edit-2509)
@@ -36,26 +36,26 @@ python image_edit.py \
--image img1.png img2.png \
--prompt "Combine these images into a single scene" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0 \
--guidance_scale 1.0
--num-inference-steps 50 \
--cfg-scale 4.0 \
--guidance-scale 1.0
```
Key arguments:
- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support.
- `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
- `--prompt` / `--negative_prompt`: text description (string).
- `--cfg_scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
- `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--prompt` / `--negative-prompt`: text description (string).
- `--cfg-scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
- `--guidance-scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--output`: path to save the generated PNG.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
## Example materials
docs/user_guide/examples/offline_inference/image_to_video.md
+22 -22
View File
@@ -20,15 +20,15 @@ python image_to_video.py \
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
--image cherry_blossom.jpg \
--prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
--negative_prompt "<optional quality filter>" \
--negative-prompt "<optional quality filter>" \
--height 480 \
--width 832 \
--num_frames 48 \
--guidance_scale 5.0 \
--guidance_scale_high 6.0 \
--num_inference_steps 40 \
--boundary_ratio 0.875 \
--flow_shift 12.0 \
--num-frames 48 \
--guidance-scale 5.0 \
--guidance-scale-high 6.0 \
--num-inference-steps 40 \
--boundary-ratio 0.875 \
--flow-shift 12.0 \
--fps 16 \
--output i2v_output.mp4
```
@@ -40,13 +40,13 @@ python image_to_video.py \
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--image cherry_blossom.jpg \
--prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
--negative_prompt "<optional quality filter>" \
--negative-prompt "<optional quality filter>" \
--height 480 \
--width 832 \
--num_frames 48 \
--guidance_scale 4.0 \
--num_inference_steps 40 \
--flow_shift 12.0 \
--num-frames 48 \
--guidance-scale 4.0 \
--num-inference-steps 40 \
--flow-shift 12.0 \
--fps 16 \
--output i2v_output.mp4
```
@@ -57,20 +57,20 @@ Key arguments:
- `--image`: Path to input image (required).
- `--prompt`: Text description of desired motion/animation.
- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
- `--num_frames`: Number of frames (default 81).
- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high-noise stages for MoE).
- `--negative_prompt`: Optional list of artifacts to suppress.
- `--boundary_ratio`: Boundary split ratio for two-stage MoE models.
- `--flow_shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
- `--num_inference_steps`: Number of denoising steps (default 50).
- `--num-frames`: Number of frames (default 81).
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
- `--negative-prompt`: Optional list of artifacts to suppress.
- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
- `--num-inference-steps`: Number of denoising steps (default 50).
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--output`: Path to save the generated video.
- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
## Example materials
docs/user_guide/examples/offline_inference/lora_inference.md
+4 -4
View File
@@ -26,7 +26,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--lora-path /path/to/lora/ \
--lora-scale 1.0 \
--num_inference_steps 50 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_preloaded.png
@@ -43,7 +43,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--lora-request-path /path/to/lora/ \
--lora-scale 1.0 \
--num_inference_steps 50 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_per_request.png
@@ -56,7 +56,7 @@ If no LoRA request is provided, we will use the base model without any LoRA adap
```bash
python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--num_inference_steps 50 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_no_lora.png
@@ -77,7 +77,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
- `--seed`: Random seed for reproducibility (default: 42)
- `--height`: Image height in pixels (default: 1024)
- `--width`: Image width in pixels (default: 1024)
- `--num_inference_steps`: Number of denoising steps (default: 50)
- `--num-inference-steps`: Number of denoising steps (default: 50)
- `--output`: Output file path (default: `lora_output.png`)
## How LoRA Works
docs/user_guide/examples/offline_inference/text_to_audio.md
+8 -8
View File
@@ -21,22 +21,22 @@ If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you
python text_to_audio.py \
--model stabilityai/stable-audio-open-1.0 \
--prompt "The sound of a hammer hitting a wooden surface" \
--negative_prompt "Low quality" \
--negative-prompt "Low quality" \
--seed 42 \
--guidance_scale 7.0 \
--audio_length 10.0 \
--num_inference_steps 100 \
--guidance-scale 7.0 \
--audio-length 10.0 \
--num-inference-steps 100 \
--output stable_audio_output.wav
```
Key arguments:
- `--prompt`: text description (string).
- `--negative_prompt`: negative prompt for classifier-free guidance.
- `--negative-prompt`: negative prompt for classifier-free guidance.
- `--seed`: integer seed for deterministic generation.
- `--guidance_scale`: classifier-free guidance scale.
- `--audio_length`: audio duration in seconds.
- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
- `--guidance-scale`: classifier-free guidance scale.
- `--audio-length`: audio duration in seconds.
- `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
- `--output`: path to save the generated WAV file.
## Example materials
docs/user_guide/examples/offline_inference/text_to_image.md
+10 -10
View File
@@ -82,9 +82,9 @@ python text_to_image.py \
--model Tongyi-MAI/Z-Image-Turbo \
--prompt "a cup of coffee on the table" \
--seed 42 \
--cfg_scale 4.0 \
--num_images_per_prompt 1 \
--num_inference_steps 50 \
--cfg-scale 4.0 \
--num-images-per-prompt 1 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output outputs/coffee.png
@@ -94,17 +94,17 @@ Key arguments:
- `--prompt`: text description (string).
- `--seed`: integer seed for deterministic sampling.
- `--cfg_scale`: true CFG scale (model-specific guidance strength).
- `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--cfg-scale`: true CFG scale (model-specific guidance strength).
- `--num-images-per-prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--height/--width`: output resolution (defaults 1024x1024).
- `--output`: path to save the generated PNG.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
docs/user_guide/examples/offline_inference/text_to_video.md
+14 -14
View File
@@ -10,14 +10,14 @@ The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text
```bash
python text_to_video.py \
--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
--negative_prompt "<optional quality filter>" \
--negative-prompt "<optional quality filter>" \
--height 480 \
--width 832 \
--num_frames 33 \
--guidance_scale 4.0 \
--guidance_scale_high 3.0 \
--flow_shift 12.0 \
--num_inference_steps 40 \
--num-frames 33 \
--guidance-scale 4.0 \
--guidance-scale-high 3.0 \
--flow-shift 12.0 \
--num-inference-steps 40 \
--fps 16 \
--output t2v_out.mp4
```
@@ -26,18 +26,18 @@ Key arguments:
- `--prompt`: text description (string).
- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
- `--num_frames`: Number of frames (Wan default is 81).
- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
- `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
- `--num-frames`: Number of frames (Wan default is 81).
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--output`: path to save the generated video.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
## Example materials
examples/offline_inference/bagel/end2end.py
+1 -1
View File
@@ -19,7 +19,7 @@ def parse_args():
default=None,
help="Path to a .txt file with one prompt per line (preferred).",
)
parser.add_argument("--prompt_type", default="text", choices=["text"])
parser.add_argument("--prompt-type", default="text", choices=["text"])
parser.add_argument(
"--modality",
examples/offline_inference/image_to_image/image_edit.py
+44 -44
View File
@@ -9,34 +9,34 @@ Usage (single image):
--image input.png \
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0 \
--guidance_scale 1.0
--num-inference-steps 50 \
--cfg-scale 4.0 \
--guidance-scale 1.0
Usage (multiple images):
python image_edit.py \
--image input1.png input2.png input3.png \
--prompt "Combine these images into a single scene" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0 \
--guidance_scale 1.0
--num-inference-steps 50 \
--cfg-scale 4.0 \
--guidance-scale 1.0
Usage (with cache-dit acceleration):
python image_edit.py \
--image input.png \
--prompt "Edit description" \
--cache_backend cache_dit \
--cache_dit_max_continuous_cached_steps 3 \
--cache_dit_residual_diff_threshold 0.24 \
--cache_dit_enable_taylorseer
--cache-backend cache_dit \
--cache-dit-max-continuous-cached-steps 3 \
--cache-dit-residual-diff-threshold 0.24 \
--cache-dit-enable-taylorseer
Usage (with tea_cache acceleration):
python image_edit.py \
--image input.png \
--prompt "Edit description" \
--cache_backend tea_cache \
--tea_cache_rel_l1_thresh 0.25
--cache-backend tea_cache \
--tea-cache-rel-l1-thresh 0.25
Usage (layered):
python image_edit.py \
@@ -44,8 +44,8 @@ Usage (layered):
--image input.png \
--prompt "" \
--output "layered" \
--num_inference_steps 50 \
--cfg_scale 4.0 \
--num-inference-steps 50 \
--cfg-scale 4.0 \
--layers 4 \
--color-format "RGBA"
@@ -53,17 +53,17 @@ Usage (with CFG Parallel):
python image_edit.py \
--image input.png \
--prompt "Edit description" \
--cfg_parallel_size 2 \
--num_inference_steps 50 \
--cfg_scale 4.0
--cfg-parallel-size 2 \
--num-inference-steps 50 \
--cfg-scale 4.0
Usage (disable torch.compile):
python image_edit.py \
--image input.png \
--prompt "Edit description" \
--enforce_eager \
--num_inference_steps 50 \
--cfg_scale 4.0
--enforce-eager \
--num-inference-steps 50 \
--cfg-scale 4.0
For more options, run:
python image_edit.py --help
@@ -109,7 +109,7 @@ def parse_args() -> argparse.Namespace:
help="Text prompt describing the edit to make to the image.",
)
parser.add_argument(
"--negative_prompt",
"--negative-prompt",
type=str,
default=None,
required=False,
@@ -121,7 +121,7 @@ def parse_args() -> argparse.Namespace:
help="Random seed for deterministic results.",
)
parser.add_argument(
"--cfg_scale",
"--cfg-scale",
type=float,
default=4.0,
help=(
@@ -132,12 +132,12 @@ def parse_args() -> argparse.Namespace:
),
)
parser.add_argument(
"--guidance_scale",
"--guidance-scale",
type=float,
default=1.0,
help=(
"Guidance scale for guidance-distilled models (default: 1.0, disabled). "
"Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale "
"Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale "
"directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models."
),
)
@@ -148,19 +148,19 @@ def parse_args() -> argparse.Namespace:
help=("Path to save the edited image (PNG). Or prefix for Qwen-Image-Layered model save images(PNG)."),
)
parser.add_argument(
"--num_outputs_per_prompt",
"--num-outputs-per-prompt",
type=int,
default=1,
help="Number of images to generate for the given prompt.",
)
parser.add_argument(
"--num_inference_steps",
"--num-inference-steps",
type=int,
default=50,
help="Number of denoising steps for the diffusion sampler.",
)
parser.add_argument(
"--cache_backend",
"--cache-backend",
type=str,
default=None,
choices=["cache_dit", "tea_cache"],
@@ -171,19 +171,19 @@ def parse_args() -> argparse.Namespace:
),
)
parser.add_argument(
"--ulysses_degree",
"--ulysses-degree",
type=int,
default=1,
help="Number of GPUs used for ulysses sequence parallelism.",
)
parser.add_argument(
"--ring_degree",
"--ring-degree",
type=int,
default=1,
help="Number of GPUs used for ring sequence parallelism.",
)
parser.add_argument(
"--tensor_parallel_size",
"--tensor-parallel-size",
type=int,
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
@@ -205,56 +205,56 @@ def parse_args() -> argparse.Namespace:
# Cache-DiT specific parameters
parser.add_argument(
"--cache_dit_fn_compute_blocks",
"--cache-dit-fn-compute-blocks",
type=int,
default=1,
help="[cache-dit] Number of forward compute blocks. Optimized for single-transformer models.",
)
parser.add_argument(
"--cache_dit_bn_compute_blocks",
"--cache-dit-bn-compute-blocks",
type=int,
default=0,
help="[cache-dit] Number of backward compute blocks.",
)
parser.add_argument(
"--cache_dit_max_warmup_steps",
"--cache-dit-max-warmup-steps",
type=int,
default=4,
help="[cache-dit] Maximum warmup steps (works for few-step models).",
)
parser.add_argument(
"--cache_dit_residual_diff_threshold",
"--cache-dit-residual-diff-threshold",
type=float,
default=0.24,
help="[cache-dit] Residual diff threshold. Higher values enable more aggressive caching.",
)
parser.add_argument(
"--cache_dit_max_continuous_cached_steps",
"--cache-dit-max-continuous-cached-steps",
type=int,
default=3,
help="[cache-dit] Maximum continuous cached steps to prevent precision degradation.",
)
parser.add_argument(
"--cache_dit_enable_taylorseer",
"--cache-dit-enable-taylorseer",
action="store_true",
default=False,
help="[cache-dit] Enable TaylorSeer acceleration (not suitable for few-step models).",
)
parser.add_argument(
"--cache_dit_taylorseer_order",
"--cache-dit-taylorseer-order",
type=int,
default=1,
help="[cache-dit] TaylorSeer polynomial order.",
)
parser.add_argument(
"--cache_dit_scm_steps_mask_policy",
"--cache-dit-scm-steps-mask-policy",
type=str,
default=None,
choices=[None, "slow", "medium", "fast", "ultra"],
help="[cache-dit] SCM mask policy: None (disabled), slow, medium, fast, ultra.",
)
parser.add_argument(
"--cache_dit_scm_steps_policy",
"--cache-dit-scm-steps-policy",
type=str,
default="dynamic",
choices=["dynamic", "static"],
@@ -263,30 +263,30 @@ def parse_args() -> argparse.Namespace:
# TeaCache specific parameters
parser.add_argument(
"--tea_cache_rel_l1_thresh",
"--tea-cache-rel-l1-thresh",
type=float,
default=0.2,
help="[tea_cache] Threshold for accumulated relative L1 distance.",
)
parser.add_argument(
"--cfg_parallel_size",
"--cfg-parallel-size",
type=int,
default=1,
choices=[1, 2],
help="Number of GPUs used for classifier free guidance parallel size.",
)
parser.add_argument(
"--enforce_eager",
"--enforce-eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
parser.add_argument(
"--vae_use_slicing",
"--vae-use-slicing",
action="store_true",
help="Enable VAE slicing for memory optimization.",
)
parser.add_argument(
"--vae_use_tiling",
"--vae-use-tiling",
action="store_true",
help="Enable VAE tiling for memory optimization.",
)
examples/offline_inference/image_to_image/image_to_image.md
+13 -13
View File
@@ -19,8 +19,8 @@ python image_edit.py \
--image qwen-bear.png \
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0
--num-inference-steps 50 \
--cfg-scale 4.0
```
### Multiple Image Editing (Qwen-Image-Edit-2509)
@@ -33,23 +33,23 @@ python image_edit.py \
--image img1.png img2.png \
--prompt "Combine these images into a single scene" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0 \
--guidance_scale 1.0
--num-inference-steps 50 \
--cfg-scale 4.0 \
--guidance-scale 1.0
```
Key arguments:
- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support.
- `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
- `--prompt` / `--negative_prompt`: text description (string).
- `--cfg_scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
- `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--prompt` / `--negative-prompt`: text description (string).
- `--cfg-scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
- `--guidance-scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--output`: path to save the generated PNG.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
examples/offline_inference/image_to_image/run_qwen_image_edit_2511.sh
+3 -3
View File
@@ -3,6 +3,6 @@ python image_edit.py \
--image qwen_bear.png \
--prompt "Add a white art board written with colorful text 'vLLM-Omni' on grassland. Add a paintbrush in the bear's hands. position the bear standing in front of the art board as if painting" \
--output output_image_edit.png \
--num_inference_steps 50 \
--cfg_scale 4.0 \
--cache_backend cache_dit \
--num-inference-steps 50 \
--cfg-scale 4.0 \
--cache-backend cache_dit \
examples/offline_inference/image_to_video/README.md
+22 -22
View File
@@ -17,15 +17,15 @@ python image_to_video.py \
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
--image cherry_blossom.jpg \
--prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
--negative_prompt "<optional quality filter>" \
--negative-prompt "<optional quality filter>" \
--height 480 \
--width 832 \
--num_frames 48 \
--guidance_scale 5.0 \
--guidance_scale_high 6.0 \
--num_inference_steps 40 \
--boundary_ratio 0.875 \
--flow_shift 12.0 \
--num-frames 48 \
--guidance-scale 5.0 \
--guidance-scale-high 6.0 \
--num-inference-steps 40 \
--boundary-ratio 0.875 \
--flow-shift 12.0 \
--fps 16 \
--output i2v_output.mp4
```
@@ -37,13 +37,13 @@ python image_to_video.py \
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--image cherry_blossom.jpg \
--prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
--negative_prompt "<optional quality filter>" \
--negative-prompt "<optional quality filter>" \
--height 480 \
--width 832 \
--num_frames 48 \
--guidance_scale 4.0 \
--num_inference_steps 40 \
--flow_shift 12.0 \
--num-frames 48 \
--guidance-scale 4.0 \
--num-inference-steps 40 \
--flow-shift 12.0 \
--fps 16 \
--output i2v_output.mp4
```
@@ -54,17 +54,17 @@ Key arguments:
- `--image`: Path to input image (required).
- `--prompt`: Text description of desired motion/animation.
- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
- `--num_frames`: Number of frames (default 81).
- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high-noise stages for MoE).
- `--negative_prompt`: Optional list of artifacts to suppress.
- `--boundary_ratio`: Boundary split ratio for two-stage MoE models.
- `--flow_shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
- `--num_inference_steps`: Number of denoising steps (default 50).
- `--num-frames`: Number of frames (default 81).
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
- `--negative-prompt`: Optional list of artifacts to suppress.
- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
- `--num-inference-steps`: Number of denoising steps (default 50).
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--output`: Path to save the generated video.
- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
examples/offline_inference/image_to_video/image_to_video.py
+11 -11
View File
@@ -42,31 +42,31 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument("--image", required=True, help="Path to input image.")
parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
parser.add_argument("--negative_prompt", default="", help="Negative prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance_scale", type=float, default=5.0, help="CFG scale.")
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
parser.add_argument(
"--guidance_scale_high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)."
"--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)."
)
parser.add_argument(
"--height", type=int, default=None, help="Video height (auto-calculated from image if not set)."
)
parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).")
parser.add_argument("--num_frames", type=int, default=81, help="Number of frames.")
parser.add_argument("--num_inference_steps", type=int, default=50, help="Sampling steps.")
parser.add_argument("--boundary_ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
parser.add_argument(
"--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
)
parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
parser.add_argument(
"--vae_use_slicing",
"--vae-use-slicing",
action="store_true",
help="Enable VAE slicing for memory optimization.",
)
parser.add_argument(
"--vae_use_tiling",
"--vae-use-tiling",
action="store_true",
help="Enable VAE tiling for memory optimization.",
)
@@ -87,14 +87,14 @@ def parse_args() -> argparse.Namespace:
help="Number of ready layers (blocks) to keep on GPU during generation.",
)
parser.add_argument(
"--cfg_parallel_size",
"--cfg-parallel-size",
type=int,
default=1,
choices=[1, 2],
help="Number of GPUs used for classifier free guidance parallel size.",
)
parser.add_argument(
"--enforce_eager",
"--enforce-eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
examples/offline_inference/lora_inference/README.md
+4 -4
View File
@@ -23,7 +23,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--lora-path /path/to/lora/ \
--lora-scale 1.0 \
--num_inference_steps 50 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_preloaded.png
@@ -40,7 +40,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--lora-request-path /path/to/lora/ \
--lora-scale 1.0 \
--num_inference_steps 50 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_per_request.png
@@ -53,7 +53,7 @@ If no LoRA request is provided, we will use the base model without any LoRA adap
```bash
python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--num_inference_steps 50 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_no_lora.png
@@ -74,7 +74,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
- `--seed`: Random seed for reproducibility (default: 42)
- `--height`: Image height in pixels (default: 1024)
- `--width`: Image width in pixels (default: 1024)
- `--num_inference_steps`: Number of denoising steps (default: 50)
- `--num-inference-steps`: Number of denoising steps (default: 50)
- `--output`: Output file path (default: `lora_output.png`)
## How LoRA Works
examples/offline_inference/lora_inference/lora_inference.py
+1 -1
View File
@@ -18,7 +18,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
parser.add_argument(
"--num_inference_steps",
"--num-inference-steps",
type=int,
default=50,
help="Number of denoising steps for the diffusion sampler.",
examples/offline_inference/text_to_audio/README.md
+8 -8
View File
@@ -18,20 +18,20 @@ If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you
python text_to_audio.py \
--model stabilityai/stable-audio-open-1.0 \
--prompt "The sound of a hammer hitting a wooden surface" \
--negative_prompt "Low quality" \
--negative-prompt "Low quality" \
--seed 42 \
--guidance_scale 7.0 \
--audio_length 10.0 \
--num_inference_steps 100 \
--guidance-scale 7.0 \
--audio-length 10.0 \
--num-inference-steps 100 \
--output stable_audio_output.wav
```
Key arguments:
- `--prompt`: text description (string).
- `--negative_prompt`: negative prompt for classifier-free guidance.
- `--negative-prompt`: negative prompt for classifier-free guidance.
- `--seed`: integer seed for deterministic generation.
- `--guidance_scale`: classifier-free guidance scale.
- `--audio_length`: audio duration in seconds.
- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
- `--guidance-scale`: classifier-free guidance scale.
- `--audio-length`: audio duration in seconds.
- `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
- `--output`: path to save the generated WAV file.
examples/offline_inference/text_to_audio/text_to_audio.py
+9 -9
View File
@@ -9,8 +9,8 @@ the Stable Audio Open model with vLLM-Omni.
Usage:
python text_to_audio.py --prompt "The sound of a dog barking"
python text_to_audio.py --prompt "A piano playing a gentle melody" --audio_length 10.0
python text_to_audio.py --prompt "Thunder and rain sounds" --negative_prompt "Low quality"
python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0
python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality"
"""
import argparse
@@ -38,7 +38,7 @@ def parse_args() -> argparse.Namespace:
help="Text prompt for audio generation.",
)
parser.add_argument(
"--negative_prompt",
"--negative-prompt",
default="Low quality.",
help="Negative prompt for classifier-free guidance.",
)
@@ -49,31 +49,31 @@ def parse_args() -> argparse.Namespace:
help="Random seed for deterministic results.",
)
parser.add_argument(
"--guidance_scale",
"--guidance-scale",
type=float,
default=7.0,
help="Classifier-free guidance scale.",
)
parser.add_argument(
"--audio_start",
"--audio-start",
type=float,
default=0.0,
help="Audio start time in seconds.",
)
parser.add_argument(
"--audio_length",
"--audio-length",
type=float,
default=10.0,
help="Audio length in seconds (max ~47s for stable-audio-open-1.0).",
)
parser.add_argument(
"--num_inference_steps",
"--num-inference-steps",
type=int,
default=100,
help="Number of denoising steps for the diffusion sampler.",
)
parser.add_argument(
"--num_waveforms",
"--num-waveforms",
type=int,
default=1,
help="Number of audio waveforms to generate for the given prompt.",
@@ -85,7 +85,7 @@ def parse_args() -> argparse.Namespace:
help="Path to save the generated audio (WAV format).",
)
parser.add_argument(
"--sample_rate",
"--sample-rate",
type=int,
default=44100,
help="Sample rate for output audio (Stable Audio uses 44100 Hz).",
examples/offline_inference/text_to_image/README.md
+10 -10
View File
@@ -79,9 +79,9 @@ python text_to_image.py \
--model Tongyi-MAI/Z-Image-Turbo \
--prompt "a cup of coffee on the table" \
--seed 42 \
--cfg_scale 4.0 \
--num_images_per_prompt 1 \
--num_inference_steps 50 \
--cfg-scale 4.0 \
--num-images-per-prompt 1 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output outputs/coffee.png
@@ -91,17 +91,17 @@ Key arguments:
- `--prompt`: text description (string).
- `--seed`: integer seed for deterministic sampling.
- `--cfg_scale`: true CFG scale (model-specific guidance strength).
- `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--cfg-scale`: true CFG scale (model-specific guidance strength).
- `--num-images-per-prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
- `--height/--width`: output resolution (defaults 1024x1024).
- `--output`: path to save the generated PNG.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
examples/offline_inference/text_to_image/gradio_demo.py
+1 -1
View File
@@ -40,7 +40,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--default-seed", type=int, default=42, help="Initial seed shown in UI.")
parser.add_argument("--default-cfg-scale", type=float, default=4.0, help="Initial CFG scale shown in UI.")
parser.add_argument(
"--num_inference_steps",
"--num-inference-steps",
type=int,
default=50,
help="Default number of denoising steps shown in the UI.",
examples/offline_inference/text_to_image/text_to_image.py
+49 -14
View File
@@ -5,6 +5,7 @@ import argparse
import os
import time
from pathlib import Path
from typing import Any
import torch
@@ -25,19 +26,19 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
parser.add_argument(
"--negative_prompt",
"--negative-prompt",
default=None,
help="negative prompt for classifier-free conditional guidance.",
)
parser.add_argument("--seed", type=int, default=142, help="Random seed for deterministic results.")
parser.add_argument(
"--cfg_scale",
"--cfg-scale",
type=float,
default=4.0,
help="True classifier-free guidance scale specific to Qwen-Image.",
)
parser.add_argument(
"--guidance_scale",
"--guidance-scale",
type=float,
default=1.0,
help="Classifier-free guidance scale.",
@@ -51,19 +52,19 @@ def parse_args() -> argparse.Namespace:
help="Path to save the generated image (PNG).",
)
parser.add_argument(
"--num_images_per_prompt",
"--num-images-per-prompt",
type=int,
default=1,
help="Number of images to generate for the given prompt.",
)
parser.add_argument(
"--num_inference_steps",
"--num-inference-steps",
type=int,
default=50,
help="Number of denoising steps for the diffusion sampler.",
)
parser.add_argument(
"--cache_backend",
"--cache-backend",
type=str,
default=None,
choices=["cache_dit", "tea_cache"],
@@ -79,26 +80,26 @@ def parse_args() -> argparse.Namespace:
help="Enable cache-dit summary logging after diffusion forward passes.",
)
parser.add_argument(
"--ulysses_degree",
"--ulysses-degree",
type=int,
default=1,
help="Number of GPUs used for ulysses sequence parallelism.",
)
parser.add_argument(
"--ring_degree",
"--ring-degree",
type=int,
default=1,
help="Number of GPUs used for ring sequence parallelism.",
)
parser.add_argument(
"--cfg_parallel_size",
"--cfg-parallel-size",
type=int,
default=1,
choices=[1, 2],
help="Number of GPUs used for classifier free guidance parallel size.",
)
parser.add_argument(
"--enforce_eager",
"--enforce-eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
@@ -119,23 +120,41 @@ def parse_args() -> argparse.Namespace:
help="Number of ready layers (blocks) to keep on GPU during generation.",
)
parser.add_argument(
"--vae_use_slicing",
"--quantization",
type=str,
default=None,
choices=["fp8"],
help="Quantization method for the transformer. "
"Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs). "
"Default: None (no quantization, uses BF16).",
)
parser.add_argument(
"--ignored-layers",
type=str,
default=None,
help="Comma-separated list of layer name patterns to skip quantization. "
"Only used when --quantization is set. "
"Available layers: to_qkv, to_out, add_kv_proj, to_add_out, img_mlp, txt_mlp, proj_out. "
"Example: --ignored-layers 'add_kv_proj,to_add_out'",
)
parser.add_argument(
"--vae-use-slicing",
action="store_true",
help="Enable VAE slicing for memory optimization.",
)
parser.add_argument(
"--vae_use_tiling",
"--vae-use-tiling",
action="store_true",
help="Enable VAE tiling for memory optimization.",
)
parser.add_argument(
"--tensor_parallel_size",
"--tensor-parallel-size",
type=int,
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
)
parser.add_argument(
"--vae_patch_parallel_size",
"--vae-patch-parallel-size",
type=int,
default=1,
help="Number of ranks used for VAE patch/tile parallelism (decode/encode).",
@@ -188,6 +207,18 @@ def main():
# Check if profiling is requested via environment variable
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
# Build quantization kwargs: use quantization_config dict when
# ignored_layers is specified so the list flows through OmniDiffusionConfig
quant_kwargs: dict[str, Any] = {}
ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
if args.quantization and ignored_layers:
quant_kwargs["quantization_config"] = {
"method": args.quantization,
"ignored_layers": ignored_layers,
}
elif args.quantization:
quant_kwargs["quantization"] = args.quantization
omni = Omni(
model=args.model,
enable_layerwise_offload=args.enable_layerwise_offload,
@@ -200,6 +231,7 @@ def main():
parallel_config=parallel_config,
enforce_eager=args.enforce_eager,
enable_cpu_offload=args.enable_cpu_offload,
**quant_kwargs,
)
if profiler_enabled:
@@ -212,6 +244,9 @@ def main():
print(f" Model: {args.model}")
print(f" Inference steps: {args.num_inference_steps}")
print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
print(f" Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
if ignored_layers:
print(f" Ignored layers: {ignored_layers}")
print(
f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
examples/offline_inference/text_to_video/text_to_video.md
+14 -14
View File
@@ -7,14 +7,14 @@ The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text
```bash
python text_to_video.py \
--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
--negative_prompt "<optional quality filter>" \
--negative-prompt "<optional quality filter>" \
--height 480 \
--width 832 \
--num_frames 33 \
--guidance_scale 4.0 \
--guidance_scale_high 3.0 \
--flow_shift 12.0 \
--num_inference_steps 40 \
--num-frames 33 \
--guidance-scale 4.0 \
--guidance-scale-high 3.0 \
--flow-shift 12.0 \
--num-inference-steps 40 \
--fps 16 \
--output t2v_out.mp4
```
@@ -23,15 +23,15 @@ Key arguments:
- `--prompt`: text description (string).
- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
- `--num_frames`: Number of frames (Wan default is 81).
- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
- `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
- `--num-frames`: Number of frames (Wan default is 81).
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
- `--output`: path to save the generated video.
- `--vae_use_slicing`: enable VAE slicing for memory optimization.
- `--vae_use_tiling`: enable VAE tiling for memory optimization.
- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
examples/offline_inference/text_to_video/text_to_video.py
+15 -15
View File
@@ -24,25 +24,25 @@ def parse_args() -> argparse.Namespace:
help="Diffusers Wan2.2 model ID or local path.",
)
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
parser.add_argument("--negative_prompt", default="", help="Negative prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance_scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
parser.add_argument("--guidance_scale_high", type=float, default=None, help="Optional separate CFG for high-noise.")
parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
parser.add_argument("--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise.")
parser.add_argument("--height", type=int, default=720, help="Video height.")
parser.add_argument("--width", type=int, default=1280, help="Video width.")
parser.add_argument("--num_frames", type=int, default=81, help="Number of frames (Wan default is 81).")
parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.")
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
parser.add_argument(
"--boundary_ratio",
"--boundary-ratio",
type=float,
default=0.875,
help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).",
)
parser.add_argument(
"--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
)
parser.add_argument(
"--cache_backend",
"--cache-backend",
type=str,
default=None,
choices=["cache_dit"],
@@ -60,17 +60,17 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
parser.add_argument(
"--vae_use_slicing",
"--vae-use-slicing",
action="store_true",
help="Enable VAE slicing for memory optimization.",
)
parser.add_argument(
"--vae_use_tiling",
"--vae-use-tiling",
action="store_true",
help="Enable VAE tiling for memory optimization.",
)
parser.add_argument(
"--enforce_eager",
"--enforce-eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
@@ -91,26 +91,26 @@ def parse_args() -> argparse.Namespace:
help="Number of ready layers (blocks) to keep on GPU during generation.",
)
parser.add_argument(
"--ulysses_degree",
"--ulysses-degree",
type=int,
default=1,
help="Number of GPUs used for ulysses sequence parallelism.",
)
parser.add_argument(
"--ring_degree",
"--ring-degree",
type=int,
default=1,
help="Number of GPUs used for ring sequence parallelism.",
)
parser.add_argument(
"--cfg_parallel_size",
"--cfg-parallel-size",
type=int,
default=1,
choices=[1, 2],
help="Number of GPUs used for classifier free guidance parallel size.",
)
parser.add_argument(
"--tensor_parallel_size",
"--tensor-parallel-size",
type=int,
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
mkdocs.yml
+1
View File
@@ -94,6 +94,7 @@ plugins:
exclude:
- "re:vllm_omni\\._.*" # Internal modules
- "vllm_omni.diffusion.models.qwen_image" # avoid importing vllm in mkdocs building
- "vllm_omni.diffusion.quantization" # avoid importing vllm in mkdocs building
- "vllm_omni.entrypoints.async_diffusion" # avoid importing vllm in mkdocs building
- "vllm_omni.entrypoints.openai" # avoid importing vllm in mkdocs building
- "vllm_omni.entrypoints.openai.protocol" # avoid importing vllm in mkdocs building
pyproject.toml
+2 -1
View File
@@ -147,7 +147,8 @@ addopts = [
]
markers = [
# ci/cd required
"core_model: Core model tests (run in each PR)",
"core_model: L1&L2 tests (run in each PR)",
"advanced_model: L3&L4 level tests (run in each merge or nightly)",
# function module markers
"diffusion: Diffusion model tests",
"omni: Omni model tests",
tests/conftest.py
+705 -48
View File
@@ -10,11 +10,14 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
if "VLLM_TARGET_DEVICE" not in os.environ:
os.environ["VLLM_TARGET_DEVICE"] = "cpu"
import concurrent.futures
import gc
import socket
import subprocess
import sys
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any
@@ -23,12 +26,22 @@ import psutil
import pytest
import torch
import yaml
from openai import OpenAI
from vllm import TextPrompt
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
from vllm.logger import init_logger
from vllm.utils.network_utils import get_open_port
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
PromptImageInput = list[Any] | Any | None
PromptVideoInput = list[Any] | Any | None
@pytest.fixture(autouse=True)
def default_vllm_config():
@@ -292,6 +305,11 @@ def generate_synthetic_audio(
if max_amp > 0:
audio_data = audio_data / max_amp * 0.8
audio_array = audio_data.copy()
result = {
"np_array": audio_array,
}
# Handle file saving
audio_bytes = None
@@ -318,16 +336,14 @@ def generate_synthetic_audio(
# Return result
base64_audio = base64.b64encode(audio_bytes).decode("utf-8")
result = {
"base64": base64_audio,
}
result["base64"] = base64_audio
if save_to_file and output_path:
result["file_path"] = output_path
return result
def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_file: bool = False) -> str:
def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_file: bool = False) -> dict[str, Any]:
"""Generate synthetic video with bouncing balls and return base64 string."""
import cv2
@@ -397,6 +413,10 @@ def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_f
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
video_frames.append(frame_rgb)
video_array = np.array(video_frames)
result = {
"np_array": video_array,
}
video_bytes = None
saved_file_path = None
@@ -451,16 +471,14 @@ def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_f
base64_video = base64.b64encode(video_bytes).decode("utf-8")
result = {
"base64": base64_video,
}
result["base64"] = base64_video
if save_to_file and saved_file_path:
result["file_path"] = saved_file_path
return result
def generate_synthetic_image(width: int, height: int, save_to_file: bool = False) -> Any:
def generate_synthetic_image(width: int, height: int, save_to_file: bool = False) -> dict[str, Any]:
"""Generate synthetic image with randomly colored squares and return base64 string."""
from PIL import Image, ImageDraw
@@ -488,6 +506,9 @@ def generate_synthetic_image(width: int, height: int, save_to_file: bool = False
# Draw square
draw.rectangle([x, y, x + square_size, y + square_size], fill=color, outline=(0, 0, 0), width=border_width)
image_array = np.array(image)
result = {"np_array": image_array.copy()}
# Handle file saving
image_bytes = None
saved_file_path = None
@@ -521,9 +542,7 @@ def generate_synthetic_image(width: int, height: int, save_to_file: bool = False
base64_image = base64.b64encode(image_bytes).decode("utf-8")
# Return result
result = {
"base64": base64_image,
}
result["base64"] = base64_image
if save_to_file and saved_file_path:
result["file_path"] = saved_file_path
@@ -594,7 +613,6 @@ def convert_audio_to_text(audio_data):
audio_file.write(audio_data)
print(f"audio data is saved: {output_path}")
text = convert_audio_file_to_text(output_path=output_path)
return text
@@ -602,7 +620,7 @@ def convert_audio_to_text(audio_data):
def convert_audio_file_to_text(output_path):
import whisper
model = whisper.load_model("base")
model = whisper.load_model("small")
text = model.transcribe(
output_path,
temperature=0.0,
@@ -641,7 +659,7 @@ def merge_base64_and_convert_to_text(base64_list):
def modify_stage_config(
yaml_path: str,
updates: dict[str, Any],
updates: dict[str, Any] = None,
deletes: dict[str, Any] = None,
) -> str:
"""
@@ -792,7 +810,7 @@ def modify_stage_config(
elif isinstance(current, dict) and last_key in current:
del current[last_key]
else:
raise KeyError(f"Path {path} does not exist")
print(f"Path {path} does not exist")
# Apply deletions first
if deletes:
@@ -829,42 +847,43 @@ def modify_stage_config(
# Delete entire key
del config[key]
# Apply updates
for key, value in updates.items():
if key == "stage_args":
if value and isinstance(value, dict):
stage_args = config.get("stage_args", [])
if not stage_args:
raise ValueError("stage_args does not exist in config")
if updates:
# Apply updates
for key, value in updates.items():
if key == "stage_args":
if value and isinstance(value, dict):
stage_args = config.get("stage_args", [])
if not stage_args:
raise ValueError("stage_args does not exist in config")
for stage_id, stage_updates in value.items():
# Find stage by ID
target_stage = None
for stage in stage_args:
if stage.get("stage_id") == stage_id:
target_stage = stage
break
for stage_id, stage_updates in value.items():
# Find stage by ID
target_stage = None
for stage in stage_args:
if stage.get("stage_id") == stage_id:
target_stage = stage
break
if target_stage is None:
available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
if target_stage is None:
available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
# Apply updates to this stage
for path, val in stage_updates.items():
# Check if this is a simple key (not dot-separated)
# Example: 'engine_input_source' vs 'engine_args.max_model_len'
if "." not in path:
# Direct key assignment (e.g., updating a list value)
target_stage[path] = val
else:
# Dot-separated path (e.g., nested dict access)
apply_update(target_stage, path, val)
elif "." in key:
# Apply using dot-separated path
apply_update(config, key, value)
else:
# Direct top-level key
config[key] = value
# Apply updates to this stage
for path, val in stage_updates.items():
# Check if this is a simple key (not dot-separated)
# Example: 'engine_input_source' vs 'engine_args.max_model_len'
if "." not in path:
# Direct key assignment (e.g., updating a list value)
target_stage[path] = val
else:
# Dot-separated path (e.g., nested dict access)
apply_update(target_stage, path, val)
elif "." in key:
# Apply using dot-separated path
apply_update(config, key, value)
else:
# Direct top-level key
config[key] = value
# Save to new file with timestamp
timestamp = int(time.time())
@@ -1012,3 +1031,641 @@ class OmniServer:
_run_pre_test_cleanup(enable_force=True)
_run_post_test_cleanup(enable_force=True)
cleanup_dist_env_and_memory()
def pytest_addoption(parser):
parser.addoption(
"--run-level",
action="store",
default="core_model",
choices=["core_model", "advanced_model"],
help="Test level to run: L2, L3",
)
@pytest.fixture(scope="session")
def run_level(request):
return request.config.getoption("--run-level")
_omni_server_lock = threading.Lock()
@pytest.fixture(scope="module")
def omni_server(request, run_level):
"""Start vLLM-Omni server as a subprocess with actual model weights.
Uses session scope so the server starts only once for the entire test session.
Multi-stage initialization can take 10-20+ minutes.
"""
with _omni_server_lock:
model, stage_config_path = request.param
if run_level == "advanced_model":
stage_config_path = modify_stage_config(
stage_config_path,
deletes={
"stage_args": {
0: ["engine_args.load_format"],
1: ["engine_args.load_format"],
2: ["engine_args.load_format"],
}
},
)
with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
print("OmniServer started successfully")
yield server
print("OmniServer stopping...")
print("OmniServer stopped")
@dataclass
class OmniResponse:
text_content: str | None = None
audio_data: list[str] | None = None
audio_content: str | None = None
similarity: float | None = None
e2e_latency: float | None = None
success: bool = False
error_message: str | None = None
def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], run_level):
"""
Validate response results.
Args:
response: OmniResponse object
Raises:
AssertionError: When the response does not meet validation criteria
"""
assert response.success, "The request failed."
e2e_latency = response.e2e_latency
if e2e_latency is not None:
print(f"the avg e2e is: {e2e_latency}")
modalities = request_config.get("modalities", ["text", "audio"])
if "audio" in modalities:
assert response.audio_content is not None, "No audio output is generated"
print(f"audio content is: {response.audio_content}")
if "text" in modalities:
assert response.text_content is not None, "No text output is generated"
print(f"text content is: {response.text_content}")
if run_level == "advanced_model":
# Verify image description
word_types = ["text", "image", "audio", "video"]
keywords_dict = request_config.get("key_words", {})
for word_type in word_types:
keywords = keywords_dict.get(word_type)
if "text" in modalities:
if keywords:
assert any(keyword in response.text_content.lower() for keyword in keywords), (
"The output does not contain any of the keywords."
)
else:
if keywords:
assert any(keyword in response.audio_content.lower() for keyword in keywords), (
"The output does not contain any of the keywords."
)
# Verify similarity
if "text" in modalities and "audio" in modalities:
assert response.similarity > 0.9, "The audio content is not same as the text"
print(f"similarity is: {response.similarity}")
class OpenAIClientHandler:
"""
OpenAI client handler class, encapsulating both streaming and non-streaming response processing logic.
This class integrates OpenAI API request sending, response handling, and validation functionality,
supporting both single request and concurrent request modes.
"""
def __init__(
self, host: str = "127.0.0.1", port: int = get_open_port(), api_key: str = "EMPTY", run_level: str = None
):
"""
Initialize the OpenAI client.
Args:
host: vLLM-Omni server host address
port: vLLM-Omni server port
api_key: API key (defaults to "EMPTY")
"""
self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key)
self.run_level = run_level
def _process_stream_response(self, chat_completion) -> OmniResponse:
"""
Process streaming responses.
Args:
chat_completion: OpenAI streaming response object
request_config: Request configuration dictionary
Returns:
OmniResponse: Processed response object
"""
result = OmniResponse()
start_time = time.perf_counter()
try:
text_content = ""
audio_data = []
for chunk in chat_completion:
for choice in chunk.choices:
# Get content data
if hasattr(choice, "delta"):
content = getattr(choice.delta, "content", None)
else:
content = None
# Get modality type
modality = getattr(chunk, "modality", None)
# Process content based on modality type
if modality == "audio" and content:
audio_data.append(content)
elif modality == "text" and content:
text_content += content if content else ""
# Calculate end-to-end latency
result.e2e_latency = time.perf_counter() - start_time
# Process audio and text content
audio_content = None
similarity = None
if audio_data or text_content:
if audio_data:
audio_content = merge_base64_and_convert_to_text(audio_data)
if audio_content and text_content:
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
# Populate result object
result.text_content = text_content
result.audio_data = audio_data
result.audio_content = audio_content
result.similarity = similarity
result.success = True
except Exception as e:
result.error_message = f"Stream processing error: {str(e)}"
print(f"Error: {result.error_message}")
return result
def _process_non_stream_response(self, chat_completion) -> OmniResponse:
"""
Process non-streaming responses.
Args:
chat_completion: OpenAI non-streaming response object
request_config: Request configuration dictionary
Returns:
OmniResponse: Processed response object
"""
result = OmniResponse()
start_time = time.perf_counter()
try:
audio_data = None
text_content = None
# Iterate through all choices
for choice in chat_completion.choices:
# Process audio data
if hasattr(choice.message, "audio") and choice.message.audio is not None:
audio_message = choice.message
audio_data = audio_message.audio.data
# Process text content
if hasattr(choice.message, "content") and choice.message.content is not None:
text_content = choice.message.content
# Calculate end-to-end latency
result.e2e_latency = time.perf_counter() - start_time
# Process audio and text content
audio_content = None
similarity = None
if audio_data or text_content:
if audio_data:
audio_content = convert_audio_to_text(audio_data)
if audio_content and text_content:
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
# Populate result object
result.text_content = text_content
result.audio_content = audio_content
result.similarity = similarity
result.success = True
except Exception as e:
result.error_message = f"Non-stream processing error: {str(e)}"
print(f"Error: {result.error_message}")
return result
def send_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
"""
Send OpenAI requests.
Args:
request_config: Request configuration dictionary containing parameters like model, messages, stream
request_num: Number of requests, defaults to 1 (single request)
Returns:
List[OmniResponse]: List of response objects
"""
responses = []
stream = request_config.get("stream", False)
if request_num == 1:
# Send single request
chat_completion = self.client.chat.completions.create(
model=request_config.get("model"), messages=request_config.get("messages"), stream=stream
)
if stream:
response = self._process_stream_response(chat_completion)
else:
response = self._process_non_stream_response(chat_completion)
assert_omni_response(response, request_config, run_level=self.run_level)
responses.append(response)
else:
# Send concurrent requests
with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
futures = []
# Submit all request tasks
for _ in range(request_num):
future = executor.submit(
self.client.chat.completions.create,
model=request_config.get("model"),
messages=request_config.get("messages"),
stream=stream,
)
futures.append(future)
# Process completed tasks
for future in concurrent.futures.as_completed(futures):
chat_completion = future.result()
if stream:
response = self._process_stream_response(chat_completion)
else:
response = self._process_non_stream_response(chat_completion)
assert_omni_response(response, request_config, run_level=self.run_level)
responses.append(response)
return responses
@pytest.fixture
def openai_client(omni_server, run_level):
"""Create OpenAIClientHandler fixture"""
return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level)
class OmniRunner:
"""
Test runner for Omni models.
"""
def __init__(
self,
model_name: str,
seed: int = 42,
stage_init_timeout: int = 300,
batch_timeout: int = 10,
init_timeout: int = 300,
shm_threshold_bytes: int = 65536,
log_stats: bool = False,
stage_configs_path: str | None = None,
**kwargs,
) -> None:
"""
Initialize an OmniRunner for testing.
Args:
model_name: The model name or path
seed: Random seed for reproducibility
stage_init_timeout: Timeout for initializing a single stage in seconds
batch_timeout: Timeout for batching in seconds
init_timeout: Timeout for initializing stages in seconds
shm_threshold_bytes: Threshold for using shared memory
log_stats: Enable detailed statistics logging
stage_configs_path: Optional path to YAML stage config file
**kwargs: Additional arguments passed to Omni
"""
cleanup_dist_env_and_memory()
_run_pre_test_cleanup(enable_force=True)
_run_post_test_cleanup(enable_force=True)
self.model_name = model_name
self.seed = seed
self.omni = Omni(
model=model_name,
log_stats=log_stats,
stage_init_timeout=stage_init_timeout,
batch_timeout=batch_timeout,
init_timeout=init_timeout,
shm_threshold_bytes=shm_threshold_bytes,
stage_configs_path=stage_configs_path,
**kwargs,
)
def get_default_sampling_params_list(self) -> list[OmniSamplingParams]:
"""
Get a list of default sampling parameters for all stages.
Returns:
List of SamplingParams with default decoding for each stage
"""
return [st.default_sampling_params for st in self.omni.stage_list]
def get_omni_inputs(
self,
prompts: list[str] | str,
system_prompt: str | None = None,
audios: PromptAudioInput = None,
images: PromptImageInput = None,
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
modalities: list[str] | None = None,
) -> list[TextPrompt]:
"""
Construct Omni input format from prompts and multimodal data.
Args:
prompts: Text prompt(s) - either a single string or list of strings
system_prompt: Optional system prompt (defaults to Qwen system prompt)
audios: Audio input(s) - tuple of (audio_array, sample_rate) or list of tuples
images: Image input(s) - PIL Image or list of PIL Images
videos: Video input(s) - numpy array or list of numpy arrays
mm_processor_kwargs: Optional processor kwargs (e.g., use_audio_in_video)
Returns:
List of prompt dictionaries suitable for Omni.generate()
"""
if system_prompt is None:
system_prompt = (
"You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
"Group, capable of perceiving auditory and visual inputs, as well as "
"generating text and speech."
)
video_padding_token = "<|VIDEO|>"
image_padding_token = "<|IMAGE|>"
audio_padding_token = "<|AUDIO|>"
if "Qwen3-Omni-30B-A3B-Instruct" in self.model_name:
video_padding_token = "<|video_pad|>"
image_padding_token = "<|image_pad|>"
audio_padding_token = "<|audio_pad|>"
if isinstance(prompts, str):
prompts = [prompts]
def _normalize_mm_input(mm_input, num_prompts):
if mm_input is None:
return [None] * num_prompts
if isinstance(mm_input, list):
if len(mm_input) != num_prompts:
raise ValueError(
f"Multimodal input list length ({len(mm_input)}) must match prompts length ({num_prompts})"
)
return mm_input
return [mm_input] * num_prompts
num_prompts = len(prompts)
audios_list = _normalize_mm_input(audios, num_prompts)
images_list = _normalize_mm_input(images, num_prompts)
videos_list = _normalize_mm_input(videos, num_prompts)
omni_inputs = []
for i, prompt_text in enumerate(prompts):
user_content = ""
multi_modal_data = {}
audio = audios_list[i]
if audio is not None:
if isinstance(audio, list):
for _ in audio:
user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
multi_modal_data["audio"] = audio
else:
user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
multi_modal_data["audio"] = audio
image = images_list[i]
if image is not None:
if isinstance(image, list):
for _ in image:
user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
multi_modal_data["image"] = image
else:
user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
multi_modal_data["image"] = image
video = videos_list[i]
if video is not None:
if isinstance(video, list):
for _ in video:
user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
multi_modal_data["video"] = video
else:
user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
multi_modal_data["video"] = video
user_content += prompt_text
full_prompt = (
f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
f"<|im_start|>user\n{user_content}<|im_end|>\n"
f"<|im_start|>assistant\n"
)
input_dict: TextPrompt = {"prompt": full_prompt}
if multi_modal_data:
input_dict["multi_modal_data"] = multi_modal_data
if modalities:
input_dict["modalities"] = modalities
if mm_processor_kwargs:
input_dict["mm_processor_kwargs"] = mm_processor_kwargs
omni_inputs.append(input_dict)
return omni_inputs
def generate(
self,
prompts: list[TextPrompt],
sampling_params_list: list[OmniSamplingParams] | None = None,
) -> list[OmniRequestOutput]:
"""
Generate outputs for the given prompts.
Args:
prompts: List of prompt dictionaries with 'prompt' and optionally
'multi_modal_data' keys
sampling_params_list: List of sampling parameters for each stage.
If None, uses default parameters.
Returns:
List of OmniRequestOutput objects from stages with final_output=True
"""
if sampling_params_list is None:
sampling_params_list = self.get_default_sampling_params_list()
return self.omni.generate(prompts, sampling_params_list)
def generate_multimodal(
self,
prompts: list[str] | str,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
audios: PromptAudioInput = None,
images: PromptImageInput = None,
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
modalities: list[str] | None = None,
) -> list[OmniRequestOutput]:
"""
Convenience method to generate with multimodal inputs.
Args:
prompts: Text prompt(s)
sampling_params_list: List of sampling parameters for each stage
system_prompt: Optional system prompt
audios: Audio input(s)
images: Image input(s)
videos: Video input(s)
mm_processor_kwargs: Optional processor kwargs
Returns:
List of OmniRequestOutput objects from stages with final_output=True
"""
omni_inputs = self.get_omni_inputs(
prompts=prompts,
system_prompt=system_prompt,
audios=audios,
images=images,
videos=videos,
mm_processor_kwargs=mm_processor_kwargs,
modalities=modalities,
)
return self.generate(omni_inputs, sampling_params_list)
def _cleanup_process(self):
try:
keywords = ["enginecore"]
for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
try:
cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
name = proc.name().lower()
is_process = any(keyword in cmdline for keyword in keywords) or any(
keyword in name for keyword in keywords
)
if is_process:
print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
try:
proc.terminate()
time.sleep(2)
except Exception:
proc.kill()
except (psutil.NoSuchProcess, psutil.AccessDenied):
pass
except Exception as e:
print(f"Error in psutil vllm cleanup: {e}")
def __enter__(self):
"""Context manager entry."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit - cleanup resources."""
if hasattr(self.omni, "close"):
self.omni.close()
self._cleanup_process()
_run_pre_test_cleanup(enable_force=True)
_run_post_test_cleanup(enable_force=True)
cleanup_dist_env_and_memory()
@pytest.fixture(scope="module")
def omni_runner(request):
with _omni_server_lock:
model, stage_config_path = request.param
with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner:
print("OmniRunner started successfully")
yield runner
print("OmniRunner stopping...")
print("OmniRunner stopped")
class OmniRunnerHandler:
def __init__(self, omni_runner):
self.runner = omni_runner
def _process_output(self, outputs: list[Any]) -> OmniResponse:
result = OmniResponse()
try:
text_content = None
audio_content = None
for stage_output in outputs:
if getattr(stage_output, "final_output_type", None) == "text":
text_content = stage_output.request_output[0].outputs[0].text
if getattr(stage_output, "final_output_type", None) == "audio":
audio_content = stage_output.request_output[0].outputs[0].multimodal_output["audio"]
result.audio_content = audio_content
result.text_content = text_content
result.success = True
except Exception as e:
result.error_message = f"Output processing error: {str(e)}"
result.success = False
print(f"Error: {result.error_message}")
return result
def send_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse:
if request_config is None:
request_config = {}
prompts = request_config.get("prompts")
videos = request_config.get("videos")
images = request_config.get("images")
audios = request_config.get("audios")
modalities = request_config.get("modalities", ["text", "audio"])
outputs = self.runner.generate_multimodal(
prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities
)
response = self._process_output(outputs)
assert_omni_response(response, request_config, run_level="L2")
return response
@pytest.fixture
def omni_runner_handler(omni_runner):
return OmniRunnerHandler(omni_runner)
tests/diffusion/quantization/__init__.py
+2
View File
@@ -0,0 +1,2 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/diffusion/quantization/test_fp8_config.py
+137
View File
@@ -0,0 +1,137 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for FP8 quantization config."""
import pytest
def test_fp8_config_creation():
"""Test that FP8 config can be created."""
from vllm_omni.diffusion.quantization import get_diffusion_quant_config
config = get_diffusion_quant_config("fp8")
assert config is not None
assert config.get_name() == "fp8"
def test_vllm_config_extraction():
"""Test that vLLM config can be extracted from diffusion config."""
from vllm_omni.diffusion.quantization import (
get_diffusion_quant_config,
get_vllm_quant_config_for_layers,
)
diff_config = get_diffusion_quant_config("fp8")
vllm_config = get_vllm_quant_config_for_layers(diff_config)
assert vllm_config is not None
assert vllm_config.activation_scheme == "dynamic"
def test_none_quantization():
"""Test that None quantization returns None config."""
from vllm_omni.diffusion.quantization import (
get_diffusion_quant_config,
get_vllm_quant_config_for_layers,
)
config = get_diffusion_quant_config(None)
assert config is None
vllm_config = get_vllm_quant_config_for_layers(config)
assert vllm_config is None
def test_invalid_quantization():
"""Test that invalid quantization method raises error."""
from vllm_omni.diffusion.quantization import get_diffusion_quant_config
with pytest.raises(ValueError, match="Unknown quantization method"):
get_diffusion_quant_config("invalid_method")
def test_fp8_config_with_custom_params():
"""Test FP8 config with custom parameters."""
from vllm_omni.diffusion.quantization import get_diffusion_quant_config
config = get_diffusion_quant_config(
"fp8",
activation_scheme="static",
ignored_layers=["proj_out"],
)
assert config is not None
assert config.activation_scheme == "static"
assert "proj_out" in config.ignored_layers
def test_supported_methods():
"""Test that supported methods list is correct."""
from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
assert "fp8" in SUPPORTED_QUANTIZATION_METHODS
def test_quantization_integration():
"""Test end-to-end quantization flow through OmniDiffusionConfig."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
# Test with quantization string only
config = OmniDiffusionConfig(model="test", quantization="fp8")
assert config.quantization_config is not None
assert config.quantization_config.get_name() == "fp8"
# Test with quantization_config dict
config2 = OmniDiffusionConfig(
model="test",
quantization_config={"method": "fp8", "activation_scheme": "static"},
)
assert config2.quantization_config is not None
assert config2.quantization_config.get_name() == "fp8"
assert config2.quantization_config.activation_scheme == "static"
# Test that vLLM config can be extracted
vllm_config = config.quantization_config.get_vllm_quant_config()
assert vllm_config is not None
def test_quantization_dict_not_mutated():
"""Test that passing a dict to quantization_config doesn't mutate it."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
original_dict = {"method": "fp8", "activation_scheme": "static"}
dict_copy = original_dict.copy()
OmniDiffusionConfig(model="test", quantization_config=original_dict)
# Original dict should be unchanged
assert original_dict == dict_copy
def test_quantization_conflicting_methods_warning(caplog):
"""Test warning when quantization and quantization_config['method'] conflict."""
import logging
from vllm_omni.diffusion.data import OmniDiffusionConfig
with caplog.at_level(logging.WARNING):
config = OmniDiffusionConfig(
model="test",
quantization="fp8", # This should be overridden
quantization_config={"method": "fp8", "activation_scheme": "static"},
)
# No warning when methods match
assert config.quantization_config is not None
def test_fp8_delegates_to_vllm_config():
"""Test that DiffusionFp8Config delegates to vLLM's Fp8Config."""
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
from vllm_omni.diffusion.quantization import DiffusionFp8Config
# Test that quant_config_cls is set correctly
assert DiffusionFp8Config.quant_config_cls is Fp8Config
# Test that get_name() delegates to vLLM
assert DiffusionFp8Config.get_name() == Fp8Config.get_name()
# Test that get_min_capability() delegates to vLLM
assert DiffusionFp8Config.get_min_capability() == Fp8Config.get_min_capability()
tests/e2e/offline_inference/conftest.py
-353
View File
@@ -1,353 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Pytest configuration and fixtures for vllm-omni tests.
"""
from typing import Any
import pytest
from vllm import TextPrompt
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
from tests.conftest import _run_post_test_cleanup, _run_pre_test_cleanup
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.outputs import OmniRequestOutput
PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
PromptImageInput = list[Any] | Any | None
PromptVideoInput = list[Any] | Any | None
class OmniRunner:
"""
Test runner for Omni models.
"""
def __init__(
self,
model_name: str,
seed: int = 42,
stage_init_timeout: int = 300,
batch_timeout: int = 10,
init_timeout: int = 300,
shm_threshold_bytes: int = 65536,
log_stats: bool = False,
stage_configs_path: str | None = None,
**kwargs,
) -> None:
"""
Initialize an OmniRunner for testing.
Args:
model_name: The model name or path
seed: Random seed for reproducibility
stage_init_timeout: Timeout for initializing a single stage in seconds
batch_timeout: Timeout for batching in seconds
init_timeout: Timeout for initializing stages in seconds
shm_threshold_bytes: Threshold for using shared memory
log_stats: Enable detailed statistics logging
stage_configs_path: Optional path to YAML stage config file
**kwargs: Additional arguments passed to Omni
"""
cleanup_dist_env_and_memory()
_run_pre_test_cleanup(enable_force=True)
_run_post_test_cleanup(enable_force=True)
self.model_name = model_name
self.seed = seed
self.omni = Omni(
model=model_name,
log_stats=log_stats,
stage_init_timeout=stage_init_timeout,
batch_timeout=batch_timeout,
init_timeout=init_timeout,
shm_threshold_bytes=shm_threshold_bytes,
stage_configs_path=stage_configs_path,
**kwargs,
)
def get_default_sampling_params_list(self) -> list[OmniSamplingParams]:
"""
Get a list of default sampling parameters for all stages.
Returns:
List of SamplingParams with default decoding for each stage
"""
return [st.default_sampling_params for st in self.omni.stage_list]
def get_omni_inputs(
self,
prompts: list[str] | str,
system_prompt: str | None = None,
audios: PromptAudioInput = None,
images: PromptImageInput = None,
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
modalities: list[str] | None = None,
) -> list[TextPrompt]:
"""
Construct Omni input format from prompts and multimodal data.
Args:
prompts: Text prompt(s) - either a single string or list of strings
system_prompt: Optional system prompt (defaults to Qwen system prompt)
audios: Audio input(s) - tuple of (audio_array, sample_rate) or list of tuples
images: Image input(s) - PIL Image or list of PIL Images
videos: Video input(s) - numpy array or list of numpy arrays
mm_processor_kwargs: Optional processor kwargs (e.g., use_audio_in_video)
Returns:
List of prompt dictionaries suitable for Omni.generate()
"""
if system_prompt is None:
system_prompt = (
"You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
"Group, capable of perceiving auditory and visual inputs, as well as "
"generating text and speech."
)
video_padding_token = "<|VIDEO|>"
image_padding_token = "<|IMAGE|>"
audio_padding_token = "<|AUDIO|>"
if self.model_name == "Qwen/Qwen3-Omni-30B-A3B-Instruct":
video_padding_token = "<|video_pad|>"
image_padding_token = "<|image_pad|>"
audio_padding_token = "<|audio_pad|>"
if isinstance(prompts, str):
prompts = [prompts]
def _normalize_mm_input(mm_input, num_prompts):
if mm_input is None:
return [None] * num_prompts
if isinstance(mm_input, list):
if len(mm_input) != num_prompts:
raise ValueError(
f"Multimodal input list length ({len(mm_input)}) must match prompts length ({num_prompts})"
)
return mm_input
return [mm_input] * num_prompts
num_prompts = len(prompts)
audios_list = _normalize_mm_input(audios, num_prompts)
images_list = _normalize_mm_input(images, num_prompts)
videos_list = _normalize_mm_input(videos, num_prompts)
omni_inputs = []
for i, prompt_text in enumerate(prompts):
user_content = ""
multi_modal_data = {}
audio = audios_list[i]
if audio is not None:
if isinstance(audio, list):
for _ in audio:
user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
multi_modal_data["audio"] = audio
else:
user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
multi_modal_data["audio"] = audio
image = images_list[i]
if image is not None:
if isinstance(image, list):
for _ in image:
user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
multi_modal_data["image"] = image
else:
user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
multi_modal_data["image"] = image
video = videos_list[i]
if video is not None:
if isinstance(video, list):
for _ in video:
user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
multi_modal_data["video"] = video
else:
user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
multi_modal_data["video"] = video
user_content += prompt_text
full_prompt = (
f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
f"<|im_start|>user\n{user_content}<|im_end|>\n"
f"<|im_start|>assistant\n"
)
input_dict: TextPrompt = {"prompt": full_prompt}
if multi_modal_data:
input_dict["multi_modal_data"] = multi_modal_data
if modalities:
input_dict["modalities"] = modalities
if mm_processor_kwargs:
input_dict["mm_processor_kwargs"] = mm_processor_kwargs
omni_inputs.append(input_dict)
return omni_inputs
def generate(
self,
prompts: list[TextPrompt],
sampling_params_list: list[OmniSamplingParams] | None = None,
) -> list[OmniRequestOutput]:
"""
Generate outputs for the given prompts.
Args:
prompts: List of prompt dictionaries with 'prompt' and optionally
'multi_modal_data' keys
sampling_params_list: List of sampling parameters for each stage.
If None, uses default parameters.
Returns:
List of OmniRequestOutput objects from stages with final_output=True
"""
if sampling_params_list is None:
sampling_params_list = self.get_default_sampling_params_list()
return self.omni.generate(prompts, sampling_params_list)
def generate_multimodal(
self,
prompts: list[str] | str,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
audios: PromptAudioInput = None,
images: PromptImageInput = None,
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
modalities: list[str] | None = None,
) -> list[OmniRequestOutput]:
"""
Convenience method to generate with multimodal inputs.
Args:
prompts: Text prompt(s)
sampling_params_list: List of sampling parameters for each stage
system_prompt: Optional system prompt
audios: Audio input(s)
images: Image input(s)
videos: Video input(s)
mm_processor_kwargs: Optional processor kwargs
Returns:
List of OmniRequestOutput objects from stages with final_output=True
"""
omni_inputs = self.get_omni_inputs(
prompts=prompts,
system_prompt=system_prompt,
audios=audios,
images=images,
videos=videos,
mm_processor_kwargs=mm_processor_kwargs,
modalities=modalities,
)
return self.generate(omni_inputs, sampling_params_list)
def generate_audio(
self,
prompts: list[str] | str,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
audios: PromptAudioInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
) -> list[OmniRequestOutput]:
"""
Convenience method to generate with multimodal inputs.
Args:
prompts: Text prompt(s)
sampling_params_list: List of sampling parameters for each stage
system_prompt: Optional system prompt
audios: Audio input(s)
mm_processor_kwargs: Optional processor kwargs
Returns:
List of OmniRequestOutput objects from stages with final_output=True
"""
omni_inputs = self.get_omni_inputs(
prompts=prompts,
system_prompt=system_prompt,
audios=audios,
mm_processor_kwargs=mm_processor_kwargs,
)
return self.generate(omni_inputs, sampling_params_list)
def generate_video(
self,
prompts: list[str] | str,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
videos: PromptVideoInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
) -> list[OmniRequestOutput]:
"""
Convenience method to generate with multimodal inputs.
Args:
prompts: Text prompt(s)
sampling_params_list: List of sampling parameters for each stage
system_prompt: Optional system prompt
videos: Video input(s)
mm_processor_kwargs: Optional processor kwargs
Returns:
List of OmniRequestOutput objects from stages with final_output=True
"""
omni_inputs = self.get_omni_inputs(
prompts=prompts,
system_prompt=system_prompt,
videos=videos,
mm_processor_kwargs=mm_processor_kwargs,
)
return self.generate(omni_inputs, sampling_params_list)
def generate_image(
self,
prompts: list[str] | str,
sampling_params_list: list[OmniSamplingParams] | None = None,
system_prompt: str | None = None,
images: PromptImageInput = None,
mm_processor_kwargs: dict[str, Any] | None = None,
) -> list[OmniRequestOutput]:
"""
Convenience method to generate with multimodal inputs.
Args:
prompts: Text prompt(s)
sampling_params_list: List of sampling parameters for each stage
system_prompt: Optional system prompt
images: Image input(s)
mm_processor_kwargs: Optional processor kwargs
Returns:
List of OmniRequestOutput objects from stages with final_output=True
"""
omni_inputs = self.get_omni_inputs(
prompts=prompts,
system_prompt=system_prompt,
images=images,
mm_processor_kwargs=mm_processor_kwargs,
)
return self.generate(omni_inputs, sampling_params_list)
def __enter__(self):
"""Context manager entry."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit - cleanup resources."""
self.close()
del self.omni
cleanup_dist_env_and_memory()
_run_post_test_cleanup(enable_force=True)
def close(self):
"""Close and cleanup the Omni instance."""
if hasattr(self.omni, "close"):
self.omni.close()
@pytest.fixture(scope="session")
def omni_runner():
return OmniRunner
tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml
-106
View File
@@ -1,106 +0,0 @@
# stage config for running qwen2.5-omni with architecture of OmniLLM.
# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
# This config is optimized for CI e2e tests.
stage_args:
- stage_id: 0
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
max_num_seqs: 1
gpu_memory_utilization: 0.8
skip_mm_profiling: true
enforce_eager: true # Now we only support eager mode
trust_remote_code: true
engine_output_type: latent
enable_prefix_caching: false
is_comprehension: true
final_output: true
final_output_type: text
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 128
seed: 42
detokenize: True
repetition_penalty: 1.1
- stage_id: 1
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
max_num_seqs: 1
gpu_memory_utilization: 0.8
skip_mm_profiling: true
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: latent
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
default_sampling_params:
temperature: 0.9
top_p: 0.8
top_k: 40
max_tokens: 128
seed: 42
detokenize: True
repetition_penalty: 1.05
stop_token_ids: [8294]
- stage_id: 2
runtime:
process: true
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio
max_num_batched_tokens: 4069
engine_input_source: [1]
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 128
seed: 42
detokenize: True
repetition_penalty: 1.1
# Top-level runtime config (concise): default windows and stage edges
runtime:
enabled: true
defaults:
window_size: -1 # Simplified: trigger downstream only after full upstream completion
max_inflight: 1 # Simplified: process serially within each stage
edges:
- from: 0 # thinker → talker: trigger only after receiving full input (-1)
to: 1
window_size: -1
- from: 1 # talker → code2wav: trigger only after receiving full input (-1)
to: 2
window_size: -1
tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml
-99
View File
@@ -1,99 +0,0 @@
# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
# Stage 0: Thinker (multimodal understanding + text generation)
# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
# The following config has been verified on 2x H100-80G GPUs.
stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: false
trust_remote_code: true
engine_output_type: latent # Output hidden states for talker
distributed_executor_backend: "mp"
enable_prefix_caching: false
hf_config_name: thinker_config
tensor_parallel_size: 1
load_format: dummy
final_output: true
final_output_type: text
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 100
seed: 42
detokenize: True
repetition_penalty: 1.05
- stage_id: 1
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
enforce_eager: true
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
# tensor_parallel_size: 2
enable_prefix_caching: false
distributed_executor_backend: "mp"
hf_config_name: talker_config
load_format: dummy
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
# final_output: true
# final_output_type: text
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 100
seed: 42
detokenize: False
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 2
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio # Final output: audio waveform
gpu_memory_utilization: 0.1
distributed_executor_backend: "mp"
max_num_batched_tokens: 1000000
hf_config_name: thinker_config
load_format: dummy
async_scheduling: false
engine_input_source: [1]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 200
seed: 42
detokenize: True
repetition_penalty: 1.1
tests/e2e/offline_inference/test_qwen2_5_omni.py
+59 -109
View File
@@ -1,138 +1,88 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E tests for Qwen2.5-Omni model with mixed modality inputs and audio output.
E2E tests for Qwen2.5-Omni model with mixed modality inputs, audio and text output.
"""
from pathlib import Path
import pytest
from vllm.assets.audio import AudioAsset
from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset
from vllm.envs import VLLM_USE_MODELSCOPE
from vllm.multimodal.image import convert_image_mode
from tests.utils import create_new_process_for_each_test, hardware_test
from tests.conftest import (
generate_synthetic_audio,
generate_synthetic_image,
generate_synthetic_video,
)
from tests.utils import hardware_test
from vllm_omni.platforms import current_omni_platform
from .conftest import OmniRunner
models = ["Qwen/Qwen2.5-Omni-3B"]
models = ["Qwen/Qwen2.5-Omni-7B"]
# CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
if current_omni_platform.is_npu():
stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
elif current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_config = str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
stage_config = str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
else:
stage_config = str(Path(__file__).parent / "stage_configs" / "qwen2_5_omni_ci.yaml")
stage_config = str(Path(__file__).parent.parent / "stage_configs" / "qwen2_5_omni_ci.yaml")
# Create parameter combinations for model and stage config
test_params = [(model, stage_config) for model in models]
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@create_new_process_for_each_test("spawn")
@pytest.mark.parametrize("test_config", test_params)
def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
"""Test processing audio, image, and video together, generating audio output."""
model, stage_config_path = test_config
with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
# Prepare multimodal inputs
question = "What is recited in the audio? What is in this image? Describe the video briefly."
audio = AudioAsset("mary_had_lamb").audio_and_sample_rate
audio = (audio[0][: 16000 * 5], audio[1]) # Trim to first 5 seconds
image = convert_image_mode(ImageAsset("cherry_blossom").pil_image.resize((128, 128)), "RGB")
if not VLLM_USE_MODELSCOPE:
video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
else:
# modelscope can't access raushan-testing-hf/videos-test, skip video input temporarily
video = None
outputs = runner.generate_multimodal(
prompts=question,
audios=audio,
images=image,
videos=video,
)
# Find and verify text output (thinker stage)
text_output = None
output_count = 0
for stage_output in outputs:
if stage_output.final_output_type == "text":
text_output = stage_output
output_count += 1
break
assert output_count > 0
assert text_output is not None
assert len(text_output.request_output) > 0
text_content = text_output.request_output[0].outputs[0].text
assert text_content is not None
assert len(text_content.strip()) > 0
# Find and verify audio output (code2wav stage)
audio_output = None
output_count = 0
for stage_output in outputs:
if stage_output.final_output_type == "audio":
audio_output = stage_output
output_count += 1
break
assert output_count > 0
assert audio_output is not None
assert len(audio_output.request_output) > 0
# Verify audio tensor exists and has content
audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"]
assert audio_tensor is not None
assert audio_tensor.numel() > 0
def get_question(prompt_type="mix"):
prompts = {
"mix": "What is recited in the audio? What is in this image? Describe the video briefly.",
"text_only": "What is the capital of China?",
}
return prompts.get(prompt_type, prompts["mix"])
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@create_new_process_for_each_test("spawn")
@pytest.mark.parametrize("test_config", test_params)
def test_mixed_modalities_to_text_only(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None:
"""Test processing audio, image, and video together, generating audio output."""
model, stage_config_path = test_config
with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
# Prepare multimodal inputs
question = "What is recited in the audio? What is in this image? Describe the video briefly."
audio = AudioAsset("mary_had_lamb").audio_and_sample_rate
audio = (audio[0][: 16000 * 5], audio[1]) # Trim to first 5 seconds
image = convert_image_mode(ImageAsset("cherry_blossom").pil_image.resize((128, 128)), "RGB")
video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
modalities = ["text"]
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
def test_mix_to_audio(omni_runner, omni_runner_handler) -> None:
"""
Test multi-modal input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
Input Modal: text + audio + video + image
Output Modal: audio
Input Setting: stream=False
Datasets: single request
"""
video = generate_synthetic_video(16, 16, 30)["np_array"]
image = generate_synthetic_image(16, 16)["np_array"]
audio = generate_synthetic_audio(1, 1, 16000)["np_array"]
if len(audio.shape) == 2:
audio = audio.squeeze()
outputs = runner.generate_multimodal(
prompts=question,
audios=audio,
images=image,
videos=video,
modalities=modalities,
)
request_config = {
"prompts": get_question(),
"videos": video,
"images": image,
"audios": (audio, 16000),
"modalities": ["audio"],
}
# Find and verify text output (thinker stage)
text_output = None
output_count = 0
for stage_output in outputs:
assert stage_output.final_output_type != "audio"
if stage_output.final_output_type == "text":
text_output = stage_output
output_count += 1
break
assert output_count > 0
# Test single completion
omni_runner_handler.send_request(request_config)
assert text_output is not None
assert len(text_output.request_output) > 0
text_content = text_output.request_output[0].outputs[0].text
assert text_content is not None
assert len(text_content.strip()) > 0
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
def test_text_to_text(omni_runner, omni_runner_handler) -> None:
"""
Test text input processing and text output generation via OpenAI API.
Deploy Setting: default yaml
Input Modal: text
Output Modal: text
Input Setting: stream=False
Datasets: single request
"""
request_config = {"prompts": get_question("text_only"), "modalities": ["text"]}
# Test single completion
omni_runner_handler.send_request(request_config)
tests/e2e/offline_inference/test_qwen3_omni.py
+18 -51
View File
@@ -1,5 +1,3 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E offline tests for Omni model with video input and audio output.
"""
@@ -12,73 +10,42 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
from pathlib import Path
import pytest
from vllm.assets.video import VideoAsset
from tests.conftest import (
generate_synthetic_video,
)
from tests.utils import hardware_test
from vllm_omni.platforms import current_omni_platform
from .conftest import OmniRunner
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
if current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
else:
stage_configs = [str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")]
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")]
# Create parameter combinations for model and stage config
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
def get_question(prompt_type="video"):
prompts = {
"video": "Describe the video briefly.",
}
return prompts.get(prompt_type, prompts["video"])
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("test_config", test_params)
def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None:
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
def test_video_to_audio(omni_runner, omni_runner_handler) -> None:
"""Test processing video, generating audio output."""
model, stage_config_path = test_config
with omni_runner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner:
# Prepare inputs
question = "Describe the video briefly."
video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
video = generate_synthetic_video(224, 224, 300)["np_array"]
outputs = runner.generate_multimodal(
prompts=question,
videos=video,
)
request_config = {"prompts": get_question(), "videos": video, "modalities": ["audio"]}
# Find and verify text output (thinker stage)
text_output = None
output_count = 0
for stage_output in outputs:
if stage_output.final_output_type == "text":
text_output = stage_output
output_count += 1
break
assert output_count > 0
assert text_output is not None
assert len(text_output.request_output) > 0
text_content = text_output.request_output[0].outputs[0].text
assert text_content is not None
assert len(text_content.strip()) > 0
# Find and verify audio output (code2wav stage)
audio_output = None
output_count = 0
for stage_output in outputs:
if stage_output.final_output_type == "audio":
audio_output = stage_output
output_count += 1
break
assert output_count > 0
assert audio_output is not None
assert len(audio_output.request_output) > 0
# Verify audio tensor exists and has content
audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"]
assert audio_tensor is not None
assert audio_tensor.numel() > 0
# Test single completion
omni_runner_handler.send_request(request_config)
tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml
-103
View File
@@ -1,103 +0,0 @@
# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
# Stage 0: Thinker (multimodal understanding + text generation)
# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
# The following config has been verified on 2x H100-80G GPUs.
stage_args:
- stage_id: 0
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "0"
max_batch_size: 5
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: false
trust_remote_code: true
engine_output_type: latent # Output hidden states for talker
distributed_executor_backend: "mp"
enable_prefix_caching: false
max_num_batched_tokens: 32768
hf_config_name: thinker_config
tensor_parallel_size: 1
load_format: dummy
final_output: true
final_output_type: text
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 100
seed: 42
detokenize: True
repetition_penalty: 1.05
- stage_id: 1
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "1"
max_batch_size: 5
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
enforce_eager: false
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
# tensor_parallel_size: 2
enable_prefix_caching: false
distributed_executor_backend: "mp"
hf_config_name: talker_config
load_format: dummy
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
# final_output: true
# final_output_type: text
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 100
seed: 42
detokenize: False
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 2
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
async_scheduling: false
enable_prefix_caching: false
engine_output_type: audio # Final output: audio waveform
gpu_memory_utilization: 0.1
distributed_executor_backend: "mp"
max_num_batched_tokens: 1000000
hf_config_name: thinker_config
load_format: dummy
engine_input_source: [1]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 200
seed: 42
detokenize: True
repetition_penalty: 1.1
tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml
-95
View File
@@ -1,95 +0,0 @@
# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
# Stage 0: Thinker (multimodal understanding + text generation)
# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
# The following config has been verified on 2x H100-80G GPUs.
stage_args:
- stage_id: 0
runtime:
devices: "0"
max_batch_size: 5
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.9
enforce_eager: false
trust_remote_code: true
engine_output_type: latent # Output hidden states for talker
distributed_executor_backend: "mp"
enable_prefix_caching: false
hf_config_name: thinker_config
tensor_parallel_size: 1
final_output: true
final_output_type: text
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 100
seed: 42
detokenize: True
repetition_penalty: 1.05
- stage_id: 1
runtime:
devices: "1"
max_batch_size: 5
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6
enforce_eager: true
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
# tensor_parallel_size: 2
enable_prefix_caching: false
distributed_executor_backend: "mp"
hf_config_name: talker_config
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
# final_output: true
# final_output_type: text
default_sampling_params:
temperature: 0.9
top_k: 50
max_tokens: 1000
seed: 42
detokenize: False
repetition_penalty: 1.05
stop_token_ids: [2150]
- stage_id: 2
runtime:
devices: "1"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio # Final output: audio waveform
gpu_memory_utilization: 0.1
distributed_executor_backend: "mp"
max_num_batched_tokens: 1000000
hf_config_name: thinker_config
async_scheduling: false
engine_input_source: [1]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 2000
seed: 42
detokenize: True
repetition_penalty: 1.1
tests/e2e/online_serving/test_qwen3_omni.py
+26 -158
View File
@@ -1,5 +1,3 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E Online tests for Qwen3-Omni model with video input and audio output.
"""
@@ -9,23 +7,15 @@ import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
import concurrent.futures
import threading
import time
from pathlib import Path
import openai
import pytest
from tests.conftest import (
OmniServer,
convert_audio_to_text,
cosine_similarity_text,
dummy_messages_from_mix_data,
generate_synthetic_audio,
generate_synthetic_image,
generate_synthetic_video,
merge_base64_and_convert_to_text,
modify_stage_config,
)
from tests.utils import hardware_test
@@ -34,13 +24,9 @@ from vllm_omni.platforms import current_omni_platform
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
def get_default_config():
return str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
def get_chunk_config():
path = modify_stage_config(
get_default_config(),
str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml"),
updates={
"async_chunk": True,
"stage_args": {
@@ -57,49 +43,17 @@ def get_chunk_config():
return path
CHUNK_CONFIG_PATH = get_chunk_config()
# CI stage config for 2xH100-80G GPUs or AMD GPU MI325
if current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
else:
stage_configs = [get_default_config(), CHUNK_CONFIG_PATH]
stage_configs = [get_chunk_config()]
# Create parameter combinations for model and stage config
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
_omni_server_lock = threading.Lock()
@pytest.fixture(scope="module")
def omni_server(request):
"""Start vLLM-Omni server as a subprocess with actual model weights.
Uses session scope so the server starts only once for the entire test session.
Multi-stage initialization can take 10-20+ minutes.
"""
with _omni_server_lock:
model, stage_config_path = request.param
print(f"Starting OmniServer with model: {model}")
with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
print("OmniServer started successfully")
yield server
print("OmniServer stopping...")
print("OmniServer stopped")
@pytest.fixture
def client(omni_server):
"""OpenAI client for the running vLLM-Omni server."""
return openai.OpenAI(
base_url=f"http://{omni_server.host}:{omni_server.port}/v1",
api_key="EMPTY",
)
def get_system_prompt():
return {
"role": "system",
@@ -116,23 +70,6 @@ def get_system_prompt():
}
def dummy_messages_from_video_data(
video_data_url: str,
content_text: str = "Describe the video briefly.",
):
"""Create messages with video data URL for OpenAI API."""
return [
get_system_prompt(),
{
"role": "user",
"content": [
{"type": "video_url", "video_url": {"url": video_data_url}},
{"type": "text", "text": content_text},
],
},
]
def get_prompt(prompt_type="text_only"):
prompts = {
"text_only": "What is the capital of China? Answer in 20 words.",
@@ -146,11 +83,12 @@ def get_max_batch_size(size_type="few"):
return batch_sizes.get(size_type, 5)
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
"""
Test multi-modal input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
@@ -160,8 +98,6 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
Datasets: single request
"""
# Test single completion
e2e_list = list()
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
@@ -173,57 +109,25 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
content_text=get_prompt("mix"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"key_words": {
"audio": ["water", "chirping"],
"image": ["square", "quadrate"],
},
}
# Test single completion
start_time = time.perf_counter()
chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, stream=True)
text_content = ""
audio_data = []
for chunk in chat_completion:
for choice in chunk.choices:
if hasattr(choice, "delta"):
content = getattr(choice.delta, "content", None)
else:
content = None
modality = getattr(chunk, "modality", None)
if modality == "audio" and content:
audio_data.append(content)
elif modality == "text" and content:
# Text chunk - accumulate text content
text_content += content if content else ""
# Verify E2E
current_e2e = time.perf_counter() - start_time
print(f"the request e2e is: {current_e2e}")
# TODO: Verify the E2E latency after confirmation baseline.
e2e_list.append(current_e2e)
print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
# Verify all completions succeeded
assert audio_data is not None, "No audio output is generated"
# Verify text output success
assert text_content is not None and len(text_content) >= 2, "No text output is generated"
assert any(
keyword in text_content.lower() for keyword in ["square", "quadrate", "sphere", "globe", "circle", "round"]
), "The output does not contain any of the keywords."
# Verify text output same as audio output
audio_content = merge_base64_and_convert_to_text(audio_data)
print(f"text content is: {text_content}")
print(f"audio content is: {audio_content}")
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
print(f"similarity is: {similarity}")
assert similarity > 0.9, "The audio content is not same as the text"
openai_client.send_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
def test_text_to_text_001(omni_server, openai_client) -> None:
"""
Test text input processing and text/audio output generation via OpenAI API.
Deploy Setting: default yaml
@@ -231,50 +135,14 @@ def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
Output Modal: text + audio
Datasets: few requests
"""
num_concurrent_requests = get_max_batch_size()
messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())
e2e_list = list()
with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
# Submit multiple completion requests concurrently
futures = [
executor.submit(client.chat.completions.create, model=omni_server.model, messages=messages)
for _ in range(num_concurrent_requests)
]
start_time = time.perf_counter()
# Wait for all requests to complete and collect results
chat_completions = list()
for future in concurrent.futures.as_completed(futures):
chat_completions.append(future.result())
# Verify E2E
current_e2e = time.perf_counter() - start_time
print(f"the request e2e is: {current_e2e}")
# TODO: Verify the E2E latency after confirmation baseline.
e2e_list.append(current_e2e)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": False,
"modalities": ["text"],
"key_words": {"text": ["beijing"]},
}
print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
# Verify all completions succeeded
assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
for chat_completion in chat_completions:
# Verify audio output success
audio_data = None
text_content = None
for choice in chat_completion.choices:
if choice.message.audio is not None:
audio_message = choice.message
audio_data = audio_message.audio.data
assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
if choice.message.content is not None:
# Verify text output success
text_content = choice.message.content
assert "beijing" in text_content.lower(), "The output do not contain keywords."
# Verify text output same as audio output
audio_content = convert_audio_to_text(audio_data)
print(f"text content is: {text_content}")
print(f"audio content is: {audio_content}")
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
print(f"similarity is: {similarity}")
assert similarity > 0.9, "The audio content is not same as the text"
openai_client.send_request(request_config, request_num=get_max_batch_size())
tests/e2e/stage_configs/qwen2_5_omni_ci.yaml
+7 -7
View File
@@ -11,10 +11,10 @@ stage_args:
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 32768
max_num_batched_tokens: 32768
max_model_len: 2400
max_num_batched_tokens: 2400
max_num_seqs: 1
gpu_memory_utilization: 0.8
skip_mm_profiling: true
@@ -41,10 +41,10 @@ stage_args:
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 32768
max_num_batched_tokens: 32768
max_model_len: 2400
max_num_batched_tokens: 2400
max_num_seqs: 1
gpu_memory_utilization: 0.8
skip_mm_profiling: true
@@ -71,7 +71,7 @@ stage_args:
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.15
enforce_eager: true
tests/e2e/stage_configs/qwen3_omni_ci.yaml
+2
View File
@@ -24,6 +24,7 @@ stage_args:
enable_prefix_caching: false
hf_config_name: thinker_config
tensor_parallel_size: 1
load_format: dummy
final_output: true
final_output_type: text
is_comprehension: true
@@ -55,6 +56,7 @@ stage_args:
max_model_len: 32768
distributed_executor_backend: "mp"
hf_config_name: talker_config
load_format: dummy
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
default_sampling_params:
tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml → tests/e2e/stage_configs/qwen3_omni_thinker_ci.yaml
View File
File renamed without changes.
tests/e2e/offline_inference/stage_configs/rocm/qwen2_5_omni_ci.yaml → tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml
+4 -4
View File
@@ -13,8 +13,8 @@ stage_args:
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
max_model_len: 2400
max_num_batched_tokens: 2400
max_num_seqs: 1
gpu_memory_utilization: 0.8
skip_mm_profiling: true
@@ -43,8 +43,8 @@ stage_args:
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 896
max_num_batched_tokens: 896
max_model_len: 2400
max_num_batched_tokens: 2400
max_num_seqs: 1
gpu_memory_utilization: 0.8
skip_mm_profiling: true
tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml → tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml
View File
File renamed without changes.
tests/engine/__init__.py
View File
Whitespace-only changes.
tests/e2e/online_serving/test_async_omni.py → tests/engine/test_async_omni_engine_abort.py
+1 -1
View File
@@ -15,7 +15,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
SEED = 42
stage_config = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_thinker_ci.yaml")
stage_config = str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen3_omni_thinker_ci.yaml")
model = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
tests/entrypoints/openai_api/test_serving_chat_metrics.py
+52
View File
@@ -0,0 +1,52 @@
# SPDX-License-Identifier: Apache-2.0
"""Unit tests for OmniChatCompletionResponse/StreamResponse metrics field."""
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_omni_chat_completion_response_metrics():
"""Test OmniChatCompletionResponse metrics field works correctly."""
from vllm.entrypoints.openai.engine.protocol import UsageInfo
from vllm_omni.entrypoints.openai.protocol.chat_completion import (
OmniChatCompletionResponse,
)
usage = UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0)
# Default is None
response = OmniChatCompletionResponse(id="test-id", created=1234567890, model="test-model", choices=[], usage=usage)
assert response.metrics is None
# Can set metrics and serialize
test_metrics = {"thinker_ttft": 0.123, "talker_ttft": 0.456}
response = OmniChatCompletionResponse(
id="test-id",
created=1234567890,
model="test-model",
choices=[],
usage=usage,
metrics=test_metrics,
)
assert response.metrics == test_metrics
assert "thinker_ttft" in response.model_dump_json()
def test_omni_chat_completion_stream_response_metrics():
"""Test OmniChatCompletionStreamResponse metrics and modality fields."""
from vllm_omni.entrypoints.openai.protocol.chat_completion import (
OmniChatCompletionStreamResponse,
)
response = OmniChatCompletionStreamResponse(
id="test-id",
created=1234567890,
model="test-model",
choices=[],
modality="audio",
metrics={"stage_latency": 0.5},
)
assert response.modality == "audio"
assert response.metrics == {"stage_latency": 0.5}
vllm_omni/diffusion/data.py
+37 -1
View File
@@ -4,7 +4,7 @@
import enum
import os
import random
from collections.abc import Callable
from collections.abc import Callable, Mapping
from dataclasses import dataclass, field, fields
from typing import Any
@@ -14,6 +14,10 @@ from typing_extensions import Self
from vllm.config.utils import config
from vllm.logger import init_logger
from vllm_omni.diffusion.quantization import (
DiffusionQuantizationConfig,
get_diffusion_quant_config,
)
from vllm_omni.diffusion.utils.network_utils import is_port_available
logger = init_logger(__name__)
@@ -369,6 +373,11 @@ class OmniDiffusionConfig:
# Omni configuration (injected from stage config)
omni_kv_config: dict[str, Any] = field(default_factory=dict)
# Quantization settings
# Supported methods: "fp8" (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs)
quantization: str | None = None
quantization_config: "DiffusionQuantizationConfig | dict[str, Any] | None" = None
def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:
"""
Find an available port with retry logic.
@@ -455,6 +464,33 @@ class OmniDiffusionConfig:
# If it's neither dict nor DiffusionCacheConfig, convert to empty config
self.cache_config = DiffusionCacheConfig()
# Convert quantization config
if self.quantization is not None or self.quantization_config is not None:
# Handle dict or DictConfig (from OmegaConf) - use Mapping for broader compatibility
if isinstance(self.quantization_config, Mapping):
# Convert DictConfig to dict if needed (OmegaConf compatibility)
config_dict = dict(self.quantization_config)
# Use get() instead of pop() to avoid mutating original dict
quant_method = config_dict.get("method", self.quantization)
# Filter out "method" key for kwargs
quant_kwargs = {k: v for k, v in config_dict.items() if k != "method"}
# Validate conflicting methods
if self.quantization is not None and quant_method is not None and quant_method != self.quantization:
logger.warning(
f"Conflicting quantization methods: quantization={self.quantization!r}, "
f"quantization_config['method']={quant_method!r}. Using quantization_config['method']."
)
self.quantization_config = get_diffusion_quant_config(quant_method, **quant_kwargs)
elif self.quantization_config is None and self.quantization is not None:
self.quantization_config = get_diffusion_quant_config(self.quantization)
elif not isinstance(self.quantization_config, DiffusionQuantizationConfig):
raise TypeError(
f"quantization_config must be a DiffusionQuantizationConfig, dict, or None, "
f"got {type(self.quantization_config)!r}"
)
if self.max_cpu_loras is None:
self.max_cpu_loras = 1
elif self.max_cpu_loras < 1:
vllm_omni/diffusion/model_loader/diffusers_loader.py
+29
View File
@@ -13,6 +13,7 @@ from torch import nn
from vllm.config import ModelConfig
from vllm.config.load import LoadConfig
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
from vllm.model_executor.model_loader.weight_utils import (
download_safetensors_index_file_from_hf,
download_weights_from_hf,
@@ -217,8 +218,36 @@ class DiffusersPipelineLoader:
logger.debug("Loading weights on %s ...", load_device)
# Quantization does not happen in `load_weights` but after it
self.load_weights(model)
# Process weights after loading for quantization (e.g., FP8 online quantization)
# This is needed for vLLM's quantization methods that need to transform weights
self._process_weights_after_loading(model, target_device)
return model.eval()
def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None:
"""Process weights after loading for quantization methods.
This handles vLLM's quantization methods that need to process weights
after loading (e.g., FP8 online quantization from BF16/FP16 weights).
"""
for _, module in model.named_modules():
quant_method = getattr(module, "quant_method", None)
if isinstance(quant_method, QuantizeMethodBase):
# Move module to target device for processing if needed
module_device = next(module.parameters(), None)
if module_device is not None:
module_device = module_device.device
needs_device_move = module_device != target_device
if needs_device_move:
module.to(target_device)
quant_method.process_weights_after_loading(module)
if needs_device_move:
module.to(module_device)
def load_weights(self, model: nn.Module) -> None:
weights_to_load = {name for name, _ in model.named_parameters()}
loaded_weights = model.load_weights(self.get_all_weights(model))
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+5 -1
View File
@@ -33,6 +33,7 @@ from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
QwenImageTransformer2DModel,
)
from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -272,7 +273,10 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin):
self.device
)
transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
self.transformer = QwenImageTransformer2DModel(
od_config=od_config, quant_config=quant_config, **transformer_kwargs
)
self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py
+42 -5
View File
@@ -1,10 +1,12 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
from collections.abc import Iterable
from functools import lru_cache
from math import prod
from typing import Any
from typing import TYPE_CHECKING, Any
import torch
import torch.nn as nn
@@ -23,6 +25,11 @@ from vllm.model_executor.layers.linear import (
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
)
from vllm_omni.diffusion.attention.backends.abstract import (
AttentionMetadata,
)
@@ -390,7 +397,16 @@ class QwenEmbedRope(nn.Module):
class ColumnParallelApproxGELU(nn.Module):
def __init__(self, dim_in: int, dim_out: int, *, approximate: str, bias: bool = True):
def __init__(
self,
dim_in: int,
dim_out: int,
*,
approximate: str,
bias: bool = True,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
):
super().__init__()
self.proj = ColumnParallelLinear(
dim_in,
@@ -398,6 +414,8 @@ class ColumnParallelApproxGELU(nn.Module):
bias=bias,
gather_output=False,
return_bias=False,
quant_config=quant_config,
prefix=prefix,
)
self.approximate = approximate
@@ -415,6 +433,8 @@ class FeedForward(nn.Module):
activation_fn: str = "gelu-approximate",
inner_dim: int | None = None,
bias: bool = True,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> None:
super().__init__()
@@ -424,13 +444,17 @@ class FeedForward(nn.Module):
dim_out = dim_out or dim
layers: list[nn.Module] = [
ColumnParallelApproxGELU(dim, inner_dim, approximate="tanh", bias=bias),
ColumnParallelApproxGELU(
dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config, prefix=prefix
),
nn.Identity(), # placeholder for weight loading
RowParallelLinear(
inner_dim,
dim_out,
input_is_parallel=True,
return_bias=False,
quant_config=quant_config,
prefix=prefix,
),
]
@@ -456,6 +480,7 @@ class QwenImageCrossAttention(nn.Module):
pre_only: bool = False,
context_pre_only: bool = False,
out_dim: int | None = None,
quant_config: QuantizationConfig | None = None,
) -> None:
super().__init__()
assert dim % num_heads == 0
@@ -471,6 +496,8 @@ class QwenImageCrossAttention(nn.Module):
hidden_size=dim,
head_size=self.head_dim,
total_num_heads=num_heads,
quant_config=quant_config,
prefix="to_qkv",
)
self.query_num_heads = self.to_qkv.num_heads
self.kv_num_heads = self.to_qkv.num_kv_heads
@@ -485,6 +512,8 @@ class QwenImageCrossAttention(nn.Module):
hidden_size=added_kv_proj_dim,
head_size=head_dim,
total_num_heads=num_heads,
quant_config=quant_config,
prefix="add_kv_proj",
)
self.add_query_num_heads = self.add_kv_proj.num_heads
self.add_kv_num_heads = self.add_kv_proj.num_kv_heads
@@ -496,6 +525,8 @@ class QwenImageCrossAttention(nn.Module):
bias=out_bias,
input_is_parallel=True,
return_bias=False,
quant_config=quant_config,
prefix="to_add_out",
)
assert not pre_only
@@ -505,6 +536,8 @@ class QwenImageCrossAttention(nn.Module):
bias=out_bias,
input_is_parallel=True,
return_bias=False,
quant_config=quant_config,
prefix="to_out",
)
self.norm_added_q = RMSNorm(head_dim, eps=eps)
@@ -637,6 +670,7 @@ class QwenImageTransformerBlock(nn.Module):
qk_norm: str = "rms_norm",
eps: float = 1e-6,
zero_cond_t: bool = False,
quant_config: QuantizationConfig | None = None,
):
super().__init__()
@@ -656,9 +690,10 @@ class QwenImageTransformerBlock(nn.Module):
added_kv_proj_dim=dim,
context_pre_only=False,
head_dim=attention_head_dim,
quant_config=quant_config,
)
self.img_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
self.img_mlp = FeedForward(dim=dim, dim_out=dim)
self.img_mlp = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix="img_mlp")
# Text processing modules
self.txt_mod = nn.Sequential(
@@ -668,7 +703,7 @@ class QwenImageTransformerBlock(nn.Module):
self.txt_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
# Text doesn't need separate attention - it's handled by img_attn joint computation
self.txt_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps)
self.txt_mlp = FeedForward(dim=dim, dim_out=dim)
self.txt_mlp = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix="txt_mlp")
self.zero_cond_t = zero_cond_t
@@ -862,6 +897,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
zero_cond_t: bool = False,
use_additional_t_cond: bool = False,
use_layer3d_rope: bool = False,
quant_config: QuantizationConfig | None = None,
):
super().__init__()
self.parallel_config = od_config.parallel_config
@@ -891,6 +927,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
zero_cond_t=zero_cond_t,
quant_config=quant_config,
)
for _ in range(num_layers)
]
vllm_omni/diffusion/models/z_image/pipeline_z_image.py
+4 -1
View File
@@ -37,6 +37,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL
from vllm_omni.diffusion.models.z_image.z_image_transformer import (
ZImageTransformer2DModel,
)
from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
@@ -173,7 +174,9 @@ class ZImagePipeline(nn.Module):
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
self._execution_device
)
self.transformer = ZImageTransformer2DModel()
# Get vLLM quantization config for linear layers
quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
self.transformer = ZImageTransformer2DModel(quant_config=quant_config)
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
# Note: Context parallelism is applied centrally in registry.initialize_model()
vllm_omni/diffusion/models/z_image/z_image_transformer.py
+36 -3
View File
@@ -18,6 +18,7 @@
import math
from collections.abc import Iterable
from typing import TYPE_CHECKING
import torch
import torch.nn as nn
@@ -32,6 +33,11 @@ from vllm.model_executor.layers.linear import (
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
)
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.cache.base import CachedTransformer
from vllm_omni.diffusion.distributed.sp_plan import (
@@ -250,6 +256,7 @@ class ZImageAttention(nn.Module):
num_kv_heads: int,
qk_norm: bool = True,
eps: float = 1e-6,
quant_config: "QuantizationConfig | None" = None,
) -> None:
super().__init__()
self.dim = dim
@@ -264,6 +271,7 @@ class ZImageAttention(nn.Module):
total_num_heads=num_heads,
total_num_kv_heads=num_kv_heads,
bias=False,
quant_config=quant_config,
)
assert qk_norm is True
@@ -281,6 +289,7 @@ class ZImageAttention(nn.Module):
bias=False,
input_is_parallel=True,
return_bias=False,
quant_config=quant_config,
)
]
)
@@ -343,13 +352,19 @@ class ZImageAttention(nn.Module):
class FeedForward(nn.Module):
def __init__(self, dim: int, hidden_dim: int):
def __init__(
self,
dim: int,
hidden_dim: int,
quant_config: "QuantizationConfig | None" = None,
):
super().__init__()
self.w13 = MergedColumnParallelLinear(
dim,
[hidden_dim] * 2,
bias=False,
return_bias=False,
quant_config=quant_config,
)
self.act = SiluAndMul()
self.w2 = RowParallelLinear(
@@ -358,6 +373,7 @@ class FeedForward(nn.Module):
bias=False,
input_is_parallel=True,
return_bias=False,
quant_config=quant_config,
)
def forward(self, x):
@@ -374,6 +390,7 @@ class ZImageTransformerBlock(nn.Module):
norm_eps: float,
qk_norm: bool,
modulation=True,
quant_config: "QuantizationConfig | None" = None,
):
super().__init__()
self.dim = dim
@@ -384,9 +401,14 @@ class ZImageTransformerBlock(nn.Module):
num_kv_heads=n_kv_heads,
qk_norm=qk_norm,
eps=1e-5,
quant_config=quant_config,
)
self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8))
self.feed_forward = FeedForward(
dim=dim,
hidden_dim=int(dim / 3 * 8),
quant_config=quant_config,
)
self.layer_id = layer_id
self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
@@ -589,6 +611,7 @@ class ZImageTransformer2DModel(CachedTransformer):
t_scale=1000.0,
axes_dims=[32, 48, 48],
axes_lens=[1024, 512, 512],
quant_config: "QuantizationConfig | None" = None,
) -> None:
super().__init__()
self.dtype = torch.bfloat16
@@ -648,6 +671,7 @@ class ZImageTransformer2DModel(CachedTransformer):
norm_eps,
qk_norm,
modulation=True,
quant_config=quant_config,
)
for layer_id in range(n_refiner_layers)
]
@@ -662,6 +686,7 @@ class ZImageTransformer2DModel(CachedTransformer):
norm_eps,
qk_norm,
modulation=False,
quant_config=quant_config,
)
for layer_id in range(n_refiner_layers)
]
@@ -677,7 +702,15 @@ class ZImageTransformer2DModel(CachedTransformer):
self.layers = nn.ModuleList(
[
ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm)
ZImageTransformerBlock(
layer_id,
dim,
n_heads,
n_kv_heads,
norm_eps,
qk_norm,
quant_config=quant_config,
)
for layer_id in range(n_layers)
]
)
vllm_omni/diffusion/quantization/__init__.py
+114
View File
@@ -0,0 +1,114 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Quantization support for diffusion models.
This module provides a unified interface for quantizing diffusion transformers
using various methods (FP8, etc.). It wraps vLLM's quantization infrastructure
while allowing diffusion-model-specific defaults and optimizations.
Example usage:
from vllm_omni.diffusion.quantization import (
get_diffusion_quant_config,
get_vllm_quant_config_for_layers,
)
# Create FP8 config for diffusion model
diff_config = get_diffusion_quant_config("fp8")
# Get vLLM config to pass to linear layers
vllm_config = get_vllm_quant_config_for_layers(diff_config)
# Use in model initialization
linear_layer = QKVParallelLinear(..., quant_config=vllm_config)
"""
from typing import TYPE_CHECKING
from vllm.logger import init_logger
from .base import DiffusionQuantizationConfig
from .fp8 import DiffusionFp8Config
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
)
logger = init_logger(__name__)
# Registry of supported quantization methods
# To add a new method, create a new config class and register it here
_QUANT_CONFIG_REGISTRY: dict[str, type[DiffusionQuantizationConfig]] = {
"fp8": DiffusionFp8Config,
}
SUPPORTED_QUANTIZATION_METHODS = list(_QUANT_CONFIG_REGISTRY.keys())
def get_diffusion_quant_config(
quantization: str | None,
**kwargs,
) -> DiffusionQuantizationConfig | None:
"""Factory function to create quantization config for diffusion models.
Args:
quantization: Quantization method name ("fp8", etc.) or None to disable
**kwargs: Method-specific parameters passed to the config constructor
Returns:
DiffusionQuantizationConfig instance or None if quantization is disabled
Raises:
ValueError: If the quantization method is not supported
Example:
# Default FP8 with dynamic activation scaling
config = get_diffusion_quant_config("fp8")
# FP8 with custom parameters
config = get_diffusion_quant_config(
"fp8",
activation_scheme="static",
ignored_layers=["proj_out"],
)
"""
if quantization is None or quantization.lower() == "none":
return None
quantization = quantization.lower()
if quantization not in _QUANT_CONFIG_REGISTRY:
raise ValueError(
f"Unknown quantization method: {quantization!r}. Supported methods: {SUPPORTED_QUANTIZATION_METHODS}"
)
config_cls = _QUANT_CONFIG_REGISTRY[quantization]
logger.info("Creating diffusion quantization config: %s", quantization)
return config_cls(**kwargs)
def get_vllm_quant_config_for_layers(
diffusion_quant_config: DiffusionQuantizationConfig | None,
) -> "QuantizationConfig | None":
"""Get the vLLM QuantizationConfig to pass to linear layers.
This extracts the underlying vLLM config from a DiffusionQuantizationConfig,
which can then be passed to vLLM linear layers (QKVParallelLinear, etc.).
Args:
diffusion_quant_config: The diffusion quantization config, or None
Returns:
vLLM QuantizationConfig instance, or None if input is None
"""
if diffusion_quant_config is None:
return None
return diffusion_quant_config.get_vllm_quant_config()
__all__ = [
"DiffusionQuantizationConfig",
"DiffusionFp8Config",
"get_diffusion_quant_config",
"get_vllm_quant_config_for_layers",
"SUPPORTED_QUANTIZATION_METHODS",
]
vllm_omni/diffusion/quantization/base.py
+61
View File
@@ -0,0 +1,61 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Base class for diffusion model quantization configurations."""
from abc import ABC
from typing import TYPE_CHECKING, ClassVar
import torch
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import (
QuantizationConfig,
)
class DiffusionQuantizationConfig(ABC):
"""Base class for diffusion model quantization configurations.
This provides a thin wrapper over vLLM's quantization configs,
allowing diffusion-model-specific defaults and future extensibility.
Subclasses should:
- Set quant_config_cls to the vLLM QuantizationConfig class
- Call super().__init__() after creating self._vllm_config
- Optionally override get_name() and get_min_capability() if needed
"""
# Subclasses should set this to the vLLM QuantizationConfig class
quant_config_cls: ClassVar[type["QuantizationConfig"] | None] = None
# The underlying vLLM config instance
_vllm_config: "QuantizationConfig | None" = None
@classmethod
def get_name(cls) -> str:
"""Return the quantization method name (e.g., 'fp8', 'int8').
By default, delegates to the underlying vLLM config class.
"""
if cls.quant_config_cls is not None:
return cls.quant_config_cls.get_name()
raise NotImplementedError("Subclass must set quant_config_cls or override get_name()")
def get_vllm_quant_config(self) -> "QuantizationConfig | None":
"""Return the underlying vLLM QuantizationConfig for linear layers."""
return self._vllm_config
@classmethod
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
"""Return supported activation dtypes."""
return [torch.bfloat16, torch.float16]
@classmethod
def get_min_capability(cls) -> int:
"""Minimum GPU compute capability required.
By default, delegates to the underlying vLLM config class.
"""
if cls.quant_config_cls is not None:
return cls.quant_config_cls.get_min_capability()
return 80 # Ampere default
vllm_omni/diffusion/quantization/fp8.py
+50
View File
@@ -0,0 +1,50 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""FP8 quantization config for diffusion transformers."""
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
from .base import DiffusionQuantizationConfig
class DiffusionFp8Config(DiffusionQuantizationConfig):
"""FP8 quantization config optimized for diffusion transformers.
Uses dynamic activation scaling (no calibration dataset needed) and
online weight quantization from BF16/FP16 checkpoints.
Device Compatibility:
- Turing (SM 75+): Weight-only FP8 via Marlin kernel
- Ada/Hopper (SM 89+): Full W8A8 FP8 with native hardware support
The kernel selection is automatic based on GPU capability.
Args:
activation_scheme: Activation quantization scheme.
- "dynamic": Per-token dynamic scaling (default, no calibration)
- "static": Single per-tensor scale (requires calibration)
weight_block_size: Block size for block-wise weight quantization.
Format: [block_n, block_k]. If None, uses per-tensor scaling.
ignored_layers: List of layer name patterns to skip quantization.
"""
# Tight coupling with vLLM's Fp8Config - delegates get_name() and get_min_capability()
quant_config_cls = Fp8Config
def __init__(
self,
activation_scheme: str = "dynamic",
weight_block_size: list[int] | None = None,
ignored_layers: list[str] | None = None,
):
self.activation_scheme = activation_scheme
self.weight_block_size = weight_block_size
self.ignored_layers = ignored_layers or []
# Create underlying vLLM FP8 config
self._vllm_config = Fp8Config(
is_checkpoint_fp8_serialized=False, # Online quantization from BF16
activation_scheme=activation_scheme,
weight_block_size=weight_block_size,
ignored_layers=ignored_layers,
)
vllm_omni/entrypoints/async_omni.py
+18 -30
View File
@@ -32,7 +32,7 @@ from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams
# Internal imports (our code)
from vllm_omni.lora.request import LoRARequest
from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats
from vllm_omni.metrics import OrchestratorAggregator
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -402,9 +402,6 @@ class AsyncOmni(OmniBase):
all_stages_finished[stage_id] = finished
if output_to_yield:
metrics.record_audio_generated_frames(
output_to_yield, engine_outputs.finished, stage_id, request_id
)
yield output_to_yield
async def _process_sequential_results(
@@ -428,9 +425,6 @@ class AsyncOmni(OmniBase):
metrics,
)
if output_to_yield:
metrics.record_audio_generated_frames(
output_to_yield, engine_outputs.finished, stage_id, request_id
)
yield output_to_yield
if not isinstance(engine_outputs, list):
engine_outputs = [engine_outputs]
@@ -504,29 +498,6 @@ class AsyncOmni(OmniBase):
finished = engine_outputs.finished
# Mark last output time
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
try:
_m: StageRequestStats = result.get("metrics")
if _m is not None:
# Accumulate generation time
metrics.accumulated_gen_time_ms[req_id][stage_id] += _m.stage_gen_time_ms
# For diffusion stages, we also accumulate diffusion time
metrics.accumulate_diffusion_metrics(stage.stage_type, req_id, engine_outputs)
if finished:
metrics.on_stage_metrics(stage_id, req_id, _m, stage.final_output_type)
except Exception as e:
logger.exception(
f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
)
logger.debug(
f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
)
output_to_yield = None
if getattr(stage, "final_output", False):
@@ -551,6 +522,23 @@ class AsyncOmni(OmniBase):
final_output_type=stage.final_output_type,
request_output=engine_outputs,
)
# Mark last output time
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
metrics.process_stage_metrics(
result=result,
stage_type=stage.stage_type,
stage_id=stage_id,
req_id=req_id,
engine_outputs=engine_outputs,
finished=finished,
final_output_type=stage.final_output_type,
output_to_yield=output_to_yield,
)
logger.debug(
f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
)
return engine_outputs, finished, output_to_yield
vllm_omni/entrypoints/chat_utils.py
+48
View File
@@ -22,6 +22,9 @@ from vllm.entrypoints.chat_utils import (
_postprocess_messages,
_ToolParser,
)
from vllm.logger import init_logger
logger = init_logger(__name__)
class OmniAsyncMultiModalItemTracker(AsyncMultiModalItemTracker):
@@ -135,6 +138,48 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser):
await asyncio.to_thread(_cleanup_file_sync, temp_video_file_path)
def _ensure_system_prompt(
messages: list[ChatCompletionMessageParam],
model_config: ModelConfig,
) -> list[ChatCompletionMessageParam]:
"""
Ensure a system prompt exists for Qwen-Omni models to preserve precision of the model.
Args:
messages: List of chat messages
model_config: Model configuration
Returns:
Messages list with system prompt
"""
model_name = getattr(model_config, "model", "").lower()
hf_config = getattr(model_config, "hf_config", None)
architectures = getattr(hf_config, "architectures", []) if hf_config else []
is_qwen_omni = ("qwen" in model_name and "omni" in model_name) or any(
"qwen" in arch.lower() and "omni" in arch.lower() for arch in architectures
)
if not is_qwen_omni:
return messages
if messages and messages[0].get("role") == "system":
return messages
default_qwen_omni_system_prompt = (
"You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
"Group, capable of perceiving auditory and visual inputs, as well as "
"generating text and speech."
)
system_message: ChatCompletionMessageParam = {
"role": "system",
"content": default_qwen_omni_system_prompt,
}
logger.info(f"injecting system prompt {default_qwen_omni_system_prompt} for Qwen-Omni model")
return [system_message] + list(messages)
def parse_chat_messages_futures(
messages: list[ChatCompletionMessageParam],
model_config: ModelConfig,
@@ -153,6 +198,9 @@ def parse_chat_messages_futures(
Tuple of (conversation, mm_future) where mm_future resolves to
(mm_data, mm_uuids) when awaited.
"""
# auto-inject system prompt for Qwen-Omni models if missing
messages = _ensure_system_prompt(messages, model_config)
conversation: list[ConversationMessage] = []
mm_tracker = OmniAsyncMultiModalItemTracker(model_config)
vllm_omni/entrypoints/omni.py
+8 -4
View File
@@ -82,8 +82,11 @@ def omni_snapshot_download(model_id) -> str:
# For other cases (Hugging Face), perform a real download to ensure all
# necessary files (including *.pt for audio/diffusion) are available locally
# before stage workers are spawned. This prevents initialization timeouts.
# Return the original model_id so that model_config.model preserves
# HuggingFace semantics (e.g. "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice")
# instead of the resolved cache path.
try:
return download_weights_from_hf_specific(
download_weights_from_hf_specific(
model_name_or_path=model_id,
cache_dir=None,
allow_patterns=["*"],
@@ -91,7 +94,7 @@ def omni_snapshot_download(model_id) -> str:
)
except huggingface_hub.errors.RepositoryNotFoundError:
logger.warning(f"Repository not found for '{model_id}'.")
return model_id
return model_id
class OmniBase:
@@ -818,7 +821,7 @@ class Omni(OmniBase):
request_output=engine_outputs,
)
# Record audio generated frames with unified signature
# Record audio generated frames (only when finished)
try:
finished = (
engine_outputs.finished
@@ -831,7 +834,8 @@ class Omni(OmniBase):
else False
)
)
metrics.record_audio_generated_frames(output_to_yield, finished, stage_id, req_id)
if finished:
metrics.record_audio_generated_frames(output_to_yield, stage_id, req_id)
except Exception as e:
logger.exception(
f"[{self._name}] Failed to record audio metrics for req {req_id} at stage {stage_id}: {e}",
vllm_omni/entrypoints/omni_stage.py
+10 -4
View File
@@ -49,6 +49,7 @@ from vllm_omni.entrypoints.stage_utils import (
maybe_dump_to_shm,
set_stage_devices,
)
from vllm_omni.entrypoints.utils import detect_pid_host
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt
from vllm_omni.metrics import count_tokens_from_outputs
from vllm_omni.outputs import OmniRequestOutput
@@ -66,12 +67,17 @@ def _sequential_init_lock(engine_args: dict[str, Any], stage_init_timeout: int =
"""
from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
if is_process_scoped_memory_available():
logger.debug("NVML process-scoped memory available — concurrent init is safe, skipping locks")
if is_process_scoped_memory_available() and detect_pid_host():
logger.debug(
"NVML process-scoped memory available and PID host is available — concurrent init is safe, skipping locks"
)
yield
return
logger.debug("NVML unavailable — using sequential init locks")
else:
logger.debug(
"NVML unavailable or PID host is not available (usually inside a container, "
"--pid=host is not set in docker run command) — using sequential init locks"
)
from vllm_omni.platforms import current_omni_platform
vllm_omni/entrypoints/openai/protocol/chat_completion.py
+8 -1
View File
@@ -1,5 +1,12 @@
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionStreamResponse
from typing import Any
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse, ChatCompletionStreamResponse
class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse):
modality: str | None = "text"
metrics: dict[str, Any] | None = None
class OmniChatCompletionResponse(ChatCompletionResponse):
metrics: dict[str, Any] | None = None
vllm_omni/entrypoints/openai/serving_chat.py
+14 -4
View File
@@ -16,6 +16,7 @@ from pydantic import TypeAdapter
from vllm.renderers.protocol import BaseRenderer
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
try:
@@ -38,7 +39,6 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionResponse,
ChatCompletionResponseChoice,
ChatCompletionResponseStreamChoice,
ChatCompletionStreamResponse,
ChatMessage,
)
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
@@ -670,6 +670,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
stream_options = request.stream_options
include_usage, include_continuous_usage = should_include_usage(stream_options, self.enable_force_include_usage)
last_metrics: dict[str, Any] | None = None
try:
async for omni_res in result_generator:
final_output_type = omni_res.final_output_type
@@ -678,6 +679,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
logger.warning(f"final output type: {final_output_type} is not needed by the request")
continue
if omni_res.metrics:
last_metrics = omni_res.metrics
if res.prompt_token_ids is not None:
num_prompt_tokens = len(res.prompt_token_ids)
if res.encoder_prompt_token_ids is not None:
@@ -1215,6 +1219,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
choices=[choice_data],
model=model_name,
modality=final_output_type,
metrics=omni_res.metrics,
)
# handle usage stats if requested & if continuous
@@ -1264,13 +1269,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
if self.enable_prompt_tokens_details and num_cached_tokens:
final_usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=num_cached_tokens)
final_usage_chunk = ChatCompletionStreamResponse(
final_usage_chunk = OmniChatCompletionStreamResponse(
id=request_id,
object=chunk_object_type,
created=created_time,
choices=[],
model=model_name,
usage=final_usage,
metrics=last_metrics,
)
final_usage_data = final_usage_chunk.model_dump_json(exclude_unset=True, exclude_none=True)
yield f"data: {final_usage_data}\n\n"
@@ -1318,7 +1324,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
tokenizer: TokenizerLike,
request_metadata: RequestResponseMetadata,
reasoning_parser: ReasoningParser | None = None,
) -> ErrorResponse | ChatCompletionResponse:
) -> ErrorResponse | OmniChatCompletionResponse:
created_time = int(time.time())
final_res: RequestOutput | None = None
@@ -1340,6 +1346,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
prompt_logprobs = None
prompt_token_ids = None
kv_transfer_params = None
response_metrics: dict[str, Any] | None = None
# Build requested modalities set for filtering
requested_modalities = (
@@ -1378,9 +1385,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
else:
logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}")
continue
if omni_outputs.metrics:
response_metrics = omni_outputs.metrics
choices.extend(choices_data)
response = ChatCompletionResponse(
response = OmniChatCompletionResponse(
id=request_id,
created=created_time,
model=model_name,
@@ -1389,6 +1398,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
prompt_logprobs=prompt_logprobs,
prompt_token_ids=prompt_token_ids,
kv_transfer_params=kv_transfer_params,
metrics=response_metrics,
)
# Log complete response if output logging is enabled
vllm_omni/entrypoints/utils.py
+60
View File
@@ -280,3 +280,63 @@ def get_final_stage_id_for_e2e(
final_stage_id_for_e2e = last_stage_id
return final_stage_id_for_e2e
# The following code detects if the process is running in a container and if
# PID host is available. If so, we can use process-scoped memory tracking;
# otherwise we need sequential init locks.
def _read_text(path: str) -> str | None:
try:
with open(path, encoding="utf-8", errors="replace") as f:
return f.read()
except (FileNotFoundError, PermissionError, OSError):
return None
def in_container() -> bool:
# Common Docker signal
if os.path.exists("/.dockerenv"):
return True
# cgroup markers (works for Docker/containerd/K8s/Podman in many setups)
cg = _read_text("/proc/1/cgroup") or ""
markers = ("docker", "containerd", "kubepods", "libpod", "podman")
return any(m in cg for m in markers)
def has_pid_host() -> bool | None:
"""
Returns:
True -> very likely running with --pid=host (host PID namespace)
False -> very likely isolated PID namespace (default)
None -> cannot determine
"""
# Strong signal: in host pid namespace, PID 2 is usually kthreadd
comm2 = _read_text("/proc/2/comm")
if comm2 is not None:
comm2 = comm2.strip()
if comm2 == "kthreadd":
return True
# If PID 2 exists and is NOT kthreadd, we're almost certainly not in host pid ns
return False
# Fallback: check for other low-numbered kernel threads (best-effort)
for pid, name in [(3, "rcu_gp"), (4, "rcu_par_gp"), (10, "ksoftirqd/0")]:
comm = _read_text(f"/proc/{pid}/comm")
if comm is not None:
if comm.strip() == name:
return True
else:
return False
return False
def detect_pid_host() -> bool:
ic = in_container()
if not ic:
return True
return has_pid_host()
vllm_omni/metrics/stats.py
+67 -2
View File
@@ -217,14 +217,12 @@ class OrchestratorAggregator:
def record_audio_generated_frames(
self,
output_to_yield: Any,
finished: bool,
stage_id: int,
request_id: str,
) -> None:
try:
if (
output_to_yield.final_output_type == "audio"
and finished
and (multimodal_output := output_to_yield.multimodal_output.get("audio")) is not None
and len(multimodal_output) > 0
):
@@ -249,6 +247,73 @@ class OrchestratorAggregator:
exc_info=True,
)
def process_stage_metrics(
self,
*,
result: dict[str, Any],
stage_type: str,
stage_id: int,
req_id: str,
engine_outputs: Any,
finished: bool,
final_output_type: str | None,
output_to_yield: Any | None,
) -> None:
"""Process and record stage metrics.
Args:
result: Result dict containing metrics from stage
stage_type: Type of the stage (e.g., 'llm', 'diffusion')
stage_id: Stage identifier
req_id: Request identifier
engine_outputs: Engine output object
finished: Whether stage processing is finished
final_output_type: Type of final output (e.g., 'text', 'audio')
output_to_yield: Output object to attach metrics to
"""
try:
_m: StageRequestStats | None = result.get("metrics")
# 1. Accumulate metrics from stage stats
if _m is not None:
self.accumulated_gen_time_ms[req_id][stage_id] += _m.stage_gen_time_ms
self.accumulate_diffusion_metrics(stage_type, req_id, engine_outputs)
if finished:
self.on_stage_metrics(stage_id, req_id, _m, final_output_type)
# 2. No output to yield, nothing more to do
if output_to_yield is None:
return
# 3. Not finished yet — empty metrics, skip audio recording
if not finished:
output_to_yield.metrics = {}
return
# 4. Finished with output: assign text metrics if available
output_to_yield.metrics = {}
stage_event = next(
(evt for evt in reversed(self.stage_events.get(req_id, [])) if evt.stage_id == stage_id),
None,
)
if stage_event is not None and stage_event.final_output_type == "text":
output_to_yield.metrics = {
"num_tokens_in": stage_event.num_tokens_in,
"num_tokens_out": stage_event.num_tokens_out,
"stage_id": stage_event.stage_id,
"final_output_type": stage_event.final_output_type,
}
# 5. Finished: record audio generated frames
self.record_audio_generated_frames(output_to_yield, stage_id, req_id)
except Exception:
logger.exception(
"Failed to process metrics for stage %s, req %s",
stage_id,
req_id,
)
def _as_stage_request_stats(
self,
stage_id: int,
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+7 -4
View File
@@ -589,6 +589,7 @@ class Qwen3OmniMoeForConditionalGeneration(
input_embeds = self.talker.embed_input_ids(input_ids)
span_len = input_ids.shape[0]
update_dict = {}
if span_len > 1:
# prefill
input_ids, input_embeds, update_dict = self.talker_preprocess_prefill(input_ids, input_embeds, **info_dict)
@@ -600,11 +601,12 @@ class Qwen3OmniMoeForConditionalGeneration(
update_dict["code_predictor_codes"] = code_predictor_codes
else:
# decode
if info_dict.get("num_processed_tokens", 0) < len(info_dict.get("thinker_input_ids", [])):
if not info_dict.get("decode_flag", False):
info_dict["num_processed_tokens"] = len(info_dict.get("thinker_input_ids", [])) + 1
update_dict["decode_flag"] = True
last_talker_hidden, text_step, update_dict = self.talker_preprocess_decode(
input_ids, input_embeds, **info_dict
input_ids, input_embeds, update_dict, **info_dict
)
update_dict["mtp_inputs"] = last_talker_hidden, text_step
@@ -874,8 +876,9 @@ class Qwen3OmniMoeForConditionalGeneration(
thinker_embed = thinker_embed[start_index : start_index + 1].to(device)
return self.talker.text_projection(thinker_embed).to(device)
def talker_preprocess_decode(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
update_dict: dict[str, dict] = {}
def talker_preprocess_decode(
self, input_ids: torch.Tensor, input_embeds: torch.Tensor, update_dict: dict, **info_dict: dict
):
last_talker_hidden = None
text_step = None
try:
vllm_omni/model_executor/models/qwen3_tts/qwen3_tts.py
+1 -1
View File
@@ -81,7 +81,7 @@ class Qwen3TTSModelForGeneration(nn.Module):
torch_dtype=torch.bfloat16,
**attn_kwargs,
)
self.task_type = model_path.split("-")[-1].strip("/")
self.task_type = model_path.split("-")[-1].split("/")[0]
# Mark that this model produces multimodal outputs
self.have_multimodal_outputs = True
vllm_omni/worker/base.py
+7 -1
View File
@@ -9,8 +9,10 @@ from vllm.logger import init_logger
from vllm.utils.mem_utils import format_gib, memory_profiling
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
from vllm_omni.entrypoints.utils import detect_pid_host
from vllm_omni.worker.gpu_memory_utils import (
get_process_gpu_memory,
is_process_scoped_memory_available,
)
logger = init_logger(__name__)
@@ -59,7 +61,11 @@ class OmniGPUWorkerBase(GPUWorker):
self.non_torch_memory = profile_result.non_torch_increase
self.peak_activation_memory = profile_result.torch_peak_increase
process_memory = get_process_gpu_memory(self.local_rank)
process_memory = (
get_process_gpu_memory(self.local_rank)
if is_process_scoped_memory_available() and detect_pid_host()
else None
)
if process_memory is not None:
# NVML available: use per-process memory