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

Merge branch 'main' into lance-model-support

Browse Source
This commit is contained in:
SYLAR authored and GitHub committed 2026-05-22 04:13:58 +08:00
2 parents 7b885f9d13 e949ccf087
commit f0cec3ecd4
224 files changed
+19202 -7026
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/scripts/hardware_ci/run-xpu-test.sh
+2 -1
View File
@@ -52,5 +52,6 @@ time timeout -k 30 30m docker run \
pip install tblib==3.1.0
cd /workspace/vllm-omni
pytest -v -s -m "core_model and xpu and B60"
pytest -v -s -m "advanced_model and xpu and B60" -k "not omni_expansion"
export VLLM_XPU_USE_SAMPLER_KERNEL=0 # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
pytest -v -s -m "advanced_model and xpu and B60"
'
.buildkite/test-nightly.yml
+1 -1
View File
@@ -959,7 +959,7 @@ steps:
- label: ":full_moon: Diffusion X2V · Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model"
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
.buildkite/test-weekly.yml
+62
View File
@@ -84,3 +84,65 @@ steps:
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips).
- group: "Reliability Test - Invalid parameters Test"
depends_on: upload-weekly-pipeline
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
steps:
- label: "Invalid parameters Test · H100"
timeout_in_minutes: 120
commands:
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
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
# Stable Audio Open, Qwen3-TTS speech/batch/voices paths marked ``L4``.
- label: "Invalid parameters Test · L4"
timeout_in_minutes: 120
commands:
- pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"
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
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
benchmarks/diffusion/README.md
+3 -2
View File
@@ -30,7 +30,7 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
**Notes**
- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
- By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
- If you run the server on another host or port, pass `--base-url` accordingly.
## 2. Supported Datasets
@@ -95,8 +95,9 @@ You can point to your own trace using `--dataset-path`.
### Basic flags
- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
- `--base-url`: Server address; `--endpoint` selects the path appended to this base URL.
- `--model`: The OpenAI-compatible `model` field.
- `--endpoint`: API endpoint path. Leading `/` is optional, e.g. `/v1/videos` or `v1/videos`.
- `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
- `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
- `--num-prompts`: Number of requests to send.
benchmarks/diffusion/backends.py
+123 -4
View File
@@ -10,6 +10,8 @@ from typing import Any
import aiohttp
from tqdm import tqdm
DEFAULT_EDITS_BOT_TASK = "think"
@dataclass
class RequestFuncInput:
@@ -27,6 +29,7 @@ class RequestFuncInput:
extra_body: dict[str, Any] = field(default_factory=dict)
image_paths: list[str] | None = None
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
@dataclass
@@ -53,6 +56,95 @@ def _encode_image_as_data_url(path: str) -> str:
return f"data:{mime};base64,{encoded}"
async def async_request_image_edits(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
enable_diffusion_pipeline_profiler: bool = False,
) -> RequestFuncOutput:
"""POST /v1/images/edits (multipart)."""
del enable_diffusion_pipeline_profiler
output = RequestFuncOutput()
output.start_time = time.perf_counter()
extra_body = dict(input.extra_body)
width = input.width or extra_body.get("width") or 1024
height = input.height or extra_body.get("height") or 1024
edits_url = input.api_url
form = aiohttp.FormData()
form.add_field("model", input.model)
form.add_field("prompt", input.prompt)
form.add_field("size", f"{width}x{height}")
form.add_field("response_format", "b64_json")
if input.num_inference_steps is not None:
form.add_field("num_inference_steps", str(input.num_inference_steps))
elif extra_body.get("num_inference_steps") is not None:
form.add_field("num_inference_steps", str(extra_body["num_inference_steps"]))
if input.seed is not None:
form.add_field("seed", str(input.seed))
elif extra_body.get("seed") is not None:
form.add_field("seed", str(extra_body["seed"]))
if extra_body.get("guidance_scale") is not None:
form.add_field("guidance_scale", str(extra_body["guidance_scale"]))
if extra_body.get("negative_prompt") is not None:
form.add_field("negative_prompt", str(extra_body["negative_prompt"]))
if extra_body.get("true_cfg_scale") is not None:
form.add_field("true_cfg_scale", str(extra_body["true_cfg_scale"]))
if extra_body.get("sys_type") is not None:
form.add_field("sys_type", str(extra_body["sys_type"]))
if extra_body.get("system_prompt") is not None:
form.add_field("system_prompt", str(extra_body["system_prompt"]))
bot_task = extra_body.get("bot_task")
if bot_task is None and input.default_bot_task is not None:
bot_task = input.default_bot_task
if bot_task is not None:
form.add_field("bot_task", str(bot_task))
assert input.image_paths is not None
for img_path in input.image_paths:
if not os.path.exists(img_path):
output.error = f"Image file not found: {img_path}"
output.success = False
if pbar:
pbar.update(1)
return output
with open(img_path, "rb") as img_f:
image_bytes = img_f.read()
form.add_field(
"image",
image_bytes,
filename=os.path.basename(img_path),
content_type=_guess_mime_type(img_path),
)
try:
async with session.post(edits_url, data=form) as response:
if response.status == 200:
resp_json = await response.json()
output.response_body = resp_json
output.success = True
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
except Exception as e:
output.error = str(e)
output.success = False
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
async def async_request_chat_completions(
input: RequestFuncInput,
session: aiohttp.ClientSession,
@@ -151,7 +243,7 @@ async def async_request_chat_completions(
return output
async def async_request_openai_images(
async def async_request_openai_image_generations(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
@@ -348,12 +440,39 @@ async def async_request_v1_videos(
return output
LEGACY_BACKEND_ENDPOINT_ALIASES = {
"vllm-omni": "/v1/chat/completions",
"openai": "/v1/images/generations",
}
def normalize_endpoint(value: str) -> str:
endpoint = str(value).strip()
if not endpoint:
raise ValueError("endpoint must not be empty.")
endpoint = LEGACY_BACKEND_ENDPOINT_ALIASES.get(
endpoint,
LEGACY_BACKEND_ENDPOINT_ALIASES.get(endpoint.lstrip("/"), endpoint),
)
if not endpoint.startswith("/"):
endpoint = f"/{endpoint}"
return endpoint
def endpoint_filename_token(value: str) -> str:
token = normalize_endpoint(value).lstrip("/")
for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
token = token.replace(bad, "_")
return token or "endpoint"
backends_function_mapping = {
"2i": {
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
"openai": (async_request_openai_images, "/v1/images/generations"),
"/v1/chat/completions": (async_request_chat_completions, "/v1/chat/completions"),
"/v1/images/generations": (async_request_openai_image_generations, "/v1/images/generations"),
"/v1/images/edits": (async_request_image_edits, "/v1/images/edits"),
},
"2v": {
"v1/videos": (async_request_v1_videos, "/v1/videos"),
"/v1/videos": (async_request_v1_videos, "/v1/videos"),
},
}
benchmarks/diffusion/diffusion_benchmark_serving.py
+81 -36
View File
@@ -6,31 +6,34 @@
Benchmark online serving for diffusion models (Image/Video Generation).
If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
Supports multiple backends:
- vllm-omni: Uses /v1/chat/completions endpoint (default)
- openai: Uses /v1/images/generations endpoint
- v1/videos: Use /v1/videos endpoint
Supports multiple endpoints:
- /v1/chat/completions: OpenAI chat-compatible image requests (e.g. t2i, Qwen i2i)
- /v1/images/edits: OpenAI image edit / IT2I (multipart; e.g. Hunyuan --bot-task think)
- /v1/images/generations: OpenAI image generation requests
- /v1/videos: Async video jobs
Legacy --backend vllm-omni and openai are aliases for chat/completions and images/generations.
Usage:
# Video (v1/videos backend)
# Video (/v1/videos endpoint)
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos --dataset vbench --task t2v --num-prompts 10 \
--endpoint /v1/videos --dataset vbench --task t2v --num-prompts 10 \
--height 480 --width 640 --fps 16 --num-frames 80
i2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos --dataset vbench --task i2v --num-prompts 10
--endpoint /v1/videos --dataset vbench --task i2v --num-prompts 10
# Image (vllm-omni backend)
# Image (/v1/chat/completions endpoint)
t2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
--endpoint /v1/chat/completions --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset random --task t2i --num-prompts 1 \
--endpoint /v1/chat/completions --dataset random --task t2i --num-prompts 1 \
--max-concurrency 1 --enable-negative-prompt \
--random-request-config '[
{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
@@ -39,20 +42,25 @@ Usage:
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
]'
i2i:
ti2i (Hunyuan / OpenAI image edit API):
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task i2i --num-prompts 10
--endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
--bot-task think
# Image (openai backend)
i2i (chat-based models such as Qwen-Image-Edit):
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
# Image (/v1/images/generations endpoint)
t2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend openai --dataset vbench --task t2i --num-prompts 10 \
--endpoint /v1/images/generations --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024 --port 3000
# Video (v1/videos)
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos --dataset random --task t2v --num-prompts 1 \
--endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \
--max-concurrency 1 --enable-negative-prompt \
--random-request-config '[
{"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
@@ -80,7 +88,12 @@ from typing import Any
import aiohttp
import numpy as np
import requests
from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
from backends import (
RequestFuncInput,
RequestFuncOutput,
backends_function_mapping,
normalize_endpoint,
)
from PIL import Image
from tqdm.asyncio import tqdm
@@ -127,7 +140,7 @@ class VBenchDataset(BaseDataset):
def _load_data(self) -> list[dict[str, Any]]:
if self.args.task == "t2v":
return self._load_t2v_prompts()
elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
elif self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
return self._load_i2v_data()
else:
return self._load_t2v_prompts()
@@ -519,7 +532,7 @@ class TraceDataset(BaseDataset):
single = row.get("image_path")
image_paths = [single] if single else None
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]:
raise ValueError(
f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
)
@@ -580,7 +593,7 @@ class RandomDataset(BaseDataset):
self._sampled_requests = None
# Random image generate
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i", "it2i"]:
self._random_image_path = self._generate_random_image_paths()
else:
self._random_image_path = None
@@ -883,13 +896,23 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None:
time.sleep(1)
def _default_endpoint_for_task(task: str) -> str:
if task in {"t2v", "i2v", "ti2v"}:
return "/v1/videos"
if task in {"i2i", "ti2i", "it2i"}:
return "/v1/images/edits"
if task == "t2i":
return "/v1/chat/completions"
raise ValueError(f"Unsupported task for endpoint resolution: {task}")
async def benchmark(args):
# Construct base_url if not provided
if args.base_url is None:
args.base_url = f"http://{args.host}:{args.port}"
VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"}
if args.task in VIDEO_TASKS:
task_type = "2v"
@@ -902,18 +925,23 @@ async def benchmark(args):
f"Valid image tasks: {sorted(IMAGE_TASKS)}"
)
valid_backends = sorted(backends_function_mapping[task_type].keys())
raw_endpoint = args.endpoint if args.endpoint is not None else args.backend
if raw_endpoint is None:
raw_endpoint = _default_endpoint_for_task(args.task)
args.endpoint = normalize_endpoint(raw_endpoint)
if args.backend not in valid_backends:
valid_endpoints = sorted(backends_function_mapping[task_type].keys())
if args.endpoint not in valid_endpoints:
logger.error(
f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
f"Valid backends for this task type: {valid_backends}\n"
f"Example usage: --task {args.task} --backend {valid_backends[0]}"
f"Invalid endpoint '{args.endpoint}' for task '{args.task}' (task type: '{task_type}').\n"
f"Valid endpoints for this task type: {valid_endpoints}\n"
f"Example usage: --task {args.task} --endpoint {valid_endpoints[0]}"
)
raise ValueError("Backend validation failed. See log above for valid options.")
raise ValueError("Endpoint validation failed. See log above for valid options.")
# Setup API URL and request function based on backend
request_func, api_url = backends_function_mapping[task_type][args.backend]
# Setup API URL and request function based on endpoint.
request_func, api_url = backends_function_mapping[task_type][args.endpoint]
api_url = f"{args.base_url}{api_url}"
if args.dataset == "vbench":
@@ -929,6 +957,10 @@ async def benchmark(args):
requests_list = dataset.get_requests()
print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
if args.endpoint == "/v1/images/edits":
for req in requests_list:
req.default_bot_task = args.bot_task
# Limit concurrency
if args.max_concurrency is not None:
semaphore = asyncio.Semaphore(args.max_concurrency)
@@ -976,15 +1008,17 @@ async def benchmark(args):
metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
# Add configuration info to metrics for JSON output
metrics["backend"] = args.backend
metrics["endpoint"] = args.endpoint
metrics["model"] = args.model
metrics["dataset"] = args.dataset
metrics["task"] = args.task
if args.endpoint == "/v1/images/edits":
metrics["bot_task"] = args.bot_task
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
# Section 1: Configuration
print("{:<40} {:<15}".format("Backend:", args.backend))
print("{:<40} {:<15}".format("Endpoint:", args.endpoint))
print("{:<40} {:<15}".format("Model:", args.model))
print("{:<40} {:<15}".format("Dataset:", args.dataset))
print("{:<40} {:<15}".format("Task:", args.task))
@@ -1047,12 +1081,17 @@ if __name__ == "__main__":
parser.add_argument("--host", type=str, default="localhost", help="Server host.")
parser.add_argument("--port", type=int, default=8091, help="Server port.")
parser.add_argument("--model", type=str, default="default", help="Model name.")
parser.add_argument(
"--endpoint",
type=str,
default=None,
help=("Endpoint path to target. Leading '/' is optional, e.g. /v1/videos or v1/videos."),
)
parser.add_argument(
"--backend",
type=str,
default="vllm-omni",
choices=["vllm-omni", "openai", "v1/videos"],
help="Backend to target the benchmark to.",
default=None,
help=argparse.SUPPRESS,
)
parser.add_argument(
"--dataset",
@@ -1065,7 +1104,7 @@ if __name__ == "__main__":
"--task",
type=str,
default="t2v",
choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "t2i"],
choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"],
help="Task type.",
)
parser.add_argument(
@@ -1172,9 +1211,15 @@ if __name__ == "__main__":
default=1,
help=(
"Number of synthetic input images to attach for image-conditioned tasks "
"(i2v, ti2v, ti2i, i2i) when using random dataset."
"(i2v, ti2v, ti2i, i2i, it2i) when using random dataset."
),
)
parser.add_argument(
"--bot-task",
type=str,
default="think",
help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."),
)
args = parser.parse_args()
asyncio.run(benchmark(args))
benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md
+3 -3
View File
@@ -58,7 +58,7 @@ Record these parameters when reporting performance results.
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni \
--endpoint /v1/chat/completions \
--dataset <DATASET_NAME> \
--task t2i \
--num-prompts <N> \
@@ -71,7 +71,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
| Parameter | Description |
| ---------------------- | --------------------------------- |
| `--backend` | Serving backend (use `vllm-omni`) |
| `--endpoint` | API endpoint (use `/v1/chat/completions`; leading `/` is optional) |
| `--dataset` | Dataset name (`random` or custom) |
| `--task` | Task type (e.g., `t2i`) |
| `--num-prompts` | Total number of requests |
@@ -129,7 +129,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni \
--endpoint /v1/chat/completions \
--dataset random \
--task t2i \
--num-prompts 1 \
benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md
+3 -3
View File
@@ -53,7 +53,7 @@ Record these parameters when reporting performance results.
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos \
--endpoint /v1/videos \
--dataset <DATASET_NAME> \
--task t2v \
--num-prompts <N> \
@@ -66,7 +66,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
| Parameter | Description |
| ---------------------- | --------------------------------- |
| `--backend` | Serving backend (use `v1/videos`) |
| `--endpoint` | API endpoint (use `/v1/videos`; leading `/` is optional) |
| `--dataset` | Dataset name (`random` or custom) |
| `--task` | Task type (e.g., `t2v`) |
| `--num-prompts` | Total number of requests |
@@ -120,7 +120,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos \
--endpoint /v1/videos \
--dataset random \
--task t2v \
--num-prompts 1 \
benchmarks/diffusion/quantization_quality.py
+31 -6
View File
@@ -34,6 +34,16 @@ Video example (text-to-video):
--height 720 --width 1280 \
--num-frames 81 --num-inference-steps 40 --seed 42
LTX-2 example (text-to-video; audio output is generated but not scored by LPIPS):
python benchmarks/diffusion/quantization_quality.py \
--model Lightricks/LTX-Video-2 \
--task t2v \
--quantization fp8 int8 \
--prompts \
"A serene lakeside sunrise with mist over the water" \
--height 704 --width 1216 \
--num-frames 121 --num-inference-steps 40 --seed 42
Multiple quantization methods:
python benchmarks/diffusion/quantization_quality.py \
--model Tongyi-MAI/Z-Image-Turbo \
@@ -204,10 +214,19 @@ def _generate_video(omni, args, prompt, seed):
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images
frames = first.images[0]
else:
raise ValueError("Could not extract video frames from output.")
# LTX-2 (and similar audio+video models) may surface a dict or (video, audio) tuple
if isinstance(frames, dict):
frames = frames.get("video") or frames.get("frames")
elif isinstance(frames, tuple) and len(frames) == 2:
frames = frames[0]
if frames is None:
raise ValueError("Could not extract video frames from output.")
if isinstance(frames, torch.Tensor):
video = frames.detach().cpu()
if video.dim() == 5:
@@ -225,9 +244,11 @@ def _generate_video(omni, args, prompt, seed):
return frames_array, elapsed, peak_mem
def _unload_omni(omni):
"""Delete Omni instance and free GPU memory."""
del omni
def _free_gpu_memory():
"""Force GC and release cached GPU memory.
Must be called AFTER the caller has dropped (i.e., via `del`)
every reference to the Omni instance"""
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
@@ -267,7 +288,9 @@ def run_benchmark(args):
bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
bl_mem = baseline_outputs[prompts[0]][2] # use first prompt's memory
_unload_omni(omni_bl)
omni_bl.shutdown()
del omni_bl
_free_gpu_memory()
# Save baseline outputs
bl_dir = output_dir / "baseline"
@@ -307,7 +330,9 @@ def run_benchmark(args):
qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
qt_mem = qt_outputs[prompts[0]][2]
_unload_omni(omni_qt)
omni_qt.shutdown()
del omni_qt
_free_gpu_memory()
# Save quantized outputs
qt_dir = output_dir / config_label.replace(" ", "_")
benchmarks/tts/README.md
+21 -4
View File
@@ -40,7 +40,7 @@ vllm bench serve --omni \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"task_type":"Base"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
--save-result --result-dir ./results
```
@@ -58,7 +58,7 @@ vllm bench serve --omni \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"voice":"Vivian","language":"English","task_type":"CustomVoice"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
--save-result --result-dir ./results
```
@@ -76,10 +76,28 @@ vllm bench serve --omni \
--num-prompts 20 --num-warmups 2 \
--extra-body '{"task_type":"VoiceDesign","language":"English"}' \
--max-concurrency 1 --request-rate inf \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun \
--save-result --result-dir ./results
```
#### Streaming continuity (`audio_underrun`)
`audio_underrun` (seconds) is the per-request worst-case buffer deficit
under a realtime-rate player simulation. The audio-speech backend records
chunk arrival times during the SSE stream and surfaces:
- `Mean / Median / P{50,99} AUDIO_UNDERRUN (s)` per the standard percentile
output - any positive value means at least one inter-chunk gap was longer
than the chunk's audio duration.
- `Streaming continuity OK rate` - fraction of requests whose worst-case
underrun stayed under the threshold (default 100 ms, the commonly cited
"audible gap" budget). Override the threshold with
`VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S=<float>`.
This captures the failure mode where `RTF_p50 < 1` (server keeps up in
aggregate) but per-stream chunk arrival is bursty enough that listeners
still hear gaps - common at high concurrency on streaming TTS pipelines.
#### Add WER / SIM / UTMOS to any of the above
Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
@@ -215,7 +233,6 @@ sentinel for regressions in this area.
benchmarks/tts/
├── README.md (this file)
├── bench_tts.py CLI — serve-mode benchmark driver
├── bench_voxcpm_offline.py CLI — offline VoxCPM benchmark (sync + streaming)
├── plot_results.py Generate per-task / per-concurrency curves
└── model_configs.yaml Model registry (supported tasks + extra body)
```
benchmarks/tts/bench_tts.py
+1 -1
View File
@@ -118,7 +118,7 @@ def build_bench_args(
"--num-warmups",
"2",
"--percentile-metrics",
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun",
]
if resolved_dataset_path:
benchmarks/tts/bench_voxcpm_offline.py
-922
View File
@@ -1,922 +0,0 @@
"""Offline VoxCPM benchmark for vLLM Omni.
Supports both:
- sync one-shot (Omni.generate)
- streaming (AsyncOmni.generate with async_chunk config)
- text-only synthesis
- voice cloning
- text/clone batch inputs from txt or jsonl
Usage::
# Sync (default voice)
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--text "Hello world" \\
--output-dir results/audio/
# Streaming (async_chunk)
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \\
--txt-prompts prompts.txt \\
--output-dir results/audio/
# Voice cloning batch via JSONL
python benchmarks/tts/bench_voxcpm_offline.py \\
--model /path/to/VoxCPM \\
--jsonl-prompts prompts.jsonl \\
--output-dir results/audio/
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import tempfile
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
def _find_repo_root(start: Path) -> Path:
"""Walk up from ``start`` until a repo marker is found.
Falls back to ``parents[2]`` for backwards compatibility if no marker hits
(which can only happen in unusual checkouts — the tree should always have
pyproject.toml + vllm_omni/ at the top level).
"""
for candidate in [start, *start.parents]:
if (candidate / "pyproject.toml").is_file() and (candidate / "vllm_omni").is_dir():
return candidate
return start.parents[2]
REPO_ROOT = _find_repo_root(Path(__file__).resolve())
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm.yaml"
logger = logging.getLogger(__name__)
@dataclass(frozen=True, slots=True)
class PromptSpec:
text: str
label: str
ref_audio: str | None = None
ref_text: str | None = None
def _require_soundfile():
try:
import soundfile as sf # type: ignore
except ModuleNotFoundError as exc:
raise RuntimeError(
"soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile"
) from exc
return sf
def _build_prompt(
args,
*,
text: str,
ref_audio: str | None = None,
ref_text: str | None = None,
global_request_id: str | None = None,
) -> dict[str, Any]:
additional_information: dict[str, list[Any]] = {
"text": [text],
"cfg_value": [args.cfg_value],
"inference_timesteps": [args.inference_timesteps],
"min_len": [args.min_len],
"max_new_tokens": [args.max_new_tokens],
}
if args.streaming_prefix_len is not None:
additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
if ref_audio:
additional_information["ref_audio"] = [ref_audio]
if ref_text:
additional_information["ref_text"] = [ref_text]
if global_request_id is not None:
additional_information["global_request_id"] = [global_request_id]
return {
"prompt_token_ids": [1],
"additional_information": additional_information,
}
def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
audio = mm.get("audio", mm.get("model_outputs"))
if audio is None:
raise ValueError("No audio output found in multimodal output.")
if isinstance(audio, list):
parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio]
audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
if not isinstance(audio, torch.Tensor):
audio = torch.as_tensor(audio)
return audio.float().cpu().reshape(-1)
def _extract_sample_rate(mm: dict[str, Any]) -> int:
sr_raw = mm.get("sr", 24000)
if isinstance(sr_raw, list) and sr_raw:
sr_raw = sr_raw[-1]
if hasattr(sr_raw, "item"):
return int(sr_raw.item())
return int(sr_raw)
def _emit_offline_metrics(
*,
request_id: str,
elapsed_s: float,
first_audio_elapsed: float | None,
audio_duration_s: float,
) -> None:
metrics = {
"request_id": request_id,
"ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None,
"audio_duration_s": round(audio_duration_s, 6),
"rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None,
}
print(f"[OfflineMetrics] {metrics}")
def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None:
sf = _require_soundfile()
if isinstance(audio_tensor, torch.Tensor):
audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy()
else:
audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy()
sf.write(
output_path,
audio_np,
sample_rate,
format="WAV",
subtype="PCM_16",
)
def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path:
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / f"output_{request_id}.wav"
_write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm))
return output_path
def _iter_request_multimodal_outputs(request_output: Any):
outputs = getattr(request_output, "outputs", None)
if outputs:
for output in outputs:
mm = getattr(output, "multimodal_output", None)
if isinstance(mm, dict):
yield mm
mm = getattr(request_output, "multimodal_output", None)
if isinstance(mm, dict):
yield mm
def _read_non_empty_lines(path: str) -> list[str]:
with open(path, encoding="utf-8") as f:
return [line.strip() for line in f if line.strip()]
def _load_prompt_specs(args) -> list[PromptSpec]:
specs: list[PromptSpec] = []
if args.txt_prompts is not None:
texts = _read_non_empty_lines(args.txt_prompts)
if not texts:
raise ValueError(f"No prompts found in {args.txt_prompts}")
for idx, text in enumerate(texts, start=1):
specs.append(
PromptSpec(
text=text,
label=f"item{idx:03d}",
ref_audio=args.ref_audio,
ref_text=args.ref_text,
)
)
return specs
if args.jsonl_prompts is not None:
with open(args.jsonl_prompts, encoding="utf-8") as f:
for line_no, raw_line in enumerate(f, start=1):
line = raw_line.strip()
if not line:
continue
try:
item = json.loads(line)
except json.JSONDecodeError as exc:
raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc
if not isinstance(item, dict):
raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object")
text = item.get("text")
if not isinstance(text, str) or not text.strip():
raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'")
ref_audio = item.get("ref_audio", args.ref_audio)
ref_text = item.get("ref_text", args.ref_text)
if (ref_audio is None) != (ref_text is None):
raise ValueError(
f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together"
)
specs.append(
PromptSpec(
text=text.strip(),
label=f"item{len(specs) + 1:03d}",
ref_audio=ref_audio,
ref_text=ref_text,
)
)
if not specs:
raise ValueError(f"No prompts found in {args.jsonl_prompts}")
return specs
specs.append(
PromptSpec(
text=args.text,
label="item001",
ref_audio=args.ref_audio,
ref_text=args.ref_text,
)
)
return specs
def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]:
return _build_prompt(
args,
text=spec.text,
ref_audio=spec.ref_audio,
ref_text=spec.ref_text,
global_request_id=global_request_id,
)
def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int:
return sum(1 for spec in prompt_specs if spec.ref_audio is not None)
def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]:
return prompt_specs[:1]
def _extract_stream_finished(stage_output: Any) -> bool:
request_output = getattr(stage_output, "request_output", None)
request_finished = getattr(request_output, "finished", None)
if request_finished is not None:
return bool(request_finished)
return bool(getattr(stage_output, "finished", False))
def _build_profiled_stage_config(
stage_configs_path: str,
profiler_dir: str,
) -> str:
stage_config_path = Path(stage_configs_path)
yaml_text = stage_config_path.read_text(encoding="utf-8")
injected_lines: list[str] = []
injected_count = 0
for line in yaml_text.splitlines():
injected_lines.append(line)
if line.strip() != "engine_args:":
continue
indent = line[: len(line) - len(line.lstrip())]
child_indent = indent + " "
grandchild_indent = child_indent + " "
injected_lines.extend(
[
f"{child_indent}profiler_config:",
f'{grandchild_indent}profiler: "torch"',
f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"',
f"{grandchild_indent}torch_profiler_with_stack: true",
]
)
injected_count += 1
if injected_count == 0:
raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}")
tmp = tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
delete=False,
suffix=".yaml",
prefix=f"{stage_config_path.stem}_profile_",
)
tmp.write("\n".join(injected_lines) + "\n")
tmp.close()
return tmp.name
def parse_args():
parser = FlexibleArgumentParser(
description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)"
)
parser.add_argument(
"--model",
type=str,
default=os.environ.get("VOXCPM_MODEL"),
help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.",
)
parser.add_argument(
"--text",
type=str,
default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.",
help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.",
)
parser.add_argument(
"--txt-prompts",
type=str,
default=None,
help="Path to a .txt file with one synthesis text per line.",
)
parser.add_argument(
"--jsonl-prompts",
type=str,
default=None,
help=(
"Path to a .jsonl file. Each line must contain at least {'text': ...}; "
"clone rows can also set ref_audio/ref_text, and ref_text must be the "
"real transcript of ref_audio."
),
)
parser.add_argument(
"--ref-audio",
type=str,
default=None,
help=(
"Optional reference audio path for voice cloning. With --txt-prompts, "
"the same reference is applied to every line."
),
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help=(
"Real transcript of the reference audio. Placeholder text or mismatched "
"text will usually produce noisy/electronic clone audio."
),
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=str(DEFAULT_STAGE_SYNC),
help="Stage config YAML path. Routing is selected only from this path.",
)
parser.add_argument(
"--cfg-value",
type=float,
default=2.0,
help="Classifier-free guidance value for VoxCPM.",
)
parser.add_argument(
"--inference-timesteps",
type=int,
default=10,
help="Number of inference timesteps.",
)
parser.add_argument(
"--min-len",
type=int,
default=2,
help="Minimum generated token length.",
)
parser.add_argument(
"--max-new-tokens",
type=int,
default=4096,
help="Maximum generated token length.",
)
parser.add_argument(
"--streaming-prefix-len",
type=int,
default=None,
help="VoxCPM streaming window (optional, streaming mode only).",
)
parser.add_argument(
"--output-dir",
type=str,
default=None,
help="Directory for output WAV files.",
)
parser.add_argument(
"--stage-init-timeout",
type=int,
default=600,
help="Stage initialization timeout in seconds.",
)
parser.add_argument(
"--log-stats",
dest="log_stats",
action="store_true",
help="Enable vLLM Omni stats logging.",
)
parser.add_argument(
"--no-log-stats",
dest="log_stats",
action="store_false",
help="Disable vLLM Omni stats logging.",
)
parser.set_defaults(log_stats=True)
parser.add_argument(
"--num-runs",
type=int,
default=1,
help="Number of full inference runs (same prompt each time). Default 1.",
)
parser.add_argument(
"--warmup-runs",
type=int,
default=0,
help=(
"Optional number of warmup passes before measured runs. Warmup uses only "
"the first prompt and does not save outputs."
),
)
parser.add_argument(
"--enable-profiler",
action="store_true",
help=(
"Enable torch profiler for the configured stages. A temporary profiled "
"stage config is generated automatically."
),
)
parser.add_argument(
"--profiler-dir",
type=str,
default=None,
help="Directory for profiler traces. Defaults to <output-dir>/profiler when profiling is enabled.",
)
parser.add_argument(
"--profiler-stages",
type=int,
nargs="*",
default=None,
help="Optional stage ids to profile. Defaults to all stages that have profiler_config.",
)
parser.add_argument(
"--profiler-wait-seconds",
type=float,
default=30.0,
help="Seconds to wait after stop_profile for trace files to flush.",
)
args = parser.parse_args()
if not args.model:
parser.error("--model is required unless $VOXCPM_MODEL is set")
if args.txt_prompts is not None and args.jsonl_prompts is not None:
parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive")
if (args.ref_audio is None) != (args.ref_text is None):
parser.error("--ref-audio and --ref-text must be provided together")
if args.num_runs < 1:
parser.error("--num-runs must be >= 1")
if args.warmup_runs < 0:
parser.error("--warmup-runs must be >= 0")
if args.output_dir is None:
args.output_dir = (
"output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio"
)
if args.enable_profiler and args.profiler_dir is None:
args.profiler_dir = str(Path(args.output_dir) / "profiler")
try:
args.prompt_specs = _load_prompt_specs(args)
except ValueError as exc:
parser.error(str(exc))
return args
def _is_streaming_stage_config(stage_configs_path: str) -> bool:
cfg_name = Path(stage_configs_path).name.lower()
return "async_chunk" in cfg_name
async def _collect_streaming_audio(
omni: AsyncOmni,
args: Any,
spec: PromptSpec,
request_id: str,
*,
phase_label: str,
prompt_index: int,
prompt_count: int,
print_prompt: bool = False,
) -> tuple[torch.Tensor, int, float, float | None]:
prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id)
delta_chunks: list[torch.Tensor] = []
sample_rate = 24000
chunk_i = 0
prev_total_samples = 0
t_start = time.perf_counter()
first_audio_elapsed: float | None = None
if print_prompt:
print(f"---prompt---:{prompt}")
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = getattr(stage_output, "multimodal_output", None)
if not isinstance(mm, dict):
ro = getattr(stage_output, "request_output", None)
if ro is None:
continue
mm = getattr(ro, "multimodal_output", None)
if not isinstance(mm, dict) and getattr(ro, "outputs", None):
seq = ro.outputs[0]
mm = getattr(seq, "multimodal_output", None)
if not isinstance(mm, dict):
continue
sample_rate = _extract_sample_rate(mm)
try:
w = _extract_audio_tensor(mm)
n = int(w.numel())
if n == 0:
continue
finished = _extract_stream_finished(stage_output)
if n > prev_total_samples:
delta = w.reshape(-1)[prev_total_samples:]
prev_total_samples = n
elif finished and n == prev_total_samples:
delta = w.reshape(-1)[:0]
else:
delta = w.reshape(-1)
prev_total_samples += int(delta.numel())
if int(delta.numel()) > 0:
delta_chunks.append(delta)
if first_audio_elapsed is None and int(delta.numel()) > 0:
first_audio_elapsed = time.perf_counter() - t_start
logger.info(
"%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s",
phase_label,
prompt_index + 1,
prompt_count,
chunk_i,
int(delta.numel()),
n,
finished,
)
chunk_i += 1
except ValueError:
if not _extract_stream_finished(stage_output):
logger.debug("skip non-audio partial output chunk=%d", chunk_i)
if not delta_chunks:
raise RuntimeError("No audio chunks received; check stage config and logs.")
audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0)
elapsed = time.perf_counter() - t_start
return audio_cat, sample_rate, elapsed, first_audio_elapsed
async def _abort_streaming_residual_work(
omni: AsyncOmni,
request_id: str,
*,
settle_seconds: float = 0.1,
) -> None:
"""Stop any late stage-0 work once the final audio has been collected."""
await omni.engine.abort_async([request_id])
if settle_seconds > 0:
await asyncio.sleep(settle_seconds)
async def _run_streaming_single(
omni: AsyncOmni,
args: Any,
spec: PromptSpec,
output_dir: Path,
request_id: str,
*,
run_index: int,
num_runs: int,
prompt_index: int,
prompt_count: int,
) -> Path:
audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio(
omni,
args,
spec,
request_id,
phase_label=f"run={run_index + 1}/{num_runs}",
prompt_index=prompt_index,
prompt_count=prompt_count,
print_prompt=(run_index == 0 and prompt_index == 0),
)
await _abort_streaming_residual_work(omni, request_id)
output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav"
_write_audio_tensor(output_path, audio_cat, sample_rate)
audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
print(
f"Saved (streaming) run {run_index + 1}/{num_runs}, "
f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})"
)
_emit_offline_metrics(
request_id=request_id,
elapsed_s=elapsed,
first_audio_elapsed=first_audio_elapsed,
audio_duration_s=audio_duration_s,
)
return output_path
async def _run_streaming_warmup(args, omni: AsyncOmni) -> None:
if args.warmup_runs == 0:
return
warmup_specs = _get_warmup_specs(args.prompt_specs)
print(
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
)
for warmup_index in range(args.warmup_runs):
t_warmup = time.perf_counter()
tasks = []
request_ids: list[str] = []
for prompt_index, spec in enumerate(warmup_specs):
request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
request_ids.append(request_id)
tasks.append(
_collect_streaming_audio(
omni,
args,
spec,
request_id,
phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}",
prompt_index=prompt_index,
prompt_count=len(warmup_specs),
)
)
results = await asyncio.gather(*tasks)
for request_id in request_ids:
await _abort_streaming_residual_work(omni, request_id)
total_samples = sum(int(audio.numel()) for audio, _, _, _ in results)
warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None]
ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else ""
print(
f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: "
f"{len(results)} prompt(s), {total_samples} sample(s) "
f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})"
)
async def _run_streaming(args) -> list[Path]:
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
omni = AsyncOmni(
model=args.model,
stage_configs_path=args.stage_configs_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
await _run_streaming_warmup(args, omni)
profiler_started = False
if args.enable_profiler:
profile_prefix = f"voxcpm_streaming_{int(time.time())}"
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}")
await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
profiler_started = True
t_total = time.perf_counter()
total_elapsed = 0.0
paths: list[Path] = []
prompt_specs: list[PromptSpec] = args.prompt_specs
try:
for run in range(args.num_runs):
for prompt_index, spec in enumerate(prompt_specs):
request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}"
paths.append(
await _run_streaming_single(
omni,
args,
spec,
output_dir,
request_id,
run_index=run,
num_runs=args.num_runs,
prompt_index=prompt_index,
prompt_count=len(prompt_specs),
)
)
total_elapsed = time.perf_counter() - t_total
finally:
if profiler_started:
print("Stopping profiler (streaming)...")
await omni.stop_profile(stages=args.profiler_stages)
if args.profiler_wait_seconds > 0:
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
await asyncio.sleep(args.profiler_wait_seconds)
print(
f"All streaming runs finished: {args.num_runs} run(s), "
f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total"
)
return paths
def _run_sync(args) -> list[Path]:
output_dir = Path(args.output_dir)
omni = Omni(
model=args.model,
stage_configs_path=args.stage_configs_path,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
def _run_sync_single(
spec: PromptSpec,
*,
request_prefix: str,
save_outputs: bool,
run_index: int | None = None,
) -> tuple[list[Path], int, float | None, float, float, str]:
global_request_id = f"{request_prefix}_{spec.label}"
prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id)
if save_outputs and run_index == 0 and spec.label == "item001":
print(f"---prompt---:{prompt}")
saved_paths: list[Path] = []
output_count = 0
first_audio_elapsed: float | None = None
total_audio_duration_s = 0.0
metrics_request_id = global_request_id
t_start = time.perf_counter()
for stage_outputs in omni.generate(prompt):
request_output = stage_outputs.request_output
if request_output is None:
continue
request_output_id = getattr(request_output, "request_id", None)
if isinstance(request_output_id, str) and request_output_id:
metrics_request_id = request_output_id
for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)):
output_count += 1
if first_audio_elapsed is None:
try:
audio_tensor = _extract_audio_tensor(mm)
if int(audio_tensor.numel()) > 0:
first_audio_elapsed = time.perf_counter() - t_start
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
except ValueError:
pass
else:
try:
audio_tensor = _extract_audio_tensor(mm)
total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm))
except ValueError:
pass
if not save_outputs:
continue
save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}"
saved_paths.append(_save_wav(mm, output_dir, save_stem))
if output_count == 0:
raise RuntimeError("No output from Omni.generate")
elapsed_s = time.perf_counter() - t_start
return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id
if args.warmup_runs:
warmup_specs = _get_warmup_specs(args.prompt_specs)
print(
f"Warmup: {args.warmup_runs} run(s) using the first prompt "
f"({len(warmup_specs)} prompt(s)); outputs will be discarded."
)
for warmup_index in range(args.warmup_runs):
t_warmup = time.perf_counter()
_, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single(
warmup_specs[0],
request_prefix=f"warmup_sync{warmup_index + 1}",
save_outputs=False,
)
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
print(
f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: "
f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})"
)
profiler_started = False
if args.enable_profiler:
profile_prefix = f"voxcpm_sync_{int(time.time())}"
stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured"
print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}")
omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages)
profiler_started = True
t_total = time.perf_counter()
total_elapsed = 0.0
saved_paths: list[Path] = []
prompt_specs: list[PromptSpec] = args.prompt_specs
try:
for run in range(args.num_runs):
t_run = time.perf_counter()
run_paths: list[Path] = []
for prompt_index, spec in enumerate(prompt_specs):
prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = (
_run_sync_single(
spec,
request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}",
save_outputs=True,
run_index=run,
)
)
run_paths.extend(prompt_paths)
ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else ""
rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else ""
print(
f"Saved (sync) run {run + 1}/{args.num_runs}, "
f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}"
)
_emit_offline_metrics(
request_id=metrics_request_id,
elapsed_s=elapsed_s,
first_audio_elapsed=first_audio_elapsed,
audio_duration_s=audio_duration_s,
)
saved_paths.extend(run_paths)
print(
f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)"
)
for path in run_paths:
print(f" {path}")
total_elapsed = time.perf_counter() - t_total
finally:
if profiler_started:
print("Stopping profiler (sync)...")
omni.stop_profile(stages=args.profiler_stages)
if args.profiler_wait_seconds > 0:
print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...")
time.sleep(args.profiler_wait_seconds)
print(
f"All sync runs finished: {args.num_runs} run(s), "
f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total"
)
return saved_paths
def main(args) -> int:
logging.basicConfig(level=logging.INFO)
profiled_stage_config_path: str | None = None
original_stage_config_path = args.stage_configs_path
if args.enable_profiler:
Path(args.profiler_dir).mkdir(parents=True, exist_ok=True)
profiled_stage_config_path = _build_profiled_stage_config(
args.stage_configs_path,
str(Path(args.profiler_dir).resolve()),
)
args.stage_configs_path = profiled_stage_config_path
is_streaming = _is_streaming_stage_config(args.stage_configs_path)
voice_clone_count = _count_voice_clone_prompts(args.prompt_specs)
print(f"Model: {args.model}")
print(f"Stage config: {original_stage_config_path}")
print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)")
print(f"Prompt count: {len(args.prompt_specs)}")
print("Batch mode: sequential (aligned with native VoxCPM)")
print(f"Warmup runs: {args.warmup_runs}")
print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}")
if args.enable_profiler:
print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})")
print(f"Profiled stage config: {args.stage_configs_path}")
if voice_clone_count:
print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.")
print(f"Num runs: {args.num_runs}")
try:
if is_streaming:
asyncio.run(_run_streaming(args))
else:
_run_sync(args)
finally:
if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path):
os.unlink(profiled_stage_config_path)
return 0
if __name__ == "__main__":
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
raise SystemExit(main(parse_args()))
docker/Dockerfile.xpu
+5
View File
@@ -133,6 +133,11 @@ ENTRYPOINT []
FROM ${VLLM_BASE} AS vllm-omni
RUN apt-get update && \
apt-get install -y espeak-ng ffmpeg git jq && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
WORKDIR /workspace/vllm-omni
COPY . .
docs/.nav.yml
+4 -1
View File
@@ -6,10 +6,11 @@ nav:
- getting_started/installation/*
- Serving:
- OpenAI-Compatible API:
- Diffusion Chat API: serving/diffusion_chat_api.md
- Chat Completions API: serving/chat_completions_api.md
- Image Generation: serving/image_generation_api.md
- Image Edit: serving/image_edit_api.md
- Text to Speech: serving/speech_api.md
- Videos: serving/videos_api.md
- Audio Generation: serving/audio_generate_api.md
- Streaming Video Input: serving/video_stream_api.md
- Examples:
@@ -79,6 +80,8 @@ nav:
- Quantized KV Cache: user_guide/quantization/quantized_kvcache.md
- FP8 W8A8: user_guide/quantization/fp8.md
- Int8 W8A8: user_guide/quantization/int8.md
- MXFP8 W8A8: user_guide/quantization/mxfp8.md
- MXFP4 W4A4: user_guide/quantization/mxfp4.md
- ModelOpt: user_guide/quantization/modelopt.md
- GGUF: user_guide/quantization/gguf.md
- AutoRound: user_guide/quantization/autoround.md
docs/api/README.md
-1
View File
@@ -121,7 +121,6 @@ Configuration classes.
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
- [vllm_omni.transformers_utils.configs.voxcpm.VoxCPMConfig][]
- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
## Workers
docs/community/governance.md
+26 -3
View File
@@ -39,11 +39,34 @@ Committers have write access and merge rights. They typically have deep expertis
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
## Diffusion Workload Division
This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
### Runtime-related
- Diffusion models: @RuixiangMa, @wtomin
### Optimization
- Parallel: @RuixiangMa, @wtomin
- Attention backend: @ZJY0516
- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
- Offload: @yuanheng-zhao
- Quantization: @lishunyang12, @david6666666
- Kernel fusion / communication-computation: @ZJY0516
### Docs & Test
- Diffusion docs: @RuixiangMa, @wtomin
## Meetings
docs/design/feature/diffusion_continuous_batching.md
+4 -1
View File
@@ -83,7 +83,10 @@ There are two important details:
Today that compatibility rule is still shape-sensitive. `height`, `width`,
`num_frames`, and CFG-related fields remain part of the key, so different
resolutions or incompatible guidance settings do **not** co-batch yet.
resolutions or incompatible guidance settings do **not** co-batch yet. The
key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests
targeting different adapters or scales run in separate batches and the
worker can activate exactly one adapter per step.
The current batching unit is one `OmniDiffusionRequest`. Requests with
multiple prompts do not participate in batching today.
docs/design/feature/disaggregated_inference.md
+2
View File
@@ -9,6 +9,7 @@ Backend-specific setup lives in separate docs:
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
## Overview
@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
## Core API
docs/design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
+211
View File
@@ -0,0 +1,211 @@
# YuanrongTransferEngineConnector
## When to Use
Use `YuanrongTransferEngineConnector` for high-performance peer-to-peer KV cache
transfer on Ascend NPU. It is different from `YuanrongConnector`:
- `YuanrongConnector` uses Yuanrong Datasystem as a distributed KV store and
requires Datasystem workers plus etcd.
- `YuanrongTransferEngineConnector` uses Yuanrong TransferEngine directly. It
registers local memory in each stage worker and lets the receiver pull data
from the sender with TransferEngine.
The implementation currently lives under
`vllm_omni/platforms/npu/omni_connectors/` because the supported backend is
Ascend NPU only. The connector is still registered with the generic
OmniConnector factory under the name `YuanrongTransferEngineConnector`.
For Ascend P2P transfer, the connector should use NPU memory as its transfer
pool. The Mooncake-style CPU pool is not supported for this connector.
## Prerequisites
Install Yuanrong Datasystem/TransferEngine Python bindings in the runtime
environment. The packaged install is:
```bash
pip install openyuanrong-datasystem
```
To build from source, follow the upstream project instructions:
- Source repository: [openeuler/yuanrong-datasystem](https://atomgit.com/openeuler/yuanrong-datasystem)
- Installation guide: [Yuanrong Datasystem Linux installation](https://pages.openeuler.openatom.cn/openyuanrong-datasystem/docs/zh-cn/latest/installation/installation_linux.html)
The Ascend driver, CANN runtime, HCCL/HCCP, and NPU device network must also be
configured on every machine participating in the transfer.
## Check NPU Device IPv4
Yuanrong TransferEngine with `protocol: "ascend"` uses the NPU device network,
not only the host management IP. Each participating NPU must have an IPv4 HCCN
device IP.
Check the device IP, link, and health state:
```bash
for i in {0..7}; do
echo "===== NPU $i ====="
hccn_tool -i "$i" -ip -g
hccn_tool -i "$i" -link -g
hccn_tool -i "$i" -net_health -g
done
```
Expected output for each NPU should include an IPv4 address and healthy link:
```text
ipaddr:10.86.1.80
netmask:255.255.255.0
link status: UP
net health status: Success
```
If the output says `Get ipconf failed, because no ip was preset there!`, the NPU
device IP is not configured and Ascend P2P will fail. If the address is IPv6,
configure IPv4 for this connector path or validate that the underlying Yuanrong
TransferEngine build supports the IPv6 HCCN path before using it.
## Check Device-to-Device Connectivity
After collecting the target-stage NPU IPv4 addresses, ping them from each
source-stage NPU:
```bash
# Example: AR uses NPU 0..3 and DiT uses NPU 4..7.
# Replace the destination IPs with the values reported by hccn_tool -i 4..7 -ip -g.
for src in 0 1 2 3; do
for dst_ip in 10.86.1.78 10.86.1.79 10.86.1.80 10.86.1.81; do
echo "===== NPU $src -> $dst_ip ====="
hccn_tool -i "$src" -ping -g address "$dst_ip" || true
done
done
```
The expected result is `0.00% packet loss` for every required source/target pair.
This verifies basic HCCN IPv4 connectivity. TransferEngine can still fail later
if RDMA/QP/HCCL setup is broken, but failed ping means the device network must be
fixed first.
## YAML Configuration
For Ascend P2P, configure the connector with `protocol: "ascend"` and an NPU
memory pool:
```yaml
runtime:
enabled: true
defaults:
window_size: -1
max_inflight: 1
connectors:
yuanrong_te_connector:
name: YuanrongTransferEngineConnector
extra:
host: "auto"
zmq_port: 50051
rpc_port: "auto"
protocol: "ascend"
device_name: "auto"
memory_pool_size: 1073741824
memory_pool_device: "npu"
edges:
- from: 0
to: 1
window_size: -1
```
Important fields:
| Parameter | Recommended Value | Notes |
|---|---|---|
| `host` | `"auto"` | Advertises a routable local host IP for ZMQ metadata exchange. |
| `zmq_port` | `50051` | Base ZMQ port. The runtime applies stage/rank offsets. |
| `rpc_port` | `"auto"` | Lets TransferEngine choose an available RPC port. |
| `protocol` | `"ascend"` | Uses the Ascend TransferEngine backend. |
| `device_name` | `"auto"` | Each TP worker resolves to its local logical NPU, for example `npu:0`, `npu:1`, etc. |
| `memory_pool_device` | `"npu"` | Required for the Ascend fast path. This registers NPU memory with TransferEngine. |
| `memory_pool_size` | `1073741824` | 1 GiB per worker is a practical starting point for KV-cache transfer validation. Increase if the transfer pool is exhausted. |
`memory_pool_device: "cpu"` is rejected for
`YuanrongTransferEngineConnector` with `protocol: "ascend"`. CPU pool registration
can lead to Ascend P2P failures such as HCCL RA init errors, QP timeouts, or
FFTS/SDMA runtime errors. `memory_pool_device: "cpu"` is common for
`MooncakeTransferEngineConnector`, but Yuanrong TE on Ascend only supports
`"npu"`.
## 4-Card AR-to-DiT TP Example
For AR 4-way TP sending KV cache to DiT 4-way TP, set the stage TP topology in
both stage `omni_kv_config` blocks:
```yaml
stage_args:
- stage_id: 0
runtime:
devices: "0,1,2,3"
engine_args:
tensor_parallel_size: 4
omni_kv_config:
need_send_cache: true
rank_mapping:
from_tp: 4
to_tp: 4
output_connectors:
to_stage_1: yuanrong_te_connector
- stage_id: 1
runtime:
devices: "4,5,6,7"
engine_args:
parallel_config:
tensor_parallel_size: 4
omni_kv_config:
need_recv_cache: true
rank_mapping:
from_tp: 4
to_tp: 4
input_connectors:
from_stage_0: yuanrong_te_connector
```
In this topology each DiT rank receives the corresponding AR rank shard:
```text
DiT rank 0 <- AR rank 0
DiT rank 1 <- AR rank 1
DiT rank 2 <- AR rank 2
DiT rank 3 <- AR rank 3
```
## Validate a Successful Run
Search the inference log for Yuanrong TE receive-side messages:
```bash
grep -E "\[YR GET\]|Successfully received KV cache|transfer_engine get failed" inference.log
```
A successful Ascend fast-path transfer contains lines like:
```text
[YR GET] ... fast_path, zero-copy
Successfully received KV cache ... across 1 key(s)
Applied CFG KV caches
[KV Reuse] ...
```
Sender-side `KV transfer OK` only means that the sender registered the payload
and metadata. The actual P2P data transfer happens on the receiver side during
`[YR GET]`.
## Troubleshooting
| Symptom | Likely Cause | Action |
|---|---|---|
| `IPv4 device IP not found` | NPU HCCN IPv4 is missing. | Run `hccn_tool -i <id> -ip -g` and configure IPv4 for every participating NPU. |
| `P2PCommInitRootInfo failed`, `ra init failed`, `RA_QP_STATUS_TIMEOUT` | Ascend P2P/RDMA/QP setup failed. | First verify device IPv4, link health, and device-to-device ping. Then test a minimal Yuanrong TE pair outside vLLM. |
| `fftsplus sdma error`, `context is abort` | Runtime SDMA/FFTS failure, often after an invalid or failed P2P transfer. | Use `memory_pool_device: "npu"` for Ascend TE, reduce `memory_pool_size` if OOM occurs, and isolate the failing NPU pair with a Yuanrong TE demo. |
| Sender logs `KV transfer OK` but receiver never logs `[YR GET]` | Sender `put()` succeeded, receiver P2P read failed or timed out. | Check receiver logs and search for `transfer_engine get failed`. |
| Pool allocation fails or OOM | Transfer pool is too small or too large for available NPU memory. | Start with 1 GiB (`1073741824`), reduce to 512 MiB for validation, or increase for larger KV payloads. |
docs/design/feature/pipeline_parallel.md
+356
View File
@@ -0,0 +1,356 @@
# Pipeline Parallel
This section describes how to add Pipeline Parallelism (PP) to a diffusion pipeline. We use the Wan2.2 text-to-video and
image-to-video pipelines as the reference implementations.
## Implementation Checklist
Adding Pipeline Parallel support requires:
1. ✅ **Inherit the mixin** - add `PipelineParallelMixin` before `CFGParallelMixin` in the pipeline class
2. ✅ **Make stage forward resumable** - support `intermediate_tensors` in the transformer forward path
3. ✅ **Return the right object type** - `IntermediateTensors` on non-last PP ranks, final model output on the last rank
4. ✅ **Use the standard denoising loop** - call `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
5. ✅ **Test parity** - compare PP results against the single-GPU baseline
---
## Table of Contents
- [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
- [Testing](#testing)
- [Troubleshooting](#troubleshooting)
- [Reference Implementations](#reference-implementations)
---
## Overview
### What is Pipeline Parallelism?
Pipeline Parallelism splits the denoising transformer into multiple sequential stages and places each stage on a
different rank. Instead of every rank holding the full DiT, each PP rank owns only a slice of the layers.
For each denoising step:
1. Rank 0 starts the forward pass using the current latents.
2. Each intermediate rank receives hidden states from the previous rank, runs its local layer slice, and forwards the
intermediate tensors downstream.
3. The last PP rank produces the final noise prediction.
4. The last PP rank applies the scheduler step and sends the updated latents back to rank 0 for the next timestep.
This reduces per-rank model memory and enables larger diffusion transformers to run across multiple GPUs. It can also be
combined with CFG-Parallel, where each PP pipeline carries one CFG branch.
### Architecture
vLLM-Omni provides `PipelineParallelMixin` to encapsulate the PP communication pattern for diffusion pipelines.
| Method | Purpose | Automatic Behavior |
|-----------------------------------|-------------------------------|---------------------------------------------------------------------------------|
| `diffuse()` | Denoising loop boundary | Wrapped by `PipelineParallelMixin` to flush pending async PP sends on exit |
| `predict_noise_maybe_with_cfg()` | Predict noise with PP support | Runs partial forwards on non-last PP ranks, combines with CFG logic when needed |
| `scheduler_step_maybe_with_cfg()` | Step scheduler with PP sync | Runs scheduler on the last PP rank, returns updated latents to rank 0 |
| `_sync_pp_send()` | Flush pending async sends | Waits for outstanding `isend` handles before later collectives or decode |
`PipelineParallelMixin` is intentionally a pipeline-level abstraction. Your model-specific `predict_noise()` still
defines how a local stage executes.
### How It Works
`predict_noise_maybe_with_cfg()` automatically switches between these modes:
- **PP disabled** (`pipeline_parallel_size == 1`):
- Falls back to `CFGParallelMixin.predict_noise_maybe_with_cfg()`
- **PP only** (`pipeline_parallel_size > 1`, `cfg_parallel_size == 1`):
- Rank 0 starts with the input latents
- Middle ranks receive `intermediate_tensors`, run their local layer range, and asynchronously send downstream
- The last rank returns the final noise prediction
- When CFG is enabled in this mode (sequential CFG), both the positive and negative branches run through the same
PP pipeline, doubling the communication volume per denoising step. Prefer PP + CFG-Parallel when
`cfg_parallel_size > 1` is available to avoid this overhead.
- **PP + CFG-Parallel** (`pipeline_parallel_size > 1`, `cfg_parallel_size > 1`):
- Each PP pipeline carries one CFG branch
- The last PP rank all-gathers across the CFG group
- CFG combination happens on every last PP rank in the CFG group, matching the non-PP CFG-parallel behavior
`scheduler_step_maybe_with_cfg()` keeps the denoising loop consistent:
- **PP disabled**:
- Falls back to `scheduler_step_maybe_with_cfg()`
- **PP enabled**:
- Only the last PP rank has `noise_pred` and runs the scheduler step
- The resulting latents are sent back to rank 0
- Rank 0 receives an `AsyncLatents` wrapper that resolves only when the tensor is next consumed
This asynchronous design avoids unnecessary blocking between denoising steps.
When a pipeline class is defined, `PipelineParallelMixin` wraps its `diffuse()` method and calls `_sync_pp_send()` in a
`finally` block after `diffuse()` returns or raises. This keeps model code free of explicit PP cleanup while still
ensuring that the final non-blocking PP send completes before decode or later collectives.
---
## Step-by-Step Implementation
### Step 1: Inherit `PipelineParallelMixin` and `CFGParallelMixin`
`PipelineParallelMixin` **requires** `CFGParallelMixin` and must appear before it in the class MRO. This is enforced at
class definition time via `__init_subclass__`: defining a pipeline that inherits `PipelineParallelMixin` without
`CFGParallelMixin`, or after `CFGParallelMixin`, raises a `TypeError` immediately on import.
`PipelineParallelMixin` delegates noise prediction, CFG combination, and scheduler stepping to `CFGParallelMixin`, which
supplies `predict_noise()`, `predict_noise_maybe_with_cfg()`, `scheduler_step_maybe_with_cfg()`, and
`combine_cfg_noise()`.
**Example:**
```python
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
import torch.nn as nn
class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
...
```
The order matters: `PipelineParallelMixin` must be listed before `CFGParallelMixin` so calls to
`predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()` resolve to the PP-aware wrappers, while their
`super()` calls delegate to the CFG implementation when PP is disabled or after the last PP stage.
### Step 2: Make model forward and `predict_noise()` PP-aware
The PP mixin injects `intermediate_tensors` into the normal `predict_noise()` call. Your model forward path must support
two inputs:
- Normal input from rank 0, typically passed as `hidden_states` or `x`
- `intermediate_tensors` from upstream PP ranks
The standard model forward pattern is:
1. If `intermediate_tensors` is present, read the local hidden state from it.
2. Run only this rank's layer slice.
3. Return `IntermediateTensors(...)` on non-last PP ranks.
4. Return the final model output on the last PP rank. `CFGParallelMixin.predict_noise()` already follows this contract
for common pipelines.
**Minimal example:**
```python
from vllm.sequence import IntermediateTensors
from vllm_omni.diffusion.distributed.parallel_state import get_pp_group
def forward(self, hidden_states=None, intermediate_tensors=None, **kwargs):
if intermediate_tensors is not None:
hidden_states = intermediate_tensors["hidden_states"]
for i in range(self.start_layer, self.end_layer):
hidden_states = self.layers[i](hidden_states)
pp_group = get_pp_group()
if not pp_group.is_last_rank:
return IntermediateTensors({"hidden_states": hidden_states})
return (hidden_states,)
```
### Step 3: Partition the transformer layers
The local module on each PP rank must expose only that rank's layer slice. In the reference tests and vLLM model
implementations, this is typically done with vLLM utilities such as `make_layers(...)`, which fill missing ranges with
`PPMissingLayer`. Besides partitioning the layers themselves, model authors should also wire
`make_empty_intermediate_tensors_factory(...)` for intermediate tensor allocation and `is_pp_missing_parameter(...)` for
PP-aware weight loading.
To prepare a transformer for PP, implement the following pieces in order.
#### 3.1 Split the transformer layers across PP ranks
Each PP rank should own only its local layer range, typically exposed as `[start_layer, end_layer)`. In practice this is
usually done with `make_layers(...)`, which constructs local layers and fills missing ranges with `PPMissingLayer`.
The goal is:
- every PP rank knows its `[start_layer, end_layer)` range
- non-local layers do not execute on this rank
- the forward can resume from incoming `intermediate_tensors`
By default, layers are distributed evenly across PP ranks using `get_pp_indices()`. When the number of layers is not
divisible by the PP size, remaining layers are assigned to middle partitions to balance compute and memory. You can
override this with the `VLLM_PP_LAYER_PARTITION` environment variable to specify an exact per-rank layer count:
```bash
# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
export VLLM_PP_LAYER_PARTITION=8,12,12,8
```
The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
the total number of transformer layers.
#### 3.2 Expose `make_empty_intermediate_tensors`
Your transformer module should expose `self.make_empty_intermediate_tensors`, typically created with
`make_empty_intermediate_tensors_factory(...)`.
This is important because PP ranks need a consistent way to allocate placeholder `IntermediateTensors` with the expected
keys and hidden dimension.
**Example:**
```python
from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
["hidden_states"],
inner_dim,
)
```
For Wan2.2, the intermediate payload between PP stages is the token sequence stored under `"hidden_states"`, so the
factory is created with that key and the transformer hidden size.
#### 3.3 Return `IntermediateTensors` on non-last PP ranks
Your model `forward()` or custom `predict_noise()` implementation should consume `intermediate_tensors` on non-first
ranks and return `IntermediateTensors(...)` on non-last ranks.
This lets each PP stage resume from the upstream hidden states and pass the local result to the next stage.
#### 3.4 Skip non-local weights during `load_weights()`
When a model is PP-partitioned, many parameters in the checkpoint belong to layers that do not exist on the current
rank. `load_weights()` must therefore skip parameters for missing PP stages using `is_pp_missing_parameter(...)`.
If this is not done, weight loading will either fail or incorrectly try to load tensors into `PPMissingLayer`
placeholders.
The Wan2.2 transformer is the best reference here: it uses `is_pp_missing_parameter(...)` before loading both remapped
and fused parameters.
If your model has multiple transformer variants, PP still works as long as each selected transformer obeys the same
contract.
### Step 4: Use the standard denoising contract
vLLM-Omni diffusion pipelines already route denoising through `diffuse()`, `predict_noise_maybe_with_cfg()`, and
`scheduler_step_maybe_with_cfg()`. `PipelineParallelMixin` overrides those standard helpers when
`pipeline_parallel_size > 1`, so model integrations should not add separate PP-specific helper names or manual post-loop
synchronization.
In PP mode:
- Non-last PP ranks return `None` from `predict_noise_maybe_with_cfg()` after sending `IntermediateTensors` downstream.
- The last PP rank returns the final prediction, and with CFG-Parallel every last PP rank in the CFG group receives the
combined result.
- Rank 0 receives `AsyncLatents` from `scheduler_step_maybe_with_cfg()`, which resolves only when the latents are
consumed.
- Pending non-blocking sends are flushed automatically when `diffuse()` exits.
---
## Testing
Use an offline inference script with `pipeline_parallel_size > 1`.
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2.mp4 \
--pipeline-parallel-size=2
```
For PP + CFG-Parallel together:
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2_cfg2.mp4 \
--pipeline-parallel-size=2 \
--cfg-parallel-size=2
```
**Verify:**
1. The run completes without hangs at the PP stage boundary
2. Output quality matches the non-PP baseline within normal numerical variance
3. Peak memory per GPU drops compared with the single-rank model
4. There are no outstanding communication errors before decode
---
## Troubleshooting
### Issue: `TypeError` on import — `CFGParallelMixin` missing
**Symptoms:** Importing a pipeline that inherits `PipelineParallelMixin` raises:
```
TypeError: YourPipeline inherits PipelineParallelMixin but not CFGParallelMixin.
```
or:
```
TypeError: YourPipeline must inherit PipelineParallelMixin before CFGParallelMixin ...
```
**Cause:** `PipelineParallelMixin` enforces via `__init_subclass__` that the subclass also inherits `CFGParallelMixin`
and lists `PipelineParallelMixin` first in the MRO.
**Solution:** Add `CFGParallelMixin` to the pipeline's base classes after `PipelineParallelMixin`:
```python
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.pipeline_parallel import PipelineParallelMixin
class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin):
...
```
### Issue: Non-last PP ranks crash when calling `predict_noise`
**Symptoms:** Shape or missing-input errors on ranks other than the first or last PP rank.
**Cause:** The model forward path assumes direct input tensors and ignores `intermediate_tensors`.
**Solution:** Update the transformer `forward()` or custom `predict_noise()` path to load hidden states from
`intermediate_tensors` when present.
### Issue: PP outputs differ from single-GPU baseline
**Symptoms:** The PP run finishes but produces numerically inconsistent outputs.
**Causes & Solutions:**
- **Local layer partitioning is wrong**
- Verify each rank runs only its `[start_layer, end_layer)` slice
- **Non-last ranks return a plain tensor instead of `IntermediateTensors`**
- Return `IntermediateTensors({...})` until the last PP stage
- **CFG branch wiring is incorrect**
- When CFG is enabled, confirm positive and negative kwargs are passed exactly as in the non-PP path
---
## Reference Implementations
Complete examples in the codebase:
| Component | Path | Notes |
|-------------------------|------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------|
| `PipelineParallelMixin` | `vllm_omni/diffusion/distributed/pipeline_parallel.py` | Core PP communication and scheduler helpers |
| `CFGParallelMixin` | `vllm_omni/diffusion/distributed/cfg_parallel.py` | Default `predict_noise()` tuple normalization and CFG helper fallback |
| Wan2.2 transformer | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference for layer partitioning, `IntermediateTensors`, `make_empty_intermediate_tensors`, and PP-aware weight loading |
| Wan2.2 T2V pipeline | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` | Reference PP + CFG integration for text-to-video |
| Wan2.2 I2V pipeline | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py` | Reference PP + CFG integration for image-to-video |
| PP tests | `tests/diffusion/distributed/test_pipeline_parallel.py` | Baseline parity and async communication tests |
docs/serving/diffusion_chat_api.md → docs/serving/chat_completions_api.md
+1 -1
View File
@@ -1,4 +1,4 @@
# Diffusion Chat Completions API
# Chat Completions API
vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
endpoint using diffusion models. This page explains how to pass generation
docs/serving/speech_api.md
+7
View File
@@ -5,6 +5,9 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano.
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
@@ -26,6 +29,10 @@ vllm serve fishaudio/s2-pro --omni --port 8091
# Voxtral TTS
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
# CosyVoice3 (voice cloning only — supply ref_audio + ref_text per request)
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 \
--omni --port 8091 --trust-remote-code
```
### Generate Speech
docs/serving/videos_api.md
+152
View File
@@ -0,0 +1,152 @@
# Videos API
vLLM-Omni provides an OpenAI-compatible video generation API for diffusion
video models. The API supports asynchronous video jobs through `/v1/videos` and
a synchronous benchmark-oriented endpoint through `/v1/videos/sync`.
Each server instance runs a single model specified at startup with
`vllm serve <model> --omni`.
## Quick Start
### Start the Server
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
```
### Create a Video Job
```bash
create_response=$(curl -s http://localhost:8091/v1/videos \
-F "prompt=A cinematic tracking shot of a mountain lake at sunrise" \
-F "width=1280" \
-F "height=720" \
-F "num_frames=80" \
-F "fps=16" \
-F "num_inference_steps=40")
video_id=$(echo "${create_response}" | jq -r '.id')
```
### Poll and Download
```bash
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
```
## API Reference
### Endpoints
| Endpoint | Method | Description |
|----------|--------|-------------|
| `/v1/videos` | `POST` | Create an asynchronous video generation job |
| `/v1/videos/sync` | `POST` | Generate a video synchronously and return raw video bytes |
| `/v1/videos/{video_id}` | `GET` | Retrieve job status and metadata |
| `/v1/videos` | `GET` | List stored video jobs |
| `/v1/videos/{video_id}/content` | `GET` | Download generated video content |
| `/v1/videos/{video_id}` | `DELETE` | Delete a video job and stored output |
### Request Parameters
`POST /v1/videos` and `POST /v1/videos/sync` accept `multipart/form-data`.
#### OpenAI-style fields
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `prompt` | string | **required** | Text prompt for video generation |
| `model` | string | server's model | Optional model name |
| `seconds` | string | null | Requested clip duration in seconds |
| `size` | string | null | Requested output size in `WIDTHxHEIGHT` format |
| `user` | string | null | Optional user identifier |
#### vLLM-Omni extension fields
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
| `width` | integer | model default | Output video width |
| `height` | integer | model default | Output video height |
| `num_frames` | integer | model default | Number of generated frames |
| `fps` | integer | model default | Output frames per second |
| `num_inference_steps` | integer | model default | Number of diffusion steps |
| `guidance_scale` | number | null | CFG guidance scale for the low-noise stage |
| `guidance_scale_2` | number | null | CFG guidance scale for the high-noise stage |
| `boundary_ratio` | number | null | Boundary split ratio for multi-stage denoising |
| `flow_shift` | number | null | Scheduler flow-shift value |
| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
| `seed` | integer | null | Random seed for reproducibility |
| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
| `frame_interpolation_scale` | number | null | RIFE inference scale |
| `frame_interpolation_model_path` | string | null | Local path or Hugging Face repo for the interpolation model |
| `lora` | string | null | JSON-encoded LoRA configuration object |
| `extra_params` | string | null | JSON-encoded object for additional model-specific parameters |
### Create Response
`POST /v1/videos` returns a job record:
```json
{
"id": "video-123",
"status": "queued",
"created_at": 1701234567
}
```
The final content is available from `/v1/videos/{video_id}/content` after the
job status becomes `completed`.
### Synchronous Response
`POST /v1/videos/sync` blocks until generation finishes and returns raw video
bytes. It is useful for benchmarks and simple scripts that do not need job
storage or polling.
## Examples
### Image-to-Video
```bash
curl -s http://localhost:8091/v1/videos \
-F "prompt=animate this image with subtle camera movement" \
-F "input_reference=@input.png" \
-F "width=1280" \
-F "height=720" \
-F "num_frames=80" \
-F "fps=16"
```
### Synchronous Generation
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=A small robot walking through a neon city" \
-F "width=854" \
-F "height=480" \
-F "num_frames=80" \
-F "fps=16" \
-o output.mp4
```
## Storage
Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are
stored:
```bash
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
```
## Model-Specific Examples
For complete text-to-video and image-to-video walkthroughs, see:
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
docs/user_guide/diffusion/frame_interpolation.md
-1
View File
@@ -28,7 +28,6 @@ Frame interpolation is currently supported for:
- `WanPipeline` (Wan2.2 text-to-video)
- `WanImageToVideoPipeline`
- `Wan22TI2VPipeline`
## Request Parameters
docs/user_guide/diffusion/parallelism/overview.md
+9 -8
View File
@@ -4,13 +4,14 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion
## Supported Methods
| Method | Description |
|--------|-------------|
| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
| Method | Description |
|----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------|
| **[Tensor Parallelism](tensor_parallel.md)** | Shards DiT weights across GPUs to reduce per-GPU memory |
| **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos |
| **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation |
| **[Pipeline Parallelism](pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory |
| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory |
| **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs |
| **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) |
See [Supported Models](../../diffusion_features.md#supported-models) for per-model compatibility.
docs/user_guide/diffusion/parallelism/pipeline_parallel.md
+205
View File
@@ -0,0 +1,205 @@
# Pipeline Parallelism Guide
## Table of Content
- [Overview](#overview)
- [Quick Start](#quick-start)
- [Example Script](#example-script)
- [Configuration Parameters](#configuration-parameters)
- [Best Practices](#best-practices)
- [Troubleshooting](#troubleshooting)
- [Summary](#summary)
---
## Overview
Pipeline Parallelism splits the denoising transformer block-wise into sequential stages across GPUs. Each rank owns only
part of the transformer, which reduces per-GPU model memory and enables larger diffusion models to run across multiple
devices.
It can also be combined with other distributed methods such as CFG-Parallel, Tensor Parallelism, and Sequence
Parallelism.
See supported models list in [Supported Models](../../diffusion_features.md#supported-models).
---
## Quick Start
### Basic Usage
Simplest working example:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(
model="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
parallel_config=DiffusionParallelConfig(
pipeline_parallel_size=2,
),
)
outputs = omni.generate(
{"prompt": "A cinematic drone shot over snowy mountains"},
OmniDiffusionSamplingParams(
num_inference_steps=40,
num_frames=81,
height=704,
width=1280,
),
)
```
---
## Example Script
### Offline Inference
Use python scripts under:
- `examples/offline_inference/text_to_video/text_to_video.py`
- `examples/offline_inference/image_to_video/image_to_video.py`
Text-to-video example:
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2.mp4 \
--pipeline-parallel-size=2
```
Pipeline Parallelism can also be combined with CFG-Parallel:
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model=Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--width=1280 \
--height=704 \
--guidance-scale=5.0 \
--prompt="Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage" \
--output=t2v_5B_pp2_cfg2.mp4 \
--pipeline-parallel-size=2 \
--cfg-parallel-size=2
```
### Online Serving
Enable Pipeline Parallelism in online serving:
```bash
# Default PP configuration
vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 --pipeline-parallel-size 2
# PP + CFG-Parallel
vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni --port 8091 \
--pipeline-parallel-size 2 \
--cfg-parallel-size 2
```
---
## Configuration Parameters
In `DiffusionParallelConfig`
| Parameter | Type | Default | Description |
|--------------------------|------|---------|------------------------------------------------------------------------------------------------------------------|
| `pipeline_parallel_size` | int | 1 | Number of pipeline-parallel stages. Set to a value greater than 1 to split the denoising transformer across GPUs |
> [!NOTE]
> Total GPU count is the product of all enabled distributed dimensions, for example
`pipeline_parallel_size * cfg_parallel_size * tensor_parallel_size * ulysses_degree * ring_degree`.
### Manual Layer Partitioning
By default, transformer layers are distributed evenly across PP ranks. You can override this with the
`VLLM_PP_LAYER_PARTITION` environment variable to assign a specific number of layers to each rank:
```bash
# Example: 40 layers across 4 PP ranks, assigning 8 / 12 / 12 / 8 layers
export VLLM_PP_LAYER_PARTITION=8,12,12,8
```
The value must be a comma-separated list of integers whose length equals `pipeline_parallel_size` and whose sum equals
the total number of transformer layers. This is useful when you want to balance memory or compute asymmetrically across
ranks.
---
## Best Practices
### When to Use
**Good for:**
- Large diffusion transformers that do not fit comfortably on one GPU
- Multi-GPU setups where reducing per-GPU model memory is more important than minimizing communication
- Combining with CFG-Parallel or other distributed methods on supported models
**Not for:**
- Single GPU setups
- Models that do not support Pipeline Parallelism
(check [supported models](../../diffusion_features.md#supported-models))
- Very small models where inter-stage communication overhead may outweigh the benefit
### Expected Behavior
Pipeline Parallelism primarily reduces per-GPU memory usage by splitting the transformer into stage-local blocks.
Depending on the model, topology, and resolution, it may also help execution fit into available hardware, but it is not
primarily a latency optimization.
---
## Troubleshooting
### Common Issue 1: No benefit from Pipeline Parallelism
**Symptoms**: PP is enabled, but latency does not improve or becomes slightly worse.
**Solutions**:
1. **Check your goal:**
```python
# PP is mainly for memory scaling, not guaranteed latency speedup
parallel_config = DiffusionParallelConfig(pipeline_parallel_size=2)
```
2. **Check model support:**
- Verify your model in [supported models](../../diffusion_features.md#supported-models)
- PP is currently validated only on selected pipelines
3. **Combine with other methods when appropriate:**
- PP can be combined with CFG-Parallel, Tensor Parallelism, or Sequence Parallelism on supported models
### Common Issue 2: PP pipeline fails at import
**Symptoms**: Importing a custom pipeline raises a `TypeError` about `CFGParallelMixin` or mixin order.
**Solutions**:
1. Inherit both mixins.
2. Put `PipelineParallelMixin` before `CFGParallelMixin` in the class base list.
3. Use `class YourPipeline(nn.Module, PipelineParallelMixin, CFGParallelMixin): ...` as the reference pattern.
---
## Summary
1. ✅ **Enable Pipeline Parallelism** - Set `pipeline_parallel_size > 1` in `DiffusionParallelConfig`
2. ✅ **Use Supported Models** - Verify your model supports PP in
[supported models](../../diffusion_features.md#supported-models)
3. ✅ **Combine When Needed** - PP can be combined with CFG-Parallel and other distributed methods on supported pipelines
4. ✅ **Scale for Memory** - Use PP primarily to reduce per-GPU model memory and fit larger transformers
docs/user_guide/diffusion/step_execution.md
+3
View File
@@ -60,6 +60,9 @@ batching only becomes relevant when `--max-num-seqs > 1`.
- `cache_backend` is not supported together with step execution.
- Unsupported pipelines fail early during model loading.
- Request-mode extras such as KV transfer are not wired into step mode yet.
- LoRA is supported in step mode, but each batch must use a single adapter:
requests with different `lora_request` or `lora_scale` are scheduled into
separate batches.
## When To Use It
docs/user_guide/diffusion_features.md
+50 -49
View File
@@ -35,14 +35,15 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).
| Method | Description | Best For |
|--------|-------------|----------|
| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
| Method | Description | Best For |
|------------------------------------------------------------------------|--------------------------------------------------------------------------|-------------------------------------------------------------------|
| **[Ulysses-SP](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via all-to-all communication | High-resolution images (>1536px) or long videos with 2-8 GPUs |
| **[Ring-Attention](diffusion/parallelism/sequence_parallel.md)** | Sequence parallelism via ring-based communication | Videos, very long sequences, memory-constrained, with 2-8 GPUs |
| **[CFG-Parallel](diffusion/parallelism/cfg_parallel.md)** | Splits CFG positive/negative branches across devices | Image editing with CFG guidance (true_cfg_scale > 1) on 2 GPUs |
| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
| **[Pipeline Parallelism](diffusion/parallelism/pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages | Large diffusion transformers that need lower per-GPU model memory |
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
#### Startup Optimization
@@ -105,32 +106,32 @@ The following tables show which models support each feature:
### ImageGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
> Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
@@ -138,15 +139,15 @@ The following tables show which models support each feature:
### VideoGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
**Frame Interpolation Support**
@@ -155,9 +156,9 @@ The following tables show which models support each feature:
### AudioGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-----------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
## Feature Compatibility
@@ -183,7 +184,7 @@ The following tables show which models support each feature:
| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ✅ | |
!!! info
@@ -192,7 +193,7 @@ The following tables show which models support each feature:
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
4. CPU Offloading (Layerwise) supports single-card for now.
5. Using FP8-Quant as an example of qunatization methods.
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT). LoRA is supported, but each scheduled batch must use a single adapter (requests with different `lora_request` or `lora_scale` are kept in separate batches).
## Multi-Thread Weight Loading
@@ -257,7 +258,7 @@ Measured on NVIDIA H800:
**Parallelism Methods:**
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, HSDP, and Expert Parallelism
- **[Parallelism Overview](diffusion/parallelism/overview.md)** - Tensor Parallelism, Sequence Parallelism, CFG Parallelism, Pipeline Parallelism, HSDP, and Expert Parallelism
**Memory Optimization:**
docs/user_guide/examples/offline_inference/text_to_speech.md
-60
View File
@@ -17,7 +17,6 @@ list of supported architectures across all modalities, see
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|---|---|---|---|---|---|---|
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz |
| VoxCPM | local model dir | split | ✓ | ✓ (`voxcpm_async_chunk.yaml`) | — | 24 kHz |
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz |
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
@@ -72,61 +71,6 @@ python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
---
## VoxCPM
Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
### Prerequisites
```bash
pip install voxcpm soundfile
# or use a local source tree:
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
```
If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
```bash
export VOXCPM_MODEL=/path/to/voxcpm-model
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
```
### Quick start
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model "$VOXCPM_MODEL" \
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
```
### Voice cloning
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model "$VOXCPM_MODEL" \
--text "This sentence is synthesized with a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in reference.wav."
```
### Streaming
Pass the async-chunk stage config:
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model "$VOXCPM_MODEL" \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
```
### Notes
- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming.
- Streaming is currently single-request oriented.
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
- For online serving, see [`examples/online_serving/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/voxcpm/README.md). For benchmark reporting, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
---
## CosyVoice3
2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
@@ -381,10 +325,6 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
``````py
--8<-- "examples/offline_inference/text_to_speech/qwen3_tts/end2end.py"
``````
??? abstract "voxcpm/end2end.py"
``````py
--8<-- "examples/offline_inference/text_to_speech/voxcpm/end2end.py"
``````
??? abstract "voxcpm2/end2end.py"
``````py
--8<-- "examples/offline_inference/text_to_speech/voxcpm2/end2end.py"
docs/user_guide/examples/online_serving/glm_image.md
+1 -1
View File
@@ -101,7 +101,7 @@ guides.
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Diffusion Chat API guide](../../../serving/diffusion_chat_api.md)).
[Chat Completions API guide](../../../serving/chat_completions_api.md)).
When using the dedicated [`/v1/images/generations`](../../../serving/image_generation_api.md)
or [`/v1/images/edits`](../../../serving/image_edit_api.md) endpoints, pass
the supported generation controls as top-level fields directly. For image
docs/user_guide/examples/online_serving/image_to_image.md
+2 -2
View File
@@ -208,7 +208,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
```
For details on how generation parameters are handled across different clients, see the
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
[Chat Completions API guide](../../../../serving/chat_completions_api.md).
### Layered Image Generation (Qwen-Image-Layered)
@@ -361,7 +361,7 @@ Provide multiple images in `content` (order matters):
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
[Chat Completions API guide](../../../../serving/chat_completions_api.md)).
When using the dedicated [`/v1/images/edits`](../../../../serving/image_edit_api.md)
endpoint, pass the supported generation controls as top-level form fields
directly. For image dimensions and count, use `size` and `n` rather than
docs/user_guide/examples/online_serving/text_to_image.md
+2 -2
View File
@@ -216,7 +216,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
```
For details on how generation parameters are handled across different clients, see the
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
[Chat Completions API guide](../../../../serving/chat_completions_api.md).
### Multimodal Input (Text + Structured Content)
@@ -237,7 +237,7 @@ Wrap generation parameters inside `extra_body` in the request JSON:
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
[Chat Completions API guide](../../../../serving/chat_completions_api.md)).
When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
endpoint, pass the supported generation controls as top-level JSON fields
directly. For image dimensions and count, use `size` and `n` rather than
docs/user_guide/examples/online_serving/text_to_speech.md
+62 -74
View File
@@ -17,14 +17,14 @@ For the full list of supported architectures across all modalities, see
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
|---|---|---|---|---|---|
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
| VoxCPM | local model dir | ✓ | ✓ (PCM stream) | — | — |
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3) instead.
For offline inference of any of these models, see the [offline TTS hub](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md).
## Common Quick Start
@@ -96,6 +96,54 @@ For full request-shape documentation (all parameters, response formats, error co
---
## CosyVoice3
2-stage TTS (`talker` + flow-matching `code2wav`) at 24 kHz. Voice cloning only — every request needs `ref_audio` + `ref_text`; there are no built-in voice presets.
### Prerequisites
```bash
huggingface-cli download FunAudioLLM/Fun-CosyVoice3-0.5B-2512
```
If your downloaded checkpoint lacks `config.json`, add one with `{"model_type": "cosyvoice3", "architectures": ["CosyVoice3Model"]}` (the loader reads `model_type` to select the class).
### Launch
```bash
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
# or:
./cosyvoice3/run_server.sh
```
Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
### CLI client
The client defaults to the official upstream zero-shot prompt, so it runs without extra flags:
```bash
cd examples/online_serving/text_to_speech/cosyvoice3
python speech_client.py --text "收到好友从远方寄来的生日礼物。"
```
Pass your own reference clip and transcript for a different voice:
```bash
python speech_client.py --text "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of the reference audio."
```
Stream PCM instead of WAV:
```bash
python speech_client.py --text "Hello world" --stream --output output.pcm
```
The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-text`, `--response-format`, `--stream`, `--output`.
### Notes
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
- For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
---
## Fish Speech S2 Pro
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
@@ -196,10 +244,10 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
./qwen3_tts/run_server.sh Base
```
### Choosing an executor backend (uniproc vs mp)
Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation.
### Executor backend
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor.
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
### Sending requests
```bash
@@ -281,71 +329,11 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
---
## VoxCPM
Split-stage TTS at 24 kHz.
### Prerequisites
```bash
pip install voxcpm
# or use a local source tree:
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
```
If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir:
```bash
export VOXCPM_MODEL=/path/to/voxcpm-model
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
```
### Launch
```bash
export VOXCPM_MODEL=/path/to/voxcpm-model
./voxcpm/run_server.sh # async-chunk streaming (default)
./voxcpm/run_server.sh sync # non-streaming
```
Or directly:
```bash
vllm serve "$VOXCPM_MODEL" \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--trust-remote-code --enforce-eager --omni --port 8091
```
### Sending requests
```bash
# Basic TTS
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This is a VoxCPM online text-to-speech example."
# Voice cloning
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This sentence is synthesized with a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in reference.wav."
# Streaming PCM
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This is a streaming VoxCPM request." \
--stream --output voxcpm_stream.pcm
```
### Notes
- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`.
- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`.
- For benchmark measurement, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
---
## VoxCPM2
@@ -407,6 +395,14 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
## Example materials
??? abstract "cosyvoice3/run_server.sh"
--8<-- "examples/online_serving/text_to_speech/cosyvoice3/run_server.sh"
??? abstract "cosyvoice3/speech_client.py"
--8<-- "examples/online_serving/text_to_speech/cosyvoice3/speech_client.py"
??? abstract "fish_speech/gradio_demo.py"
``````py
--8<-- "examples/online_serving/text_to_speech/fish_speech/gradio_demo.py"
@@ -463,14 +459,6 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
``````py
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/tts_common.py"
``````
??? abstract "voxcpm/openai_speech_client.py"
``````py
--8<-- "examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py"
``````
??? abstract "voxcpm/run_server.sh"
``````sh
--8<-- "examples/online_serving/text_to_speech/voxcpm/run_server.sh"
``````
??? abstract "voxcpm2/gradio_demo.py"
``````py
--8<-- "examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py"
docs/user_guide/quantization/autoround.md
+3 -1
View File
@@ -32,7 +32,9 @@ guide. AutoRound is Intel-supported.
|-------|------------|-------|--------|---------|
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
| Wan2.2 | Not listed | Diffusion transformer | W4A16 | Not validated |
| Wan2.2-I2V | `Intel/Wan2.2-I2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
| Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
docs/user_guide/quantization/mxfp4.md
+479
View File
@@ -0,0 +1,479 @@
# W4A4 MXFP4 Quantization
## Overview
W4A4 MXFP4 (Microscaling FP4) quantizes both weights and activations to FP4
(`float4_e2m1fn_x2`, packed 2 values per byte) using the OCP MX format: groups
of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale.
vLLM-Omni provides two quantization methods with different scale structures:
| Method | Scale structure | Mode | Use case |
|--------|----------------|------|----------|
| `mxfp4` | Single-scale (per-32 fine only) | Online only | Quick accuracy baseline; no checkpoint prep needed |
| `mxfp4_dualscale` | Dual-scale (fine per-32 + coarse per-512 + per-channel `mul_scale`) | Online + Offline | Production; better accuracy; offline recommended |
!!! tip "Recommended: `mxfp4_dualscale` offline"
For production deployments, use the `mxfp4_dualscale` offline mode with a
pre-quantized checkpoint produced by msModelSlim. Offline checkpoints load
calibrated `mul_scale` tensors from disk, providing measurably better accuracy
than any online method. The one-time preprocessing cost amortises across all
subsequent inference runs.
Use `mxfp4` online only for quick experimentation where preprocessing time
is not acceptable and accuracy loss is tolerable.
!!! warning "Online single-scale ≠ Offline dual-scale"
`mxfp4_dualscale` offline mode uses `NPUMxfp4DualScaleLinearMethod`:
fine scale (per-32 K), coarse scale (per-512 K), and per-input-channel
`mul_scale` from calibration — all loaded from the checkpoint.
`mxfp4_dualscale` online mode uses `NPUMxfp4DualScaleOnlineLinearMethod`:
dual-level scales computed on the fly from BF16 weights; no calibration
`mul_scale` is available. Loading an offline checkpoint with the online
method (or vice versa) will produce incorrect results or shape errors.
## Hardware Support
| Device | Support |
|--------|---------|
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
| AMD ROCm | ⭕ |
| Intel XPU | ⭕ |
| Ascend NPU (Atlas 950 A5) | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
## Model Type Support
### Diffusion Model (Wan2.2)
| Model | Online | Offline | Notes |
|-------|--------|---------|-------|
| Wan2.2-T2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade (`transformer` + `transformer_2`); both transformers quantized with the same config |
| Wan2.2-I2V-A14B | `mxfp4` / `mxfp4_dualscale` | `mxfp4_dualscale` | MoE cascade; same scheme as T2V-A14B |
| Wan2.2-TI2V-5B | ❌ | ❌ | Parameter count too small; W4A4 causes unacceptable accuracy loss |
The choice between `mxfp4` and `mxfp4_dualscale` in **online mode** is about
quantization quality, not model compatibility — both work on cascade (A14B) and
single-transformer models alike, the same as `mxfp8` online:
- `mxfp4`: single-scale, lower overhead, simpler compute, online only
- `mxfp4_dualscale`: dual-scale + optional BF16 fallback, better accuracy, online **and** offline
**Offline** checkpoints for A14B are always in `mxfp4_dualscale` format (produced
by the merge script); there is no offline `mxfp4` single-scale format.
!!! note "Per-layer BF16 fallback in offline cascade models"
The A14B offline checkpoint uses `quant_method: mxfp4_dualscale`. Most
linear layers are stored as W4A4 MXFP4 DualScale; precision-sensitive layers
retain their original BF16 weights and are listed in `ignored_layers` inside
each transformer's `config.json`. The two transformers may have different
`ignored_layers` sets — the pipeline reads each transformer's own `config.json`
and rebuilds the config locally when they differ, so routing is always
per-transformer-accurate.
!!! warning "TI2V-5B not supported"
Wan2.2-TI2V-5B is excluded from W4A4 quantization. Its smaller parameter
count makes it significantly more sensitive to 4-bit quantization noise,
resulting in unacceptable accuracy loss. Use [MXFP8](mxfp8.md) for TI2V-5B.
## Configuration
### `mxfp4` — Single-Scale Online Mode
Online mode requires no pre-processing. vLLM-Omni quantizes BF16 weights to
MXFP4 at load time using `npu_dynamic_mx_quant`. A single block scale
(`float8_e8m0fnu`, one per 32 K elements) is computed on the fly; no
calibration `mul_scale` is available. Applies equally to single-transformer
and cascade (A14B) models — both transformers in a cascade receive the same
quantization config automatically.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(model="<your-model>", quantization="mxfp4")
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
```bash
# Single-transformer or cascade model — same command
python text_to_video.py --model <your-model> --quantization mxfp4
# Online serving
vllm serve <your-model> --omni --quantization mxfp4
```
### `mxfp4_dualscale` — DualScale Online Mode
Online DualScale mode computes both fine and coarse scales on the fly from BF16
weights using `npu_dynamic_dual_level_mx_quant`. Applies equally to
single-transformer and cascade (A14B) models. Compared to `mxfp4` online,
DualScale provides better quantization accuracy at higher compute cost.
The default configuration keeps the leading 5 transformer blocks in BF16
(`num_bf16_fallback_layers=5`). Accuracy evaluation on Wan2.2-A14B shows this
is sufficient to meet quality requirements and is the recommended setting.
```python
omni = Omni(model="<your-model>", quantization="mxfp4_dualscale")
```
```bash
python text_to_video.py --model <your-model> --quantization mxfp4_dualscale
```
If accuracy debugging identifies additional precision-sensitive layers, they
can be pinned to BF16 via the Python API:
```python
omni = Omni(
model="<your-model>",
quantization_config={
"method": "mxfp4_dualscale",
"ignored_layers": ["blocks.10.attn1.to_q"], # explicit per-layer override
},
)
```
BF16 fallback routing in online mode applies two rules in priority order:
1. **`ignored_layers`** (explicit per-layer override): any layer whose prefix
matches is kept in BF16 regardless of block index.
2. **`num_bf16_fallback_layers`** (coarse leading-block rule): the first N
transformer blocks (`blocks.0` … `blocks.N-1`) fall back to BF16. Defaults
to `5` (recommended). Layers outside `blocks.N.*`
(e.g. `condition_embedder`) are always quantized.
### `mxfp4_dualscale` — DualScale Offline Mode (Recommended)
Offline mode loads a pre-quantized DualScale checkpoint from msModelSlim. A
preprocessing step converts the raw quantized output to the diffusers format
expected by vLLM-Omni and injects the quantization config into each
`transformer/config.json` so that vLLM-Omni auto-detects the offline path
without a `--quantization` flag.
BF16 fallback layers may be interleaved anywhere in the transformer — they are
not restricted to leading blocks. The merge script detects them from
`quant_model_description.json` and writes their prefixes into `ignored_layers`
inside `config.json`. At runtime, each layer's prefix is matched against
`ignored_layers` to decide BF16 vs. MXFP4 DualScale.
#### Checkpoint tensor layout
Each quantized linear layer stores four tensors:
| Tensor | Shape | dtype | Description |
|--------|-------|-------|-------------|
| `weight` | `(N, K)` | float8_e4m3fn | FP4 packed (2 values per byte) |
| `weight_scale` | `(N, K//32)` | uint8 | Fine block scale (`float8_e8m0fnu` bit pattern) |
| `weight_dual_scale` | `(N, K//512, 1)` | float32 | Coarse block scale |
| `mul_scale` | `(K,)` | float32 | Per-input-channel smooth pre-scale (from calibration) |
BF16 fallback layers have no quantization tensors; only the original `weight`
(and optional `bias`) are present, loaded directly from the base checkpoint.
#### Step 1 — Quantize with msModelSlim
```bash
msmodelslim quant \
--model_path /path/to/Wan2.2-T2V-A14B-Diffusers \
--save_path /path/to/wan2_2_t2v_quantized_raw \
--device npu \
--model_type Wan2_2 \
--config_path /path/to/wan2_2_w4a4_mxfp4_dualscale.yaml \
--trust_remote_code True
```
After this step, `--save_path` contains raw quantized safetensors files,
scale files, and a metadata JSON (`quant_model_description*.json`).
For cascade MoE models (T2V-A14B, I2V-A14B), msModelSlim outputs two
subdirectories: `high_noise_model/` (transformer) and `low_noise_model/`
(transformer_2).
#### Step 2 — Preprocess with merge_mxfp4_dualscale_checkpoint.py
The script (`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`):
1. Copies the original diffusers model to `--output-path` (VAE, text encoder,
scheduler, etc. are preserved).
2. Remaps tensor names from msModelSlim convention to diffusers convention and
strips `.linear.` / `.div.` wrappers added by the quantization tool.
3. Overlays MXFP4 tensors (weight, fine/coarse scales, `mul_scale`) onto the
BF16 base checkpoint. Non-quantized layers keep their original BF16 weights.
4. Detects all linear layers that remain in BF16 and writes their prefixes into
`ignored_layers` in `config.json`.
5. Injects `quantization_config` so vLLM-Omni auto-detects offline MXFP4
DualScale.
For cascade MoE models, steps 2–5 run separately for each transformer.
```bash
python vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py \
--model-type Wan2.2-T2V-A14B \
--original-model /path/to/Wan2.2-T2V-A14B-Diffusers \
--quant-path /path/to/wan2_2_t2v_quantized_raw \
--output-path /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
```
| Argument | Description |
|----------|-------------|
| `--model-type` | Model variant: `Wan2.2-T2V-A14B` or `Wan2.2-I2V-A14B` |
| `--original-model` | Root directory of the original BF16 diffusers model |
| `--quant-path` | Root directory of the msModelSlim quantized output |
| `--output-path` | Output directory for the merged model (created by the script) |
The script outputs a complete diffusers model directory at `--output-path`,
with each transformer subfolder containing:
- `diffusion_pytorch_model.safetensors` — MXFP4 weights + scale tensors, with BF16 fallback layers from the base checkpoint
- `config.json` — original transformer config with `quantization_config` injected
- `quant_model_description.json` — quantization metadata (reference only)
The `quantization_config` injected into `config.json` for each transformer:
```json
{
"quant_method": "mxfp4_dualscale",
"is_checkpoint_serialized": true,
"ignored_layers": [
"blocks.0.attn1.to_qkv",
"blocks.0.attn1.to_out",
"proj_out"
]
}
```
`ignored_layers` lists every linear layer that retains its original BF16 weight,
using vllm-omni model parameter names (QKV-fused, FFN underscored, `to_out`
unindexed). The exact entries are determined by the quantization tool (msModelSlim)
and may differ between `transformer` and `transformer_2` in a cascade model.
#### Step 3 — Serve
```bash
python text_to_video.py --model /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale
# Online serving
vllm serve /path/to/Wan2.2-T2V-A14B-MXFP4-DualScale --omni
```
```python
omni = Omni(model="/path/to/Wan2.2-T2V-A14B-MXFP4-DualScale")
```
!!! note
No `--quantization` flag is needed for offline mode. The preprocessing
script injects `quantization_config` into each `transformer/config.json`,
which vLLM-Omni reads automatically to activate the correct offline path.
## Parameters
### `mxfp4` (single-scale, online only)
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `method` | str | — | `"mxfp4"` |
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16 |
### `mxfp4_dualscale` (dual-scale, online + offline)
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `method` | str | — | `"mxfp4_dualscale"` |
| `is_checkpoint_serialized` | bool | `False` | `True` for offline DualScale checkpoints; auto-set from `config.json` when using the preprocessing script |
| `ignored_layers` | list[str] | `[]` | Layer prefixes to keep in BF16. **Works in both modes**: offline — populated by the merge script for interleaved sensitive layers; online — user-supplied for explicit per-layer precision override |
| `num_bf16_fallback_layers` | int | `5` | **Online mode only**: leading N transformer blocks (`blocks.0` … `blocks.N-1`) kept in BF16. Applied after `ignored_layers`; ignored in offline mode. Default of `5` is the evaluated recommended value for Wan2.2-A14B |
#### BF16 fallback priority (online mode)
```
for each linear layer:
if prefix in ignored_layers → BF16 (explicit override, highest priority)
elif block_idx < num_bf16_fallback_layers → BF16 (coarse leading-block rule)
else → MXFP4 DualScale online
```
Layers outside `blocks.N.*` (e.g. `condition_embedder.*`) are always quantized
unless they appear in `ignored_layers`.
## Validation and Notes
1. **Online single-scale (`mxfp4`)** quantizes BF16 weights at load time using
`npu_dynamic_mx_quant` (single-scale). No calibration `mul_scale` is
available — all output partitions receive an identity pre-scale. No offline
checkpoint format exists for this method.
2. **Online dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=False`)**
quantizes BF16 weights using `npu_dynamic_dual_level_mx_quant` (fine + coarse
scales computed on the fly). No calibration `mul_scale`; leading blocks or
explicit `ignored_layers` stay in BF16 for accuracy.
3. **Offline dual-scale (`mxfp4_dualscale`, `is_checkpoint_serialized=True`)** —
**recommended for production** — loads four tensors per quantized layer: FP4
weight, fine scale (`uint8` reinterpreted as `float8_e8m0fnu`), coarse scale
(`float32`), and per-input-channel `mul_scale` (`float32`). BF16 fallback
layers have no quantization tensors and are routed via `ignored_layers`.
4. **Scale dtype**: fine scales are stored as `uint8` in safetensors (same bit
layout as `float8_e8m0fnu`) and reinterpreted at load time without a lossy
float32 round-trip.
5. **Cascade model config propagation**: in a cascade model (transformer +
transformer_2), vLLM-Omni reads each transformer's own `config.json` and
rebuilds the quant config locally when `ignored_layers` differs between
transformers, ensuring per-layer routing is accurate for each. The first
transformer's config is propagated to `od_config` so the second transformer
can reuse it as a starting point.
6. **Self-attention QKV fusion**: Q, K, V projection weights are fused into a
single `QKVParallelLinear` layer at runtime. `ignored_layers` entries use the
fused name (`attn1.to_qkv`), written automatically by the merge script.
7. W4A4 carries higher quantization noise than W8A8 (16 vs 256 levels). The
DualScale offline method mitigates this with calibrated `mul_scale` smooth
quantization. Use `ignored_layers` and `num_bf16_fallback_layers` to trade
off compression vs. accuracy for precision-sensitive layers.
## Adapting MXFP4 for a New Model
This section is aimed at developers who want to add MXFP4 support to a model
other than Wan2.2. The three integration points are: (1) discovering the correct
runtime layer names, (2) wiring `ignored_layers` into the model, and (3) writing
a merge script for offline checkpoints.
### Step 1 — Discover runtime layer names
`ignored_layers` entries must match the **runtime parameter names** used inside
vllm-omni, which may differ from the names stored in the diffusers checkpoint.
The canonical source of truth is the model's own `named_parameters()`.
```python
from vllm_omni import Omni
# Load the model without quantization to inspect parameter names.
omni = Omni(model="/path/to/your-model") # no --quantization flag
for name, _ in omni.pipeline.transformer.named_parameters():
if "weight" in name and "scale" not in name:
print(name)
```
Compare the printed names against the diffusers checkpoint keys
(`safetensors.safe_open` or `torch.load`) to identify any renames your model
applies. Common patterns that differ in Wan2.2 (and may appear in other
models):
| Diffusers checkpoint name | vllm-omni runtime name | Reason |
|---------------------------|------------------------|--------|
| `attn1.to_q`, `attn1.to_k`, `attn1.to_v` | `attn1.to_qkv` | Self-attention Q/K/V fused into `QKVParallelLinear` |
| `ffn.net.0.proj` | `ffn.net_0.proj` | Dots in sub-module names replaced with underscores |
| `ffn.net.2` | `ffn.net_2` | Same underscore rule |
| `to_out.0` | `to_out` | Sequential index stripped |
If your model has different fusion patterns, inspect `packed_modules_mapping`
on the model class — this dict records how checkpoint keys are mapped to
fused runtime parameters.
!!! warning "Partial QKV fallback is not allowed"
If your model fuses Q, K, V into a single layer, `ignored_layers` must
include **all three or none**. A partial fallback (e.g. `to_q` in BF16 but
`to_k`, `to_v` quantized) cannot be expressed at runtime because they share
one `QKVParallelLinear`. The merge script enforces this and raises an error
if only some of the trio appear as non-quantized.
### Step 2 — Add ignored_layers to the model
#### Online mode
Pass `ignored_layers` directly in the quantization config using the **runtime
names** discovered in Step 1. No code changes to the model are required.
```python
omni = Omni(
model="/path/to/your-model",
quantization={
"method": "mxfp4_dualscale",
"ignored_layers": [
"blocks.0.attn1.to_qkv", # runtime name, not diffusers name
"blocks.0.attn1.to_out",
"blocks.0.ffn.net_0.proj",
],
},
)
```
```bash
# CLI does not support list-typed ignored_layers directly.
# Use the Python API or set ignored_layers in config.json (offline).
python your_script.py --model /path/to/your-model --quantization mxfp4_dualscale
```
The `num_bf16_fallback_layers` coarse rule is an alternative to listing layers
individually: set it to N to keep all linear layers in blocks 0 … N-1 in BF16.
The right value depends on the model's sensitivity; evaluate on a validation
set and pick the smallest N that meets your accuracy target.
#### Offline mode
For offline checkpoints, `ignored_layers` is written into each transformer's
`config.json` by the merge script (see Step 3). No manual editing is needed if
the merge script is correct. The injected block:
```json
{
"quant_method": "mxfp4_dualscale",
"is_checkpoint_serialized": true,
"ignored_layers": [
"blocks.0.attn1.to_qkv",
"blocks.0.attn1.to_out"
]
}
```
To add a layer manually (e.g. to pin an additional layer to BF16 without
re-running the merge script), edit `config.json` inside the transformer
subfolder. Use runtime names, not diffusers checkpoint names.
### Step 3 — Write a merge script for offline mode
The merge script for a new model mirrors
`vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py`. The four
things it must do:
1. **Remap tensor names** from the quantization tool convention to diffusers
convention (strip wrappers like `.linear.`, `.div.`; fix any prefix
differences).
2. **Collect ignored_layers**: after loading, enumerate all `*.weight` keys that
have no corresponding `*.weight_scale` (i.e. layers the tool left in BF16).
Convert diffusers names to vllm-omni runtime names (fuse QKV, rename FFN
sub-modules, etc.). Write the result to `config.json`.
3. **Inject `quantization_config`** into `config.json`:
```python
config["quantization_config"] = {
"quant_method": "mxfp4_dualscale",
"is_checkpoint_serialized": True,
"ignored_layers": ignored_layers, # runtime names
}
```
4. **Save** the merged safetensors and the updated `config.json`.
The key helper to implement is the diffusers-to-runtime name translator
(equivalent to `_diffusers_to_vllm_ignored` in the Wan2.2 merge script).
For each non-quantized diffusers weight key, apply your model's specific
renaming rules and collect the results.
!!! tip "Validate before serving"
After producing the offline checkpoint, load it without a `--quantization`
flag and verify that vLLM-Omni auto-detects the correct method. Check that
the layer count reported in the startup log matches expectations: quantized
layer count + `ignored_layers` count should equal total linear layer count.
Any mismatch indicates a name-mapping bug in the merge script.
docs/user_guide/quantization/online.md
+26 -11
View File
@@ -9,20 +9,24 @@ separate quantized checkpoint.
This mode is different from pre-quantized checkpoint formats such as GGUF,
AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
prepared before serving and are documented in their method-specific guides.
For MXFP8 and MXFP4, use this page for load-time quantization from BF16
checkpoints, and use the method-specific pages for offline checkpoints produced
by msModelSlim and the merge tools.
## Hardware Support
| Device | FP8 W8A8 | Int8 W8A8 |
|--------|----------|-----------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ |
| AMD ROCm | ⭕ | ⭕ |
| Intel XPU | ⭕ | ⭕ |
| Ascend NPU | ❌ | ✅ |
| Device | FP8 W8A8 | Int8 W8A8 | MXFP8 W8A8 | MXFP4 W4A4 |
|--------|----------|-----------|------------|------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ⭕ | ⭕ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ |
| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available.
guide. FP8 on Ampere may use a weight-only path where available. MXFP8 and
MXFP4 are documented for the Ascend NPU path.
## Model Type Support
@@ -32,6 +36,8 @@ guide. FP8 on Ampere may use a weight-only path where available.
|--------|-------|----------------|--------|
| FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
| Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Ascend NPU only |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
@@ -53,6 +59,9 @@ from vllm_omni import Omni
omni_fp8 = Omni(model="<your-model>", quantization="fp8")
omni_int8 = Omni(model="<your-model>", quantization="int8")
omni_mxfp8 = Omni(model="<your-model>", quantization="mxfp8")
omni_mxfp4 = Omni(model="<your-model>", quantization="mxfp4")
omni_mxfp4_dualscale = Omni(model="<your-model>", quantization="mxfp4_dualscale")
```
CLI:
@@ -60,6 +69,9 @@ CLI:
```bash
vllm serve <your-model> --omni --quantization fp8
vllm serve <your-model> --omni --quantization int8
vllm serve <your-model> --omni --quantization mxfp8
vllm serve <your-model> --omni --quantization mxfp4
vllm serve <your-model> --omni --quantization mxfp4_dualscale
```
Per-component routing:
@@ -77,10 +89,11 @@ config = build_quant_config({
| Parameter | Methods | Description |
|-----------|---------|-------------|
| `method` | FP8, Int8 | Quantization method, usually `"fp8"` or `"int8"` |
| `ignored_layers` | FP8, Int8 | Layer name patterns to keep in BF16/FP16 |
| `method` | FP8, Int8, MXFP8, MXFP4 | Quantization method: `"fp8"`, `"int8"`, `"mxfp8"`, `"mxfp4"`, or `"mxfp4_dualscale"` |
| `ignored_layers` | FP8, Int8, MXFP8, MXFP4 | Layer name patterns to keep in BF16/FP16 |
| `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
| `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
| `num_bf16_fallback_layers` | MXFP4 DualScale | Leading transformer blocks to keep in BF16 for online `mxfp4_dualscale`; defaults to `5` |
## Validation and Notes
@@ -91,3 +104,5 @@ config = build_quant_config({
model as supported.
4. If a model already ships quantized weights, use the matching pre-quantized
method guide instead of online quantization.
5. For Ascend MXFP4 deployments, prefer offline `mxfp4_dualscale` checkpoints
when production quality is more important than avoiding preprocessing.
docs/user_guide/quantization/overview.md
+109 -11
View File
@@ -9,20 +9,20 @@ type has a different quantization scope.
| Mode | Guide | Description | Methods |
|------|-------|-------------|---------|
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8 |
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 |
| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8 |
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
## Hardware Support
| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | GGUF | AutoRound | msModelSlim |
|--------|----------|-----------|----------|------------|------|-----------|-------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ |
| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim |
|--------|----------|-----------|----------|------------|------------|------|-----------|-------------|
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
guide. FP8 on Ampere may use a weight-only path where available.
@@ -42,6 +42,7 @@ otherwise.
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
@@ -58,6 +59,7 @@ in BF16 unless the model guide explicitly adds support.
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
@@ -73,6 +75,7 @@ attached to the intended stage rather than applied globally.
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
@@ -100,7 +103,7 @@ config = build_quant_config({
| Component | Default quantized? | Notes |
|-----------|--------------------|-------|
| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, GGUF, AutoRound, and msModelSlim |
| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim |
| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
| VAE | No | Keep BF16; storage-only paths are method-specific |
| Scheduler/tokenizer | No | Loaded from the base model repository |
@@ -138,3 +141,98 @@ build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
build_quant_config(None)
```
## Output Similarity Comparison Tool
Use `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
to compare a reference diffusion run with a quantized candidate run using the
same prompt, seed, resolution, scheduler settings, and inference steps. The
tool compares final decoded images or video frames, and also reports generation
latency and worker-reported peak memory when available.
This is useful when validating whether online quantization, an offline
pre-quantized checkpoint, or a new `ignored_layers` choice keeps generation
quality close to the BF16 reference.
### Online Quantization Example
```bash
python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
--task t2i \
--model Qwen/Qwen-Image \
--candidate-quantization fp8 \
--ignored-layers img_mlp \
--prompt "a cup of coffee on the table" \
--height 512 --width 512 \
--num-inference-steps 20 \
--seed 142 \
--output-json /tmp/qwen_image_fp8_similarity/result.json \
--save-output-dir /tmp/qwen_image_fp8_similarity/images \
--enforce-eager
```
### Offline Checkpoint Example
Use `--candidate-model` when the candidate is already quantized or lives at a
different model path:
```bash
python -m vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity \
--task t2i \
--reference-model Qwen/Qwen-Image \
--candidate-model /path/to/qwen-image-fp8-checkpoint \
--prompt "a cup of coffee on the table" \
--height 512 --width 512 \
--num-inference-steps 20 \
--seed 142 \
--output-json /tmp/qwen_image_fp8_checkpoint_similarity/result.json
```
If the checkpoint does not include a loadable quantization config, pass one
explicitly:
```bash
--candidate-quantization-config-json '{"method":"fp8"}'
```
### Output Metrics
The output JSON includes `output_metrics`, `reference_generation`, and
`candidate_generation`.
| Metric | Direction | Meaning |
|--------|-----------|---------|
| `cosine_similarity` | Higher is better | Vector direction similarity between output pixels or frames. Useful as a broad sanity check. |
| `mae` | Lower is better | Mean absolute pixel or frame error. For decoded outputs, values are in uint8 pixel units. |
| `mse` / `rmse` | Lower is better | Squared error and its square root. These penalize localized large differences more than `mae`. |
| `max_abs` | Lower is better | Worst single-element absolute error. Treat it as an outlier/debug signal, not as a release gate. |
| `l2` / `relative_l2` | Lower is better | Absolute and reference-normalized L2 distance. `relative_l2` is easier to compare across resolutions. |
| `psnr_db` | Higher is better | Pixel-space signal-to-noise ratio in dB for uint8 images or frames. |
| `avg_generation_time_s` | Lower is better | Average wall-clock generation time across measured runs. |
| `max_peak_memory_mb` | Lower is better | Maximum worker-reported peak device memory across measured runs, when the worker reports it. |
Recommended starting thresholds for same-seed diffusion comparisons:
| Metric | Smoke threshold | Stricter target | Notes |
|--------|-----------------|-----------------|-------|
| `psnr_db` | `>= 20.0` | `>= 25.0` | Good for quick image or frame regression checks. |
| `mae` | `<= 12.0` | `<= 6.0` | Interpreted in decoded uint8 pixel units. |
| `cosine_similarity` | `>= 0.98` | `>= 0.995` | Less sensitive to global scale than L2-style metrics. |
| `relative_l2` | `<= 0.20` | `<= 0.08` | Useful when comparing across prompts or resolutions. |
These thresholds are heuristics. Tune them by model family, task, resolution,
quantization method, and deployment tolerance. For release gating, pair the
numeric report with visual inspection of saved reference and candidate outputs.
The tool intentionally reports separate quality, latency, and memory metrics
instead of a single consolidated similarity score. A single score can hide
important tradeoffs, for example a candidate with good PSNR but a meaningful
memory regression, or a candidate with low average error but localized visual
artifacts. If you need a project-specific pass/fail gate, define it as an
explicit policy over the individual metrics.
Pixel-level metrics do not measure semantic consistency. For higher-cost
evaluation, you can complement this report with a vision-language judge that
describes the reference and candidate outputs and compares those descriptions.
Keep that semantic check separate from this lightweight tool so users can
choose whether the additional model cost and latency are appropriate.
examples/offline_inference/hunyuan_image3/end2end.py
+24
View File
@@ -77,6 +77,24 @@ def parse_args():
parser.add_argument("--log-stats", action="store_true", default=False)
parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
parser.add_argument(
"--diffusion-kv-cache-dtype",
type=str,
default=None,
help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.",
)
parser.add_argument(
"--diffusion-kv-cache-skip-steps",
type=str,
default=None,
help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.",
)
parser.add_argument(
"--diffusion-kv-cache-skip-layers",
type=str,
default=None,
help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.",
)
parser.add_argument(
"--additional-config",
type=str,
@@ -140,6 +158,9 @@ def main():
"init_timeout": args.init_timeout,
"enforce_eager": args.enforce_eager,
"mode": _MODALITY_MODE[args.modality],
"diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype,
"diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps,
"diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers,
}
if additional_config is not None:
@@ -237,6 +258,9 @@ def main():
print(f" Inference steps: {args.steps}")
print(f" Guidance scale: {args.guidance_scale}")
print(f" Seed: {args.seed}")
print(f" diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}")
print(f" diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}")
print(f" diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}")
if args.modality == "text2img":
print(f" Output size: {args.width}x{args.height}")
if args.image_path:
examples/offline_inference/image_to_video/image_to_video.py
+43 -33
View File
@@ -142,6 +142,42 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--enforce-eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
default=24000,
help="Sample rate for audio output when saved (default: 24000).",
)
parser.add_argument(
"--cache-backend",
type=str,
default=None,
choices=["cache_dit", "tea_cache"],
help=(
"Cache backend to use for acceleration. "
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
"Default: None (no cache acceleration)."
),
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
)
# Distributed and parallel execution
parser.add_argument(
"--ulysses-degree",
type=int,
@@ -173,28 +209,6 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
)
parser.add_argument(
"--enforce-eager",
action="store_true",
help="Disable torch.compile and force eager execution.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
default=24000,
help="Sample rate for audio output when saved (default: 24000).",
)
parser.add_argument(
"--cache-backend",
type=str,
default=None,
choices=["cache_dit", "tea_cache"],
help=(
"Cache backend to use for acceleration. "
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
"Default: None (no cache acceleration)."
),
)
parser.add_argument(
"--use-hsdp",
action="store_true",
@@ -219,16 +233,10 @@ def parse_args() -> argparse.Namespace:
),
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "int8", "gguf"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
"--pipeline-parallel-size",
type=int,
default=1,
help="Number of pipeline parallel stages.",
)
parser.add_argument(
"--profiler-config",
@@ -326,6 +334,7 @@ def main():
use_hsdp=args.use_hsdp,
hsdp_shard_size=args.hsdp_shard_size,
hsdp_replicate_size=args.hsdp_replicate_size,
pipeline_parallel_size=args.pipeline_parallel_size,
)
omni_kwargs = dict(
model=args.model,
@@ -366,7 +375,8 @@ def main():
print(f" diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}")
print(
f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size},"
f" pipeline_parallel_size={args.pipeline_parallel_size}"
)
print(f" Video size: {args.width}x{args.height}")
print(f"{'=' * 60}\n")
examples/offline_inference/text_to_speech/README.md
+2 -53
View File
@@ -14,13 +14,12 @@ list of supported architectures across all modalities, see
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|---|---|---|---|---|---|---|
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 22.05 kHz |
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
| VoxCPM | `openbmb/VoxCPM-0.5B` | split | ✓ | ✓ | — | 24 kHz |
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
@@ -41,7 +40,7 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
## CosyVoice3
2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz.
### Prerequisites
```bash
@@ -308,56 +307,6 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
---
## VoxCPM
Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
### Prerequisites
```bash
pip install voxcpm soundfile
# or use a local source tree:
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
```
### Quick start
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model openbmb/VoxCPM-0.5B \
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
```
### Voice cloning
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model openbmb/VoxCPM-0.5B \
--text "This sentence is synthesized with a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in reference.wav."
```
### Streaming
Pass `--streaming` together with the legacy async-chunk stage config:
```bash
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
--model openbmb/VoxCPM-0.5B \
--streaming \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
```
### Persistent HF config (optional)
The engine auto-patches a missing `model_type` in the checkpoint's `config.json` in-memory at startup (see `vllm_omni/engine/arg_utils.py`), so no manual setup is required.
If you'd rather maintain a persistent patched config dir (e.g. to share across runs or to add fields the auto-patcher doesn't set), export `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` to its location — both the bundled deploy and async-chunk yamls read it via `${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}`.
### Notes
- Non-streaming auto-loads `vllm_omni/deploy/voxcpm.yaml` from the model directory name / HF `model_type`. Streaming uses the legacy `voxcpm_async_chunk.yaml` passed via `--stage-configs-path`.
- Streaming is currently single-request oriented.
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
- For online serving, see the [VoxCPM section in the online hub](../../online_serving/text_to_speech/README.md#voxcpm). For benchmark reporting, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).
---
## VoxCPM2
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
examples/offline_inference/text_to_speech/voxcpm/end2end.py
-220
View File
@@ -1,220 +0,0 @@
"""Minimal offline VoxCPM example for vLLM Omni."""
from __future__ import annotations
import asyncio
import time
from pathlib import Path
from typing import Any
import soundfile as sf
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni import AsyncOmni, Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
def _build_prompt(args) -> dict[str, Any]:
additional_information: dict[str, list[Any]] = {
"text": [args.text],
"cfg_value": [args.cfg_value],
"inference_timesteps": [args.inference_timesteps],
"min_len": [args.min_len],
"max_new_tokens": [args.max_new_tokens],
}
if args.streaming_prefix_len is not None:
additional_information["streaming_prefix_len"] = [args.streaming_prefix_len]
if args.ref_audio is not None:
additional_information["ref_audio"] = [args.ref_audio]
if args.ref_text is not None:
additional_information["ref_text"] = [args.ref_text]
return {
"prompt_token_ids": [1],
"additional_information": additional_information,
}
def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor:
audio = mm.get("audio", mm.get("model_outputs"))
if audio is None:
raise ValueError("No audio output found in multimodal output.")
if isinstance(audio, list):
parts = [torch.as_tensor(item).float().cpu().reshape(-1) for item in audio]
audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0)
if not isinstance(audio, torch.Tensor):
audio = torch.as_tensor(audio)
return audio.float().cpu().reshape(-1)
def _extract_sample_rate(mm: dict[str, Any]) -> int:
sr_raw = mm.get("sr", 24000)
if isinstance(sr_raw, list) and sr_raw:
sr_raw = sr_raw[-1]
if hasattr(sr_raw, "item"):
return int(sr_raw.item())
return int(sr_raw)
def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path:
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / f"output_{request_id}.wav"
sf.write(
output_path,
audio.float().cpu().clamp(-1.0, 1.0).numpy(),
sample_rate,
format="WAV",
subtype="PCM_16",
)
return output_path
async def _run_streaming(args) -> Path:
prompt = _build_prompt(args)
output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio_streaming")
request_id = "streaming_example"
sample_rate = 24000
buffered_samples = 0
chunks: list[torch.Tensor] = []
started = time.perf_counter()
omni = AsyncOmni(
model=args.model,
stage_configs_path=args.stage_configs_path,
deploy_config=args.deploy_config,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
try:
async for stage_output in omni.generate(prompt, request_id=request_id):
mm = getattr(stage_output, "multimodal_output", None)
if not isinstance(mm, dict):
request_output = getattr(stage_output, "request_output", None)
if request_output is None:
continue
mm = getattr(request_output, "multimodal_output", None)
if not isinstance(mm, dict) and getattr(request_output, "outputs", None):
mm = getattr(request_output.outputs[0], "multimodal_output", None)
if not isinstance(mm, dict):
continue
audio = _extract_audio_tensor(mm)
if audio.numel() == 0:
continue
sample_rate = _extract_sample_rate(mm)
if audio.numel() > buffered_samples:
delta = audio[buffered_samples:]
buffered_samples = int(audio.numel())
else:
delta = audio
buffered_samples += int(delta.numel())
if delta.numel() > 0:
chunks.append(delta)
if not chunks:
raise RuntimeError("No streaming audio chunks received from VoxCPM.")
output_audio = torch.cat(chunks, dim=0)
output_path = _save_audio(output_audio, sample_rate, output_dir, request_id)
print(f"Saved streaming audio to: {output_path} ({time.perf_counter() - started:.2f}s)")
return output_path
finally:
omni.shutdown()
def _run_sync(args) -> Path:
prompt = _build_prompt(args)
output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio")
request_id = "sync_example"
started = time.perf_counter()
last_mm: dict[str, Any] | None = None
omni = Omni(
model=args.model,
stage_configs_path=args.stage_configs_path,
deploy_config=args.deploy_config,
log_stats=args.log_stats,
stage_init_timeout=args.stage_init_timeout,
)
for stage_outputs in omni.generate(prompt):
request_output = getattr(stage_outputs, "request_output", None)
if request_output is None:
continue
outputs = getattr(request_output, "outputs", None)
if outputs:
for output in outputs:
mm = getattr(output, "multimodal_output", None)
if isinstance(mm, dict):
last_mm = mm
mm = getattr(request_output, "multimodal_output", None)
if isinstance(mm, dict):
last_mm = mm
if last_mm is None:
raise RuntimeError("No audio output received from VoxCPM.")
output_path = _save_audio(
_extract_audio_tensor(last_mm),
_extract_sample_rate(last_mm),
output_dir,
request_id,
)
print(f"Saved audio to: {output_path} ({time.perf_counter() - started:.2f}s)")
return output_path
def parse_args():
parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.")
parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help=("Override the deploy config path."),
)
parser.add_argument(
"--stage-configs-path",
type=str,
default=None,
help=(
"Legacy stage_args yaml path. Required for streaming "
"(vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml); "
"leave unset for non-streaming to use the auto-resolved deploy config."
),
)
parser.add_argument(
"--streaming",
action="store_true",
default=False,
help="Stream audio chunks as they arrive via AsyncOmni.",
)
parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.")
parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
parser.add_argument("--ref-text", type=str, default=None, help="Transcript of the reference audio.")
parser.add_argument("--output-dir", type=str, default=None, help="Output directory for generated wav files.")
parser.add_argument("--cfg-value", type=float, default=2.0, help="Guidance value passed to VoxCPM.")
parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of diffusion timesteps.")
parser.add_argument("--min-len", type=int, default=2, help="Minimum latent length.")
parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum latent length.")
parser.add_argument(
"--streaming-prefix-len",
type=int,
default=3,
help="Streaming prefix length used by voxcpm_async_chunk.yaml.",
)
parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.")
parser.add_argument("--log-stats", action="store_true", help="Enable vLLM Omni stats logging.")
nullify_stage_engine_defaults(parser)
args = parser.parse_args()
if (args.ref_audio is None) != (args.ref_text is None):
raise ValueError("Voice cloning requires --ref-audio and --ref-text together.")
return args
def main(args) -> None:
route = "streaming" if args.streaming else "sync"
print(f"Model: {args.model}")
print(f"Deploy config: {args.deploy_config or '<auto from HF model_type>'}")
print(f"Stage configs path: {args.stage_configs_path or '<unused>'}")
print(f"Route: {route}")
if route == "streaming":
asyncio.run(_run_streaming(args))
else:
_run_sync(args)
if __name__ == "__main__":
main(parse_args())
examples/offline_inference/text_to_video/text_to_video.py
+60 -25
View File
@@ -142,6 +142,32 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
default=24000,
help="Sample rate for audio output when saved (default: 24000).",
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).",
)
# Distributed and parallel execution
parser.add_argument(
"--ulysses-mode",
type=str,
@@ -174,41 +200,23 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
)
parser.add_argument(
"--audio-sample-rate",
type=int,
default=24000,
help="Sample rate for audio output when saved (default: 24000).",
)
parser.add_argument(
"--vae-patch-parallel-size",
type=int,
default=1,
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
)
parser.add_argument(
"--pipeline-parallel-size",
type=int,
default=1,
help="Number of pipeline parallel stages.",
)
parser.add_argument(
"--enable-expert-parallel",
action="store_true",
help="Enable expert parallelism for MoE layers.",
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
"--quantization",
type=str,
default=None,
choices=["fp8", "mxfp8", "int8", "gguf"],
help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 online quant (NPU). fp8: online FP8 (GPU).",
)
parser.add_argument(
"--use-hsdp",
action="store_true",
@@ -229,6 +237,27 @@ def parse_args() -> argparse.Namespace:
return parser.parse_args()
def _extract_peak_memory_mb(result: Any) -> float:
"""Pull worker-reported peak VRAM (MiB) generation result.
Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
"""
if isinstance(result, list):
result = result[0] if result else None
if result is None:
return 0.0
val = getattr(result, "peak_memory_mb", 0.0)
if not val:
inner = getattr(result, "request_output", None)
if isinstance(inner, list):
inner = inner[0] if inner else None
val = getattr(inner, "peak_memory_mb", 0.0)
try:
return float(val or 0.0)
except (TypeError, ValueError):
return 0.0
def main():
args = parse_args()
model_class_name = args.model_class_name
@@ -262,6 +291,7 @@ def main():
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
vae_patch_parallel_size=args.vae_patch_parallel_size,
pipeline_parallel_size=args.pipeline_parallel_size,
enable_expert_parallel=args.enable_expert_parallel,
)
@@ -307,7 +337,8 @@ def main():
print(
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, pipeline_parallel_size={args.pipeline_parallel_size},"
f" enable_expert_parallel={args.enable_expert_parallel}"
)
print(f" Video size: {args.width}x{args.height}")
print(f"{'=' * 60}\n")
@@ -338,6 +369,10 @@ def main():
# Print profiling results
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
peak_mb = _extract_peak_memory_mb(frames)
if peak_mb:
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
audio = None
if isinstance(frames, list):
frames = frames[0] if frames else None
examples/online_serving/text_to_speech/README.md
+4 -67
View File
@@ -19,7 +19,6 @@ For the full list of supported architectures across all modalities, see
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
| VoxCPM | `openbmb/VoxCPM-0.5B` | ✓ | ✓ (PCM stream) | — | — |
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
@@ -294,10 +293,10 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
./qwen3_tts/run_server.sh Base
```
### Choosing an executor backend (uniproc vs mp)
Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation.
### Executor backend
Single-GPU serves now default to the uniproc executor (lower IPC overhead, the Base cloning use case from [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) / [#2604](https://github.com/vllm-project/vllm-omni/pull/2604)). `vllm_omni/deploy/qwen3_tts.yaml` is the only Qwen3-TTS deploy config; pass `--deploy-config <path>` to override.
To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor.
To opt out of chunked streaming, pass `--no-async-chunk` — the pipeline auto-dispatches to the end-to-end codec processor.
### Sending requests
```bash
@@ -375,73 +374,11 @@ python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-de
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
### Notes
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
---
## VoxCPM
Split-stage TTS at 24 kHz.
### Prerequisites
```bash
pip install voxcpm
# or use a local source tree:
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
```
If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir:
```bash
export VOXCPM_MODEL=/path/to/voxcpm-model
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
```
### Launch
```bash
export VOXCPM_MODEL=/path/to/voxcpm-model
./voxcpm/run_server.sh # async-chunk streaming (default)
./voxcpm/run_server.sh sync # non-streaming
```
Or directly:
```bash
vllm serve "$VOXCPM_MODEL" \
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
--trust-remote-code --enforce-eager --omni --port 8091
```
### Sending requests
```bash
# Basic TTS
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This is a VoxCPM online text-to-speech example."
# Voice cloning
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This sentence is synthesized with a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in reference.wav."
# Streaming PCM
python voxcpm/openai_speech_client.py \
--model "$VOXCPM_MODEL" \
--text "This is a streaming VoxCPM request." \
--stream --output voxcpm_stream.pcm
```
### Notes
- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`.
- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`.
- For benchmark measurement, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).
---
## VoxCPM2
Single-stage native AR TTS at 48 kHz.
examples/online_serving/text_to_speech/cosyvoice3/run_server.sh
Executable
+28
View File
@@ -0,0 +1,28 @@
#!/bin/bash
# Launch vLLM-Omni server for CosyVoice3 TTS
#
# Usage:
# ./run_server.sh
# CUDA_VISIBLE_DEVICES=0 ./run_server.sh
#
# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
# Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
set -e
MODEL="${MODEL:-FunAudioLLM/Fun-CosyVoice3-0.5B-2512}"
PORT="${PORT:-8091}"
EXTRA_ARGS=()
if [[ -n "${NO_ASYNC_CHUNK:-}" ]]; then
EXTRA_ARGS+=(--no-async-chunk)
fi
echo "Starting CosyVoice3 server with model: $MODEL"
vllm serve "$MODEL" \
--host 0.0.0.0 \
--port "$PORT" \
--trust-remote-code \
--omni \
"${EXTRA_ARGS[@]}"
examples/online_serving/text_to_speech/cosyvoice3/speech_client.py
+142
View File
@@ -0,0 +1,142 @@
"""Client for CosyVoice3 TTS via /v1/audio/speech endpoint.
CosyVoice3 has no built-in voice presets: every request is voice cloning
driven by ``ref_audio`` + ``ref_text``. The defaults below point at the
official upstream zero-shot prompt so the script runs out of the box.
Examples:
# Voice cloning with the default upstream prompt
python speech_client.py --text "收到好友从远方寄来的生日礼物。"
# Custom reference clip + transcript
python speech_client.py --text "Hello, this is a cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of the reference audio."
# Streaming PCM output
python speech_client.py --text "Hello world" --stream --output output.pcm
"""
import argparse
import base64
import os
import httpx
DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
DEFAULT_MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
# Official CosyVoice zero-shot prompt and its transcript.
DEFAULT_REF_AUDIO = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
DEFAULT_REF_TEXT = "希望你以后能够做的比我还好呦。"
def encode_audio_to_base64(audio_path: str) -> str:
"""Encode a local audio file to a base64 data URL."""
if not os.path.exists(audio_path):
raise FileNotFoundError(f"Audio file not found: {audio_path}")
ext = audio_path.lower().rsplit(".", 1)[-1]
mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}
mime_type = mime_map.get(ext, "audio/wav")
with open(audio_path, "rb") as f:
audio_b64 = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime_type};base64,{audio_b64}"
def run_tts(args) -> None:
"""Generate speech via the /v1/audio/speech API."""
payload = {
"model": args.model,
"input": args.text,
"response_format": args.response_format,
}
if args.ref_audio.startswith(("http://", "https://")):
payload["ref_audio"] = args.ref_audio
else:
payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
payload["ref_text"] = args.ref_text
if args.stream:
payload["stream"] = True
payload["response_format"] = "pcm"
print(f"Model: {args.model}")
print(f"Text: {args.text}")
print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}")
print("Generating audio...")
api_url = f"{args.api_base}/v1/audio/speech"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
if args.stream:
output_path = args.output or "output.pcm"
with httpx.Client(timeout=300.0) as client:
with client.stream("POST", api_url, json=payload, headers=headers) as resp:
if resp.status_code != 200:
print(f"Error: {resp.status_code}")
print(resp.read().decode())
return
total_bytes = 0
with open(output_path, "wb") as f:
for chunk in resp.iter_bytes():
f.write(chunk)
total_bytes += len(chunk)
print(f"Streamed {total_bytes} bytes to: {output_path}")
else:
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.text)
return
try:
text = response.content.decode("utf-8")
if text.startswith('{"error"'):
print(f"Error: {text}")
return
except UnicodeDecodeError:
pass
output_path = args.output or "output.wav"
with open(output_path, "wb") as f:
f.write(response.content)
print(f"Audio saved to: {output_path}")
def main():
parser = argparse.ArgumentParser(description="CosyVoice3 TTS client")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name")
parser.add_argument("--text", required=True, help="Text to synthesize")
parser.add_argument(
"--ref-audio",
default=DEFAULT_REF_AUDIO,
help="Reference audio for voice cloning (path or URL)",
)
parser.add_argument(
"--ref-text",
default=DEFAULT_REF_TEXT,
help="Transcript of the reference audio",
)
parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio format (default: wav)",
)
parser.add_argument("--output", "-o", default=None, help="Output file path")
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py
-155
View File
@@ -1,155 +0,0 @@
"""OpenAI-compatible client for VoxCPM via /v1/audio/speech.
Examples:
# Basic text-to-speech
python openai_speech_client.py --text "Hello from VoxCPM"
# Voice cloning
python openai_speech_client.py \
--text "This sentence uses the cloned voice." \
--ref-audio /path/to/reference.wav \
--ref-text "The exact transcript spoken in the reference audio."
# Streaming PCM output
python openai_speech_client.py \
--text "This is a streaming VoxCPM request." \
--stream \
--output output.pcm
"""
import argparse
import base64
import os
import httpx
DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
DEFAULT_MODEL = "OpenBMB/VoxCPM1.5"
def encode_audio_to_base64(audio_path: str) -> str:
"""Encode a local audio file to base64 data URL."""
if not os.path.exists(audio_path):
raise FileNotFoundError(f"Audio file not found: {audio_path}")
ext = audio_path.lower().rsplit(".", 1)[-1]
mime_map = {
"wav": "audio/wav",
"mp3": "audio/mpeg",
"flac": "audio/flac",
"ogg": "audio/ogg",
}
mime_type = mime_map.get(ext, "audio/wav")
with open(audio_path, "rb") as f:
audio_b64 = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime_type};base64,{audio_b64}"
def build_payload(args) -> dict[str, object]:
payload: dict[str, object] = {
"model": args.model,
"input": args.text,
"response_format": "pcm" if args.stream else args.response_format,
}
if args.ref_audio:
if args.ref_audio.startswith(("http://", "https://", "data:")):
payload["ref_audio"] = args.ref_audio
else:
payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
if args.ref_text:
payload["ref_text"] = args.ref_text
if args.max_new_tokens is not None:
payload["max_new_tokens"] = args.max_new_tokens
if args.stream:
payload["stream"] = True
return payload
def run_tts(args) -> None:
payload = build_payload(args)
api_url = f"{args.api_base}/v1/audio/speech"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
print(f"Model: {args.model}")
print(f"Text: {args.text}")
if args.ref_audio:
print("Mode: voice cloning")
print(f"Reference audio: {args.ref_audio}")
else:
print("Mode: text-to-speech")
if args.stream:
output_path = args.output or "voxcpm_output.pcm"
with httpx.Client(timeout=300.0) as client:
with client.stream("POST", api_url, json=payload, headers=headers) as response:
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.read().decode("utf-8", errors="ignore"))
return
total_bytes = 0
with open(output_path, "wb") as f:
for chunk in response.iter_bytes():
if not chunk:
continue
f.write(chunk)
total_bytes += len(chunk)
print(f"Streamed {total_bytes} bytes to: {output_path}")
return
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.text)
return
try:
text = response.content.decode("utf-8")
if text.startswith('{"error"'):
print(f"Error: {text}")
return
except UnicodeDecodeError:
pass
output_path = args.output or "voxcpm_output.wav"
with open(output_path, "wb") as f:
f.write(response.content)
print(f"Audio saved to: {output_path}")
def main():
parser = argparse.ArgumentParser(description="VoxCPM OpenAI-compatible speech client")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path")
parser.add_argument("--text", required=True, help="Text to synthesize")
parser.add_argument("--ref-audio", default=None, help="Reference audio path, URL, or data URL")
parser.add_argument(
"--ref-text",
default=None,
help="The exact transcript spoken in the reference audio",
)
parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "pcm", "flac", "mp3", "aac", "opus"],
help="Audio format for non-streaming mode (default: wav)",
)
parser.add_argument("--max-new-tokens", type=int, default=None, help="Maximum tokens to generate")
parser.add_argument("--output", "-o", default=None, help="Output file path")
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
examples/online_serving/text_to_speech/voxcpm/run_server.sh
-38
View File
@@ -1,38 +0,0 @@
#!/bin/bash
# Launch vLLM-Omni server for VoxCPM online speech serving.
#
# Usage:
# ./run_server.sh # default: async_chunk stage config
# ./run_server.sh async # async_chunk stage config
# ./run_server.sh sync # no-async-chunk stage config
# VOXCPM_MODEL=/path/to/model ./run_server.sh
set -e
MODE="${1:-async}"
MODEL="${VOXCPM_MODEL:-OpenBMB/VoxCPM1.5}"
case "$MODE" in
async)
STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml"
;;
sync)
STAGE_CONFIG="vllm_omni/deploy/voxcpm.yaml"
;;
*)
echo "Unknown mode: $MODE"
echo "Supported: async, sync"
exit 1
;;
esac
echo "Starting VoxCPM server with model: $MODEL"
echo "Stage config: $STAGE_CONFIG"
vllm serve "$MODEL" \
--stage-configs-path "$STAGE_CONFIG" \
--host 0.0.0.0 \
--port 8091 \
--trust-remote-code \
--enforce-eager \
--omni
pyproject.toml
+8 -1
View File
@@ -63,7 +63,7 @@ dev = [
"funasr>=1.0.0",
"lpips>=0.1.4", # for quantization fp8
"orjson>=3.11.7", # to make v1/embeddings API fast
"voxcpm>=2.0.2", # for voxcpm model Test
"voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
"torchdiffeq>=0.2.5" # for covo_audio Test
]
@@ -71,6 +71,11 @@ demo = [
"gradio>=6.7.0",
]
# Local-test-only extras
local = [
"FlagEmbedding",
]
docs = [
"mkdocs>=1.5.0",
"mkdocs-api-autonav",
@@ -184,6 +189,7 @@ markers = [
"core_model: L1&L2 tests (run in each PR)",
"advanced_model: L3 level tests (run on each merge)",
"full_model: L4 level tests (run nightly)",
"local_model: tests (run locally)",
# function module markers
"diffusion: Diffusion model tests",
"omni: Omni model tests",
@@ -252,3 +258,4 @@ fullset = "fullset"
Vai = "Vai"
tockens = "tockens"
CANN = "CANN"
ND = "ND"
recipes/Qwen/Qwen-Image-Edit.md
+171
View File
@@ -0,0 +1,171 @@
# Qwen-Image-Edit
> Image editing serving
## Summary
- Vendor: Qwen
- Model: `Qwen/Qwen-Image-Edit`
- Task: Text-guided single-image editing
- Mode: Online serving with the OpenAI-compatible multimodal endpoint
- Maintainer: Community
## When to use this recipe
Use this recipe as a baseline for serving `Qwen/Qwen-Image-Edit` for single-image editing on H200 141GB. Two configurations are validated below, sharing one software environment:
| Configuration | Parallelism | Notes |
| --- | --- | --- |
| `1x H200` | none | Single-card baseline; ~60 GiB peak VRAM on a 141 GB card |
| `2x H200` | TP=2 | Lower single-request latency (≈ 1.60× speedup) and lower per-GPU VRAM |
## References
- Model: <https://huggingface.co/Qwen/Qwen-Image-Edit>
- Example guide: [`examples/online_serving/image_to_image/README.md`](../../examples/online_serving/image_to_image/README.md)
- User guide: [`docs/user_guide/examples/online_serving/image_to_image.md`](../../docs/user_guide/examples/online_serving/image_to_image.md)
## Hardware Support
### GPU
The two hardware configurations below share the same software environment.
#### Environment
- OS: Linux
- Python: 3.12+
- Driver / runtime: CUDA-capable runtime matching the repository build
- vLLM: use versions from your current checkout, >=0.20.0
- vLLM-Omni: use versions from your current checkout
#### 1x H200 141GB
##### Command
```bash
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen-Image-Edit --omni \
--port 8092
```
The bundled launcher is equivalent:
```bash
bash examples/online_serving/image_to_image/run_server.sh
```
The server is ready when the log shows `INFO: Application startup complete.` (preceded by `Started server process [PID]` and `Waiting for application startup.`).
##### Verification
Run a single-image edit against the public example asset:
```bash
wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png
PROMPT="Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'"
IMG_B64=$(base64 -w0 qwen-bear.png)
curl -s http://localhost:8092/v1/chat/completions \
-H "Content-Type: application/json" \
-d "$(jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{
messages: [{
role: "user",
content: [
{type: "text", text: $prompt},
{type: "image_url", image_url: {url: ("data:image/png;base64," + $img)}}
]
}],
extra_body: {
height: 1024,
width: 1024,
num_inference_steps: 50,
guidance_scale: 1.0,
seed: 42
}
}')" \
| jq -r '.choices[0].message.content[0].image_url.url' \
| cut -d',' -f2- | base64 -d > qwen_image_edit_1card.png
```
The decoded PNG is a 1024×1024 RGB image of the bear mascot dancing under stylized moonlight, surrounded by floating stars and a "Be Kind" text bubble.
You can also run the offline CLI without standing up a server:
```bash
python examples/offline_inference/image_to_image/image_edit.py \
--image qwen-bear.png \
--prompt "Let this mascot dance under the moon, ..." \
--output qwen_image_edit_offline.png \
--num-inference-steps 50 \
--cfg-scale 4.0
```
Validated reference numbers for 1024×1024, 50 denoising steps, batch size 1, `guidance_scale=1.0`, seed 42 (`Cold` = first request after startup, `Warm` = identical subsequent request):
| Cold latency | Warm latency | Peak GPU memory |
| ---: | ---: | ---: |
| 17.80 s | 17.78 s | 61,091 MiB (~59.66 GiB) |
##### Notes
- Generation parameters live under `extra_body` for `/v1/chat/completions`, or as top-level form fields when using `/v1/images/edits`. See [`examples/online_serving/image_to_image/README.md`](../../examples/online_serving/image_to_image/README.md) for the full request / response schema.
- Memory: a 1024×1024 / 50-step edit uses ~60 GiB on a single H200, well within the 141 GB headroom.
- Multi-image variants of Qwen-Image-Edit (e.g. `Qwen-Image-Edit-2509`) are intentionally out of scope for this recipe and are not validated here.
#### 2x H200 141GB (TP=2)
##### Command
```bash
CUDA_VISIBLE_DEVICES=0,1 vllm serve Qwen/Qwen-Image-Edit --omni \
--port 8092 \
--tensor-parallel-size 2
```
In addition to the `Application startup complete.` line, the TP=2 server's `non-default args` log line echoes `'tensor_parallel_size': 2`, and a separate `Worker N: Initialized device and distributed environment` line is emitted for each rank — one per visible GPU.
##### Verification
Run the same `curl` command as in the **1x H200** Verification section above against the TP=2 server (same prompt, same seed, same image). Save the result to `qwen_image_edit_2card_tp2.png` so the two outputs sit side by side for comparison. During TP=2 inference, `nvidia-smi` shows sustained `~99 %` utilization on both GPUs, which is the expected signal that TP wiring is active.
Same workload as the 1x table above (1024×1024, 50 steps, `guidance_scale=1.0`, seed 42):
| Configuration | Cold latency | Warm latency | Peak per-GPU memory |
| --- | ---: | ---: | ---: |
| 1x H200 | 17.80 s | 17.78 s | 61,091 MiB (~59.66 GiB) |
| 2x H200 TP=2 | 11.16 s | 11.09 s | 49,377 MiB (~48.22 GiB) |
TP=2 reduces single-request warm latency from `17.78 s` to `11.09 s` (≈ **1.60×**) and drops per-GPU peak VRAM by ~11.4 GiB.
##### Notes
- TP=2 shards model weights across both cards, which is what produces both the latency reduction and the per-GPU VRAM drop above.
- CFG-Parallel (`--cfg-parallel-size 2`) is supported on the Qwen-Image family but only when `cfg_scale > 1`. The verification command above uses `guidance_scale=1.0` to match the public example image, so it does not exercise CFG-parallel; a separate run with `cfg_scale > 1` is needed if you want to benchmark that path.
- For deeper acceleration knobs (Cache-DiT, sequence parallel, HSDP) see [`docs/user_guide/diffusion/parallelism/overview.md`](../../docs/user_guide/diffusion/parallelism/overview.md); this recipe intentionally documents only the validated TP=2 baseline.
#### Profiling
The numbers below are from a warm request at 1024×1024, 50 steps, and `guidance_scale=1.0`, measured with `--enable-diffusion-pipeline-profiler` and a 250 ms `nvidia-smi --query-gpu=memory.used` sampler. For TP=2, latency uses rank max, and VRAM is per GPU.
Per-stage latency:
| Stage | TP=1 | TP=2 | Notes |
| --- | ---: | ---: | --- |
| `text_encoder.forward` | 0.132 s | 0.135 s | ~1.0× |
| `vae.encode` | 0.040 s | 0.041 s | ~1.0× |
| `diffuse` | 17.105 s | 10.382 s | **1.65×** speedup |
| `vae.decode` | 0.064 s | 0.066 s | ~1.0× |
| `forward` | 17.372 s | 10.652 s | 1.63× speedup |
Per-GPU VRAM:
| Snapshot | TP=1 | TP=2 | Notes |
| --- | ---: | ---: | --- |
| Static VRAM | 58,777 MiB | 47,065 MiB | after model load, before request |
| Peak VRAM | 61,091 MiB | 49,377 MiB | during inference |
| Request-time increment | 2,314 MiB | 2,312 MiB | peak minus static |
The latency gain mainly comes from the DiT denoising stage (`diffuse` in the table above). The other measured stages, such as text encoder and VAE encode/decode, stay essentially unchanged between TP=1 and TP=2, which is consistent with them not benefiting meaningfully from TP in this pipeline. The `forward` row's 1.63× speedup tracks `diffuse`, since `diffuse` accounts for about 98% of total pipeline time.
For memory, TP=2 saves about 11.4 GiB per GPU, mostly from static model-state memory, likely dominated by tensor-parallel sharding of the DiT weights. The request-time increment is almost unchanged, so TP=2 does not materially reduce transient inference-time memory for this single-image, batch-1 workload.
recipes/Qwen/Qwen-Image.md
+1 -1
View File
@@ -104,7 +104,7 @@ To replay the batching benefit with matched warmup:
```bash
python benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni \
--endpoint /v1/chat/completions \
--dataset vbench \
--task t2i \
--model Qwen/Qwen-Image \
recipes/README.md
+1
View File
@@ -35,6 +35,7 @@ recipes/
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
| [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB |
| [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB |
| [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) |
tests/benchmarks/metrics/test_metrics.py
+32
View File
@@ -63,6 +63,38 @@ def test_total_input_aggregated_from_output_prompt_len():
)
def test_audio_continuity_aggregation():
"""Continuity rate and underrun percentile must aggregate from per-output fields."""
bad = _make_output(100)
bad.audio_underrun_s = 0.5
bad.audio_continuity_ok = False
good_a = _make_output(100)
good_a.audio_underrun_s = 0.02
good_a.audio_continuity_ok = True
good_b = _make_output(100)
good_b.audio_underrun_s = 0.0
good_b.audio_continuity_ok = True
metrics, _ = calculate_metrics(
input_requests=[],
outputs=[bad, good_a, good_b],
dur_s=10.0,
tokenizer=None,
selected_percentiles=[50.0, 99.0],
goodput_config_dict={},
task_type=TaskType.GENERATION,
selected_percentile_metrics=["audio_underrun"],
max_concurrency=None,
request_rate=float("inf"),
benchmark_duration=10.0,
)
assert metrics.audio_continuity_ok_rate == pytest.approx(2 / 3, abs=1e-6)
# p99 of [0.5, 0.02, 0.0] is dominated by the 0.5 outlier.
p99 = dict(metrics.percentiles_audio_underrun_s).get(99.0)
assert p99 is not None and p99 > 0.4
# ============================================================================
# TTFT suppression for pure-audio (TTS) benchmarks
# ============================================================================
tests/benchmarks/test_audio_continuity.py
+123
View File
@@ -0,0 +1,123 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the audio-streaming continuity helper.
The helper is intentionally vllm-free: it operates on a flat timeline of
``(arrival_time_s, num_bytes)`` chunks plus PCM format params. These tests
exercise the underrun arithmetic against hand-crafted timelines so future
refactors of the streaming backend cannot silently regress the metric.
"""
from __future__ import annotations
import pytest
from vllm_omni.benchmarks.audio_continuity import (
ContinuityStats,
compute_continuity_stats,
)
pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
# PCM s16le mono at 24 kHz -> 48 000 bytes/sec, the Qwen3-TTS / VoxCPM2 default.
_SR = 24_000
_BPS = _SR * 2 # 48 000
def test_empty_timeline_is_continuous() -> None:
stats = compute_continuity_stats([], [], sample_rate=_SR)
assert stats == ContinuityStats(
max_underrun_s=0.0,
underrun_event_count=0,
is_continuous=True,
)
def test_single_chunk_has_no_underrun() -> None:
# One large chunk, no later arrivals to check against.
stats = compute_continuity_stats(
chunk_arrival_times_s=[0.05],
chunk_bytes=[_BPS], # one second of audio in a single chunk
sample_rate=_SR,
)
assert stats.max_underrun_s == 0.0
assert stats.underrun_event_count == 0
assert stats.is_continuous is True
def test_chunks_faster_than_realtime_have_no_underrun() -> None:
# 100ms of audio per chunk, but arrive every 50ms — buffer always builds.
stats = compute_continuity_stats(
chunk_arrival_times_s=[0.0, 0.05, 0.10, 0.15],
chunk_bytes=[_BPS // 10] * 4,
sample_rate=_SR,
)
assert stats.max_underrun_s == 0.0
assert stats.underrun_event_count == 0
assert stats.is_continuous is True
def test_single_gap_above_threshold_breaks_continuity() -> None:
# First chunk carries 100ms of audio at t=0. Next chunk arrives at t=0.5s,
# so the player ran dry from t=0.1 to t=0.5 -> 400ms underrun.
stats = compute_continuity_stats(
chunk_arrival_times_s=[0.0, 0.5],
chunk_bytes=[_BPS // 10, _BPS // 10], # 100ms each
sample_rate=_SR,
threshold_s=0.1,
)
assert stats.max_underrun_s == pytest.approx(0.4, abs=1e-6)
assert stats.underrun_event_count == 1
assert stats.is_continuous is False
def test_underrun_is_max_across_gaps_not_sum() -> None:
# Two gaps: a 200ms shortfall and a 50ms shortfall. max_underrun_s
# must surface the 200ms one, not the sum.
stats = compute_continuity_stats(
chunk_arrival_times_s=[0.0, 0.3, 0.45],
chunk_bytes=[_BPS // 10, _BPS // 10, _BPS // 10], # 100ms chunks
sample_rate=_SR,
threshold_s=0.1,
)
# At t=0.3 the player consumed 300ms but only 100ms received -> 200ms deficit
# At t=0.45 player consumed 450ms, received 200ms -> 250ms deficit.
assert stats.max_underrun_s == pytest.approx(0.25, abs=1e-6)
assert stats.underrun_event_count == 2
assert stats.is_continuous is False
def test_continuity_threshold_boundary_inclusive() -> None:
# First chunk = 100ms audio at t=0. Next arrives at t=0.2 -> 100ms deficit
# exactly. With threshold=0.1, continuity_ok stays True (boundary).
stats = compute_continuity_stats(
chunk_arrival_times_s=[0.0, 0.2],
chunk_bytes=[_BPS // 10, _BPS // 10],
sample_rate=_SR,
threshold_s=0.1,
)
assert stats.max_underrun_s == pytest.approx(0.1, abs=1e-6)
assert stats.is_continuous is True
def test_mismatched_lengths_returns_zero_stats() -> None:
# Defensive: never crash on malformed timeline.
stats = compute_continuity_stats(
chunk_arrival_times_s=[0.0, 0.1],
chunk_bytes=[100],
sample_rate=_SR,
)
assert stats == ContinuityStats(0.0, 0, True)
def test_threshold_is_configurable() -> None:
# 200ms underrun: violates 100ms threshold, OK against 500ms threshold.
arrivals = [0.0, 0.3]
bytes_list = [_BPS // 10, _BPS // 10]
strict = compute_continuity_stats(arrivals, bytes_list, sample_rate=_SR, threshold_s=0.1)
lenient = compute_continuity_stats(arrivals, bytes_list, sample_rate=_SR, threshold_s=0.5)
assert strict.is_continuous is False
assert lenient.is_continuous is True
assert strict.max_underrun_s == lenient.max_underrun_s
tests/benchmarks/test_diffusion_backends_metrics.py
+32 -1
View File
@@ -1,6 +1,11 @@
import pytest
from benchmarks.diffusion.backends import RequestFuncInput, async_request_chat_completions
from benchmarks.diffusion.backends import (
RequestFuncInput,
async_request_chat_completions,
endpoint_filename_token,
normalize_endpoint,
)
class _MockResponse:
@@ -29,6 +34,32 @@ class _MockSession:
return _MockResponse(self._payload)
@pytest.mark.core_model
@pytest.mark.benchmark
@pytest.mark.cpu
def test_endpoint_normalization_accepts_optional_leading_slash():
assert normalize_endpoint("v1/videos") == "/v1/videos"
assert normalize_endpoint("/v1/videos") == "/v1/videos"
assert normalize_endpoint("v1/chat/completions") == "/v1/chat/completions"
assert normalize_endpoint("v1/images/generations") == "/v1/images/generations"
@pytest.mark.core_model
@pytest.mark.benchmark
@pytest.mark.cpu
def test_endpoint_normalization_accepts_legacy_backend_aliases():
assert normalize_endpoint("vllm-omni") == "/v1/chat/completions"
assert normalize_endpoint("openai") == "/v1/images/generations"
@pytest.mark.core_model
@pytest.mark.benchmark
@pytest.mark.cpu
def test_endpoint_filename_token_drops_leading_slash():
assert endpoint_filename_token("/v1/videos") == "v1_videos"
assert endpoint_filename_token("v1/chat/completions") == "v1_chat_completions"
@pytest.mark.core_model
@pytest.mark.benchmark
@pytest.mark.cpu
tests/conftest.py
+2
View File
@@ -36,6 +36,7 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config):
_ASSERTION_EXPORT_NAMES = (
"assert_audio_speech_response",
"assert_diffusion_response",
"assert_http_error",
"assert_image_diffusion_response",
"assert_image_valid",
"assert_omni_response",
@@ -54,6 +55,7 @@ _MEDIA_EXPORT_NAMES = (
_STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",)
_RUNTIME_EXPORT_NAMES = (
"DiffusionResponse",
"HttpResponse",
"OmniResponse",
"OmniRunner",
"OmniRunnerHandler",
tests/dfx/conftest.py
+32 -5
View File
@@ -279,6 +279,10 @@ def create_benchmark_indices(
return indices
# Omni perf/stability: ``vllm bench serve --omni`` result template (used if bench exits without writing JSON).
OMNI_RESULT_TEMPLATE_PATH = Path(__file__).resolve().parent / "perf" / "scripts" / "result_omni_template.json"
def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
"""Make a single path segment safe for result filenames on common filesystems."""
if value is None:
@@ -289,7 +293,7 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
return s if s else default
def _resolve_baseline_value(
def resolve_baseline_value(
baseline_raw: Any,
*,
sweep_index: int | None,
@@ -330,7 +334,7 @@ def _baseline_thresholds_for_step(
) -> dict[str, Any]:
"""Resolve baseline config to one threshold per metric for this iteration."""
return {
metric_name: _resolve_baseline_value(
metric_name: resolve_baseline_value(
baseline_raw,
sweep_index=sweep_index,
max_concurrency=max_concurrency,
@@ -354,7 +358,12 @@ def run_benchmark(
random_input_len: Any | None = None,
random_output_len: Any | None = None,
) -> dict[str, Any]:
"""Run one ``vllm bench serve --omni`` iteration and return parsed metrics."""
"""Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the resolved per-metric thresholds
(when ``baseline_config`` is provided). If the benchmark exits without writing a result file,
``result_omni_template.json`` is used as a fallback.
"""
current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
ri = _safe_filename_token(random_input_len)
ro = _safe_filename_token(random_output_len)
@@ -403,8 +412,17 @@ def run_benchmark(
result_dir = "./"
result_path = os.path.join(result_dir, result_filename)
with open(result_path, encoding="utf-8") as f:
result = json.load(f)
if not os.path.exists(result_path):
with open(OMNI_RESULT_TEMPLATE_PATH, encoding="utf-8") as f:
template_result: dict[str, Any] = json.load(f)
Path(result_path).parent.mkdir(parents=True, exist_ok=True)
with open(result_path, "w", encoding="utf-8") as f:
json.dump(template_result, f, ensure_ascii=False, indent=2)
print(f"Benchmark result file not generated, fallback to template: {result_path}")
result = template_result
else:
with open(result_path, encoding="utf-8") as f:
result = json.load(f)
if baseline_config:
result["baseline"] = _baseline_thresholds_for_step(
@@ -432,3 +450,12 @@ def pytest_addoption(parser: pytest.Parser) -> None:
default=None,
help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"),
)
parser.addoption(
"--assert-baseline",
action="store_true",
default=False,
help=(
"When set, omni/diffusion perf runners compare metrics against the baseline block in the JSON config "
"(default: off)."
),
)
tests/dfx/perf/scripts/diffusion_result_template.json
+3 -2
View File
@@ -1,7 +1,7 @@
[
{
"test_name": null,
"backend": null,
"endpoint": null,
"timestamp": null,
"server_params": {
"model": null,
@@ -43,7 +43,7 @@
"stage_durations_mean": {},
"stage_durations_p50": {},
"stage_durations_p99": {},
"backend": null,
"endpoint": null,
"model": null,
"dataset": null,
"task": null
@@ -51,6 +51,7 @@
"log_file": null,
"Model": null,
"Framework": null,
"API Endpoint": null,
"Hardware": null,
"Deployment": null,
"Task": null,
tests/dfx/perf/scripts/run_benchmark.py
+20 -75
View File
@@ -12,12 +12,17 @@ from tests.dfx.conftest import (
create_unique_server_params,
get_benchmark_params_for_server,
load_configs,
resolve_baseline_value,
run_benchmark,
)
from tests.helpers.runtime import OmniServer
pytestmark = [pytest.mark.full_model]
# Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline``
# (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are
# defined in the same module.
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -46,7 +51,6 @@ if CONFIG_FILE_PATH is None:
CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE
BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
OMNI_RESULT_TEMPLATE_PATH = Path(__file__).parent / "result_omni_template.json"
DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy"
@@ -87,16 +91,6 @@ def omni_server(request):
print("OmniServer stopped")
def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
"""Make a single path segment safe for result filenames on common filesystems."""
if value is None:
return default
s = str(value).strip()
for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
s = s.replace(bad, "_")
return s if s else default
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
@@ -126,76 +120,23 @@ def benchmark_params(request, omni_server):
}
def _resolve_baseline_value(
baseline_raw: Any,
*,
sweep_index: int | None,
max_concurrency: Any = None,
request_rate: Any = None,
) -> Any:
"""Pick the baseline threshold for this sweep step.
Supported shapes per metric:
- **Scalar** — same threshold for every concurrency / QPS.
- **List** — aligned with ``max_concurrency`` / ``request_rate`` sweep order; use ``sweep_index``.
- **Dict** — keyed by concurrency or rate, e.g. ``{"1": 500, "4": 800}`` (keys are strings in JSON).
For dict lookup, ``max_concurrency`` is preferred when both are set (concurrency sweep).
"""
if baseline_raw is None:
# If no baseline is set, the maximum value will be used.
return 100000
if isinstance(baseline_raw, dict):
if max_concurrency is not None:
for key in (max_concurrency, str(max_concurrency)):
if key in baseline_raw:
return baseline_raw[key]
if request_rate is not None:
for key in (request_rate, str(request_rate)):
if key in baseline_raw:
return baseline_raw[key]
raise KeyError(
f"baseline dict has no key for max_concurrency={max_concurrency!r} "
f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
)
if isinstance(baseline_raw, (list, tuple)):
return baseline_raw[sweep_index]
return baseline_raw
def _baseline_thresholds_for_step(
baseline_data: dict[str, Any],
*,
sweep_index: int | None = None,
max_concurrency: Any = None,
request_rate: Any = None,
) -> dict[str, Any]:
"""Resolve ``test.json`` ``baseline`` block to one threshold per metric (same as ``assert_result``)."""
return {
metric_name: _resolve_baseline_value(
baseline_raw,
sweep_index=sweep_index,
max_concurrency=max_concurrency,
request_rate=request_rate,
)
for metric_name, baseline_raw in baseline_data.items()
}
def assert_result(
result,
params,
num_prompt,
*,
assert_baseline: bool,
sweep_index: int | None = None,
max_concurrency: Any = None,
request_rate: Any = None,
max_concurrency: Any | None = None,
request_rate: Any | None = None,
) -> None:
assert result["completed"] == num_prompt, "Request failures exist"
if not assert_baseline:
return
baseline_data = params.get("baseline", {})
for metric_name, baseline_raw in baseline_data.items():
current_value = result[metric_name]
baseline_value = _resolve_baseline_value(
baseline_value = resolve_baseline_value(
baseline_raw,
sweep_index=sweep_index,
max_concurrency=max_concurrency,
@@ -214,7 +155,7 @@ def assert_result(
@pytest.mark.benchmark
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
def test_performance_benchmark(omni_server, benchmark_params):
def test_performance_benchmark(omni_server, benchmark_params, request):
test_name = benchmark_params["test_name"]
params = benchmark_params["params"]
dataset_name = params.get("dataset_name", "")
@@ -226,6 +167,8 @@ def test_performance_benchmark(omni_server, benchmark_params):
print(f"Running benchmark for model: {model}")
print(f"Benchmark parameters: {benchmark_params}")
assert_baseline = request.config.getoption("--assert-baseline", default=False)
def to_list(value, default=None):
if value is None:
return [] if default is None else [default]
@@ -264,7 +207,7 @@ def test_performance_benchmark(omni_server, benchmark_params):
elif not isinstance(value, bool):
args.extend([arg_name, str(value)])
# QPS test (sweep_index aligns with qps_list / num_prompt_list for this loop)
# QPS / request-rate sweep
for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
result = run_benchmark(
@@ -283,12 +226,13 @@ def test_performance_benchmark(omni_server, benchmark_params):
assert_result(
result,
params,
num_prompt=num_prompt,
num_prompt,
assert_baseline=assert_baseline,
sweep_index=i,
request_rate=qps,
)
# concurrency test (sweep_index aligns with max_concurrency_list for separate thresholds per concurrency)
# concurrency test
for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
result = run_benchmark(
@@ -307,7 +251,8 @@ def test_performance_benchmark(omni_server, benchmark_params):
assert_result(
result,
params,
num_prompt=num_prompt,
num_prompt,
assert_baseline=assert_baseline,
sweep_index=i,
max_concurrency=concurrency,
)
tests/dfx/perf/scripts/run_diffusion_benchmark.py
+40 -35
View File
@@ -5,16 +5,17 @@ This runner separates two concepts:
1. ``server_type``: how the serving process is started.
Currently only ``vllm-omni`` is supported here.
2. ``benchmark_backend``: which serving API the benchmark client calls.
Examples: ``vllm-omni`` for ``/v1/chat/completions`` and ``v1/videos``
for async video jobs.
2. ``benchmark_endpoint``: which serving API the benchmark client calls.
Examples: ``/v1/chat/completions`` and ``/v1/videos``.
A config JSON file is REQUIRED via --test-config-file:
pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off).
All benchmark results for a session are consolidated into a single JSON file under
BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable).
Each entry in the file contains the test metadata (test_name, backend, benchmark_params,
Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params,
timestamp) together with the raw metrics returned by the benchmark script.
"""
@@ -33,6 +34,8 @@ from typing import Any, cast
import psutil
import pytest
from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -400,7 +403,7 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]
"model": cfg["server_params"]["model"],
"serve_args_dict": serve_args_dict,
"serve_args": _build_serve_args(serve_args_dict),
"benchmark_backend": cfg.get("benchmark_backend"),
"benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")),
"server_params": cfg["server_params"],
}
)
@@ -482,7 +485,7 @@ def run_benchmark(
model: str,
params: dict[str, Any],
test_name: str,
backend: str = "vllm-omni",
endpoint: str = "/v1/chat/completions",
server_cfg: dict[str, Any] | None = None,
source_file: str = "",
) -> dict[str, Any]:
@@ -490,18 +493,19 @@ def run_benchmark(
The raw metrics are written to a temporary file by the subprocess. After
the run completes the metrics are merged with full metadata (test_name,
backend, benchmark_params, timestamp, flat reporting fields) and appended
endpoint, benchmark_params, timestamp, flat reporting fields) and appended
to the session-wide aggregated JSON file (AGGREGATED_RESULT_FILE). The
temporary file is removed afterwards. Subprocess stdout/stderr are tee'd
to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in
the record.
"""
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
endpoint = normalize_endpoint(endpoint)
log_dir = BENCHMARK_RESULT_DIR / "logs"
log_dir.mkdir(parents=True, exist_ok=True)
backend_label = backend.replace("/", "_")
log_file = log_dir / f"{test_name}_{backend_label}_{timestamp}.log"
endpoint_label = endpoint_filename_token(endpoint)
log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
tmp_result_file = Path(tmp.name)
@@ -517,8 +521,8 @@ def run_benchmark(
str(port),
"--model",
model,
"--backend",
backend,
"--endpoint",
endpoint,
"--dataset",
params.get("dataset", "random"),
"--task",
@@ -545,7 +549,7 @@ def run_benchmark(
# cause truncated or out-of-order log output when stdout is piped).
cmd = [cmd[0], "-u"] + cmd[1:]
print(f"\nRunning benchmark (backend={backend}): {' '.join(cmd)}")
print(f"\nRunning benchmark (endpoint={endpoint}): {' '.join(cmd)}")
print(f" Log file: {log_file}")
# Redirect stdout + stderr directly to the log file at the OS level
@@ -597,7 +601,7 @@ def run_benchmark(
record: dict[str, Any] = {
"test_name": test_name,
"backend": backend,
"endpoint": endpoint,
"timestamp": timestamp,
"server_params": server_cfg.get("server_params"),
"benchmark_params": params,
@@ -605,7 +609,7 @@ def run_benchmark(
"log_file": str(log_file),
"Model": model,
"Framework": server_type,
"API Backend": backend,
"API Endpoint": endpoint,
"Hardware": "",
"Deployment": "",
"Task": params.get("task", "t2i"),
@@ -650,12 +654,14 @@ def run_benchmark(
# ---------------------------------------------------------------------------
def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None:
"""Assert that benchmark metrics satisfy the configured baselines."""
def assert_result(result: dict[str, Any], params: dict[str, Any], *, assert_baseline: bool) -> None:
"""Assert request completion; optionally assert metrics against ``baseline`` in *params*."""
num_prompts = params.get("num-prompts", 10)
completed = result.get("completed_requests", result.get("completed", 0))
assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
if not assert_baseline:
return
if params.get("skip-performance-assertion", False):
print("Skipping performance assertions.")
return
@@ -669,21 +675,21 @@ def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None:
assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}"
def _default_benchmark_backend_for_task(task: str) -> str:
"""Return the default client-side benchmark backend for a diffusion task."""
def _default_benchmark_endpoint_for_task(task: str) -> str:
"""Return the default client-side benchmark endpoint for a diffusion task."""
if task in {"t2v", "i2v", "ti2v"}:
return "v1/videos"
return "/v1/videos"
if task in {"t2i", "i2i", "ti2i"}:
return "vllm-omni"
raise ValueError(f"Unsupported task for benchmark backend resolution: {task}")
return "/v1/chat/completions"
raise ValueError(f"Unsupported task for benchmark endpoint resolution: {task}")
def _resolve_benchmark_backend(server_cfg: dict[str, Any], params: dict[str, Any]) -> str:
def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, Any]) -> str:
"""Resolve which serving API the benchmark client should call."""
configured = server_cfg.get("benchmark_backend")
configured = server_cfg.get("benchmark_endpoint")
if configured:
return cast(str, configured)
return _default_benchmark_backend_for_task(cast(str, params.get("task", "t2i")))
return normalize_endpoint(cast(str, configured))
return _default_benchmark_endpoint_for_task(cast(str, params.get("task", "t2i")))
# ---------------------------------------------------------------------------
@@ -697,21 +703,19 @@ def _resolve_benchmark_backend(server_cfg: dict[str, Any], params: dict[str, Any
indirect=True,
)
@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
"""Run the diffusion performance benchmark and assert against baselines.
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
"""Run the diffusion performance benchmark and verify request completion.
One server is started per unique parallel configuration (module scope).
For each server, all benchmark parameter sets defined in the config JSON
are executed sequentially; results are asserted against the baselines.
are executed sequentially; metrics are recorded to the aggregated result file.
Tracked metrics:
- throughput_qps (higher is better)
- latency_p50, latency_p99 (lower is better)
Pass ``--assert-baseline`` to compare ``throughput_qps``, ``latency_mean``, etc. to the JSON ``baseline`` block.
"""
test_name = benchmark_params["test_name"]
params = benchmark_params["params"]
server_cfg = getattr(diffusion_server, "server_cfg", {})
backend = _resolve_benchmark_backend(server_cfg, params)
endpoint = _resolve_benchmark_endpoint(server_cfg, params)
result = run_benchmark(
host=diffusion_server.host,
@@ -719,13 +723,13 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
model=diffusion_server.model,
params=params,
test_name=test_name,
backend=backend,
endpoint=endpoint,
server_cfg=server_cfg,
source_file=cast(str, CONFIG_FILE_PATH),
)
print(f"\n{'=' * 60}")
print(f"Results for {test_name} (server={diffusion_server.server_type}, backend={backend}):")
print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
for key in (
"throughput_qps",
"latency_mean",
@@ -742,4 +746,5 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
print(f"\n Aggregated results: {AGGREGATED_RESULT_FILE}")
print("=" * 60)
assert_result(result, params)
assert_baseline = request.config.getoption("--assert-baseline", default=False)
assert_result(result, params, assert_baseline=assert_baseline)
tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
+159 -37
View File
@@ -3,31 +3,67 @@
"test_name": "test_qwen3_omni_chunk",
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"extra_cli_args": ["--async-chunk"]
"extra_cli_args": [
"--async-chunk"
]
},
"benchmark_params": [
{
"dataset_name": "random",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [4, 16, 32, 64, 128],
"max_concurrency": [1, 4, 8, 16, 32],
"num_prompts": [
4,
16,
32,
64,
128
],
"max_concurrency": [
1,
4,
8,
16,
32
],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
"mean_audio_ttfp_ms": [1000, 3000, 5000, 7000, 9000],
"mean_audio_rtf": [0.2, 0.35, 0.6, 0.85, 0.9]
"mean_ttft_ms": [
400,
300,
400,
600,
900
],
"mean_e2el_ms": [
26000,
43000,
52000,
72000,
109000
],
"mean_audio_rtf": [
0.15,
0.30,
0.35,
0.50,
0.75
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [10],
"request_rate": [0.1],
"num_prompts": [
10
],
"request_rate": [
0.1
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -42,17 +78,27 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [2000],
"mean_audio_ttfp_ms": [2000],
"mean_audio_rtf": [0.25]
"mean_ttft_ms": [
600
],
"mean_e2el_ms": [
6000
],
"mean_audio_rtf": [
0.20
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [40],
"request_rate": [0.5],
"num_prompts": [
40
],
"request_rate": [
0.5
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -69,17 +115,27 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [6000],
"mean_audio_ttfp_ms": [6000],
"mean_audio_rtf": [0.7]
"mean_ttft_ms": [
300
],
"mean_e2el_ms": [
7000
],
"mean_audio_rtf": [
0.25
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [100],
"request_rate": [1.0],
"num_prompts": [
100
],
"request_rate": [
1.0
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -98,40 +154,79 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [12000],
"mean_audio_ttfp_ms": [12000],
"mean_audio_rtf": [1.0]
"mean_ttft_ms": [
700
],
"mean_e2el_ms": [
11000
],
"mean_audio_rtf": [
0.35
]
}
},
{
"dataset_name": "random",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [4, 16, 32, 64, 128],
"max_concurrency": [1, 4, 8, 16, 32],
"num_prompts": [
4,
16,
32,
64,
128
],
"max_concurrency": [
1,
4,
8,
16,
32
],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"extra_body": {
"modalities": ["text"]
"modalities": [
"text"
]
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
"mean_ttft_ms": [
100,
300,
500,
800,
1500
],
"mean_e2el_ms": [
5000,
8000,
10000,
14000,
20000
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [10],
"request_rate": [0.1],
"num_prompts": [
10
],
"request_rate": [
0.1
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"ignore_eos": true,
"extra_body": {
"modalities": ["text"]
"modalities": [
"text"
]
},
"random_mm_base_items_per_request": 1,
"random_mm_num_mm_items_range_ratio": 0.5,
@@ -143,21 +238,32 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"mean_ttft_ms": [2000]
"mean_ttft_ms": [
500
],
"mean_e2el_ms": [
1000
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [40],
"request_rate": [0.5],
"num_prompts": [
40
],
"request_rate": [
0.5
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"ignore_eos": true,
"extra_body": {
"modalities": ["text"]
"modalities": [
"text"
]
},
"random_mm_base_items_per_request": 2,
"random_mm_num_mm_items_range_ratio": 0.5,
@@ -171,21 +277,32 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"mean_ttft_ms": [6000]
"mean_ttft_ms": [
200
],
"mean_e2el_ms": [
800
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [100],
"request_rate": [1.0],
"num_prompts": [
100
],
"request_rate": [
1.0
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
"ignore_eos": true,
"extra_body": {
"modalities": ["text"]
"modalities": [
"text"
]
},
"random_mm_base_items_per_request": 3,
"random_mm_num_mm_items_range_ratio": 0.5,
@@ -201,7 +318,12 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"mean_ttft_ms": [6000]
"mean_ttft_ms": [
600
],
"mean_e2el_ms": [
1300
]
}
}
]
tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
+83 -21
View File
@@ -3,31 +3,67 @@
"test_name": "test_qwen3_omni",
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"extra_cli_args": ["--no-async-chunk"]
"extra_cli_args": [
"--no-async-chunk"
]
},
"benchmark_params": [
{
"dataset_name": "random",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [4, 16, 32, 64, 128],
"max_concurrency": [1, 4, 8, 16, 32],
"num_prompts": [
4,
16,
32,
64,
128
],
"max_concurrency": [
1,
4,
8,
16,
32
],
"random_input_len": 2500,
"random_output_len": 900,
"ignore_eos": true,
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
"mean_audio_ttfp_ms": [30000, 60000, 90000, 120000, 150000],
"mean_audio_rtf": [0.35, 0.45, 0.55, 0.65, 0.75]
"mean_ttft_ms": [
400,
300,
400,
600,
900
],
"mean_e2el_ms": [
32000,
34000,
45000,
61000,
90000
],
"mean_audio_rtf": [
0.20,
0.25,
0.30,
0.45,
0.60
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [10],
"request_rate": [0.1],
"num_prompts": [
10
],
"request_rate": [
0.1
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -42,17 +78,27 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [2000],
"mean_audio_ttfp_ms": [10000],
"mean_audio_rtf": [0.25]
"mean_ttft_ms": [
500
],
"mean_e2el_ms": [
6000
],
"mean_audio_rtf": [
0.20
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [40],
"request_rate": [0.5],
"num_prompts": [
40
],
"request_rate": [
0.5
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -69,17 +115,27 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [6000],
"mean_audio_ttfp_ms": [15000],
"mean_audio_rtf": [0.45]
"mean_ttft_ms": [
300
],
"mean_e2el_ms": [
7000
],
"mean_audio_rtf": [
0.20
]
}
},
{
"dataset_name": "random-mm",
"backend": "openai-chat-omni",
"endpoint": "/v1/chat/completions",
"num_prompts": [100],
"request_rate": [1.0],
"num_prompts": [
100
],
"request_rate": [
1.0
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -98,9 +154,15 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_ttft_ms": [12000],
"mean_audio_ttfp_ms": [18000],
"mean_audio_rtf": [0.9]
"mean_ttft_ms": [
700
],
"mean_e2el_ms": [
9000
],
"mean_audio_rtf": [
0.30
]
}
}
]
tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json
+79 -17
View File
@@ -4,7 +4,9 @@
"server_params": {
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"use_omni": false,
"extra_cli_args": ["--no-enable-prefix-caching"]
"extra_cli_args": [
"--no-enable-prefix-caching"
]
},
"benchmark_params": [
{
@@ -12,8 +14,20 @@
"backend": "vllm",
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"endpoint": "/v1/completions",
"num_prompts": [4, 16, 32, 64, 128],
"max_concurrency": [1, 4, 8, 16, 32],
"num_prompts": [
4,
16,
32,
64,
128
],
"max_concurrency": [
1,
4,
8,
16,
32
],
"random_input_len": 2500,
"random_output_len": 900,
"temperature": 0.4,
@@ -23,11 +37,26 @@
"repetition_penalty": 1.05,
"ignore_eos": true,
"extra_body": {
"modalities": ["text"]
"modalities": [
"text"
]
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
"mean_ttft_ms": [
300,
300,
1000,
600,
800
],
"mean_e2el_ms": [
6000,
8000,
11000,
13000,
17000
]
}
},
{
@@ -35,8 +64,12 @@
"backend": "vllm",
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"endpoint": "/v1/completions",
"num_prompts": [10],
"request_rate": [0.1],
"num_prompts": [
10
],
"request_rate": [
0.1
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -47,7 +80,9 @@
"repetition_penalty": 1.05,
"ignore_eos": true,
"extra_body": {
"modalities": ["text"]
"modalities": [
"text"
]
},
"random_mm_base_items_per_request": 1,
"random_mm_num_mm_items_range_ratio": 0.5,
@@ -59,7 +94,12 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"mean_ttft_ms": [2000]
"mean_ttft_ms": [
30
],
"mean_e2el_ms": [
540
]
}
},
{
@@ -67,8 +107,12 @@
"backend": "vllm",
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"endpoint": "/v1/completions",
"num_prompts": [40],
"request_rate": [0.5],
"num_prompts": [
40
],
"request_rate": [
0.5
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -79,7 +123,9 @@
"repetition_penalty": 1.05,
"ignore_eos": true,
"extra_body": {
"modalities": ["text"]
"modalities": [
"text"
]
},
"random_mm_base_items_per_request": 2,
"random_mm_num_mm_items_range_ratio": 0.5,
@@ -93,7 +139,12 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"mean_ttft_ms": [6000]
"mean_ttft_ms": [
30
],
"mean_e2el_ms": [
600
]
}
},
{
@@ -101,8 +152,12 @@
"backend": "vllm",
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
"endpoint": "/v1/completions",
"num_prompts": [100],
"request_rate": [1.0],
"num_prompts": [
100
],
"request_rate": [
1.0
],
"random_input_len": 100,
"random_output_len": 100,
"random_range_ratio": 0.0,
@@ -113,7 +168,9 @@
"repetition_penalty": 1.05,
"ignore_eos": true,
"extra_body": {
"modalities": ["text"]
"modalities": [
"text"
]
},
"random_mm_base_items_per_request": 3,
"random_mm_num_mm_items_range_ratio": 0.5,
@@ -129,7 +186,12 @@
},
"percentile-metrics": "ttft,tpot,itl,e2el",
"baseline": {
"mean_ttft_ms": [6000]
"mean_ttft_ms": [
40
],
"mean_e2el_ms": [
700
]
}
}
]
tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
+22 -22
View File
@@ -22,10 +22,10 @@
"num-input-images": 2,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.05,
"latency_mean": 18,
"peak_memory_mb_max": 78500,
"peak_memory_mb_mean": 78500
"throughput_qps": 0.06,
"latency_mean": 20,
"peak_memory_mb_max": 57000,
"peak_memory_mb_mean": 57000
}
},
{
@@ -41,9 +41,9 @@
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.01,
"latency_mean": 70,
"peak_memory_mb_max": 81000,
"peak_memory_mb_mean": 81000
"latency_mean": 60,
"peak_memory_mb_max": 67000,
"peak_memory_mb_mean": 67000
}
}
]
@@ -75,10 +75,10 @@
"num-input-images": 2,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.1,
"latency_mean": 12,
"peak_memory_mb_max": 69000,
"peak_memory_mb_mean": 69000
"throughput_qps": 0.15,
"latency_mean": 10,
"peak_memory_mb_max": 55000,
"peak_memory_mb_mean": 55000
}
},
{
@@ -93,10 +93,10 @@
"num-input-images": 2,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.03,
"latency_mean": 28,
"peak_memory_mb_max": 69000,
"peak_memory_mb_mean": 69000
"throughput_qps": 0.05,
"latency_mean": 20,
"peak_memory_mb_max": 55000,
"peak_memory_mb_mean": 55000
}
}
]
@@ -138,10 +138,10 @@
"num-input-images": 2,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.10,
"latency_mean": 12,
"peak_memory_mb_max": 73000,
"peak_memory_mb_mean": 73000
"throughput_qps": 0.21,
"latency_mean": 10,
"peak_memory_mb_max": 57000,
"peak_memory_mb_mean": 57000
}
},
{
@@ -156,10 +156,10 @@
"num-input-images": 2,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.05,
"throughput_qps": 0.08,
"latency_mean": 20,
"peak_memory_mb_max": 81000,
"peak_memory_mb_mean": 81000
"peak_memory_mb_max": 66000,
"peak_memory_mb_mean": 66000
}
}
]
tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json
+24 -24
View File
@@ -21,10 +21,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.05,
"latency_mean": 15.0,
"peak_memory_mb_max": 72500,
"peak_memory_mb_mean": 72500
"throughput_qps": 0.1,
"latency_mean": 10,
"peak_memory_mb_max": 57000,
"peak_memory_mb_mean": 57000
}
},
{
@@ -38,10 +38,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.01,
"latency_mean": 65.6,
"peak_memory_mb_max": 80777,
"peak_memory_mb_mean": 80777
"throughput_qps": 0.02,
"latency_mean": 50,
"peak_memory_mb_max": 65000,
"peak_memory_mb_mean": 65000
}
}
]
@@ -72,10 +72,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.10,
"latency_mean": 7.2,
"peak_memory_mb_max": 68100,
"peak_memory_mb_mean": 68100
"throughput_qps": 0.21,
"latency_mean": 10,
"peak_memory_mb_max": 56000,
"peak_memory_mb_mean": 56000
}
},
{
@@ -89,10 +89,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.03,
"latency_mean": 24.0,
"peak_memory_mb_max": 68100,
"peak_memory_mb_mean": 68100
"throughput_qps": 0.06,
"latency_mean": 20,
"peak_memory_mb_max": 56000,
"peak_memory_mb_mean": 56000
}
}
]
@@ -133,10 +133,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.1,
"latency_mean": 6.5,
"peak_memory_mb_max": 72600,
"peak_memory_mb_mean": 72600
"throughput_qps": 0.26,
"latency_mean": 10,
"peak_memory_mb_max": 58000,
"peak_memory_mb_mean": 58000
}
},
{
@@ -150,10 +150,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.05,
"latency_mean": 16.0,
"peak_memory_mb_max": 81000,
"peak_memory_mb_mean": 81000
"throughput_qps": 0.1,
"latency_mean": 10,
"peak_memory_mb_max": 67000,
"peak_memory_mb_mean": 67000
}
}
]
tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json
+8 -8
View File
@@ -21,10 +21,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.02,
"latency_mean": 40.0,
"peak_memory_mb_max": 70000,
"peak_memory_mb_mean": 70000
"throughput_qps": 0.05,
"latency_mean": 20,
"peak_memory_mb_max": 63000,
"peak_memory_mb_mean": 63000
}
},
{
@@ -38,10 +38,10 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.005,
"latency_mean": 80.0,
"peak_memory_mb_max": 70000,
"peak_memory_mb_mean": 70000
"throughput_qps": 0.03,
"latency_mean": 30,
"peak_memory_mb_max": 63000,
"peak_memory_mb_mean": 63000
}
}
]
tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+21 -21
View File
@@ -21,9 +21,9 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.30,
"latency_mean": 3.50,
"peak_memory_mb_mean": 67000
"throughput_qps": 0.35,
"latency_mean": 3,
"peak_memory_mb_mean": 56000
}
},
{
@@ -37,9 +37,9 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.037,
"latency_mean": 27.0,
"peak_memory_mb_mean": 74000
"throughput_qps": 0.03,
"latency_mean": 26,
"peak_memory_mb_mean": 65000
}
}
]
@@ -67,9 +67,9 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.30,
"latency_mean": 3.50,
"peak_memory_mb_mean": 67000
"throughput_qps": 0.35,
"latency_mean": 3,
"peak_memory_mb_mean": 56000
}
},
{
@@ -83,9 +83,9 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.037,
"latency_mean": 27.0,
"peak_memory_mb_mean": 74000
"throughput_qps": 0.03,
"latency_mean": 26,
"peak_memory_mb_mean": 65000
}
}
]
@@ -116,9 +116,9 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.1,
"latency_mean": 9.1,
"peak_memory_mb_mean": 61000
"throughput_qps": 0.11,
"latency_mean": 9,
"peak_memory_mb_mean": 55000
}
}
]
@@ -159,9 +159,9 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.1,
"latency_mean": 3.50,
"peak_memory_mb_mean": 67000
"throughput_qps": 0.46,
"latency_mean": 3,
"peak_memory_mb_mean": 56000
}
},
{
@@ -175,9 +175,9 @@
"max-concurrency": 1,
"enable-negative-prompt": true,
"baseline": {
"throughput_qps": 0.1,
"latency_mean": 6.15,
"peak_memory_mb_mean": 74000
"throughput_qps": 0.18,
"latency_mean": 6,
"peak_memory_mb_mean": 67000
}
}
]
tests/dfx/perf/tests/test_tts.json
+197 -52
View File
@@ -4,7 +4,9 @@
"server_params": {
"model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
"stage_overrides": {
"1": {"devices": "1"}
"1": {
"devices": "1"
}
}
},
"benchmark_params": [
@@ -15,13 +17,21 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "linyueqian/seed-tts-eval-subset",
"num_prompts": [20],
"max_concurrency": [1],
"num_prompts": [
20
],
"max_concurrency": [
1
],
"seed_tts_locale": "en",
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [350],
"median_audio_rtf": [0.25]
"median_audio_ttfp_ms": [
3000
],
"median_audio_rtf": [
0.6
]
}
},
{
@@ -31,14 +41,34 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "linyueqian/seed-tts-eval-subset",
"num_prompts": [80, 128, 128],
"max_concurrency": [8, 16, 64],
"num_prompts": [
80,
128,
128
],
"max_concurrency": [
8,
16,
64
],
"seed_tts_locale": "en",
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [3500, 30000, 7000],
"median_audio_rtf": [0.75, 10.0, 1.5],
"audio_throughput": [10.0, 1.0, 40.0]
"median_audio_ttfp_ms": [
1000,
3000,
14000
],
"median_audio_rtf": [
0.70,
1.15,
3.90
],
"audio_throughput": [
12,
14,
14
]
}
},
{
@@ -49,13 +79,19 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "linyueqian/seed-tts-eval-subset",
"num_prompts": [200],
"max_concurrency": [4],
"num_prompts": [
200
],
"max_concurrency": [
4
],
"seed_tts_locale": "en",
"seed_tts_wer_eval": true,
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_audio_rtf": [0.45]
"mean_audio_rtf": [
0.45
]
}
}
]
@@ -65,7 +101,9 @@
"server_params": {
"model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
"stage_overrides": {
"1": {"devices": "1"}
"1": {
"devices": "1"
}
}
},
"benchmark_params": [
@@ -76,14 +114,26 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [20],
"max_concurrency": [1],
"num_prompts": [
20
],
"max_concurrency": [
1
],
"seed_tts_locale": "en",
"extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
"extra_body": {
"voice": "Vivian",
"language": "English",
"task_type": "CustomVoice"
},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [150],
"median_audio_rtf": [0.15]
"median_audio_ttfp_ms": [
2000
],
"median_audio_rtf": [
0.40
]
}
},
{
@@ -93,15 +143,39 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [80, 128, 128],
"max_concurrency": [8, 16, 64],
"num_prompts": [
80,
128,
128
],
"max_concurrency": [
8,
16,
64
],
"seed_tts_locale": "en",
"extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
"extra_body": {
"voice": "Vivian",
"language": "English",
"task_type": "CustomVoice"
},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [1500, 30000, 3000],
"median_audio_rtf": [0.30, 10.0, 0.60],
"audio_throughput": [30.0, 1.0, 100.0]
"median_audio_ttfp_ms": [
300,
1000,
5000
],
"median_audio_rtf": [
0.25,
0.35,
1.10
],
"audio_throughput": [
30,
40,
40
]
}
},
{
@@ -112,14 +186,24 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [200],
"max_concurrency": [4],
"num_prompts": [
200
],
"max_concurrency": [
4
],
"seed_tts_locale": "en",
"extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
"extra_body": {
"voice": "Vivian",
"language": "English",
"task_type": "CustomVoice"
},
"seed_tts_wer_eval": true,
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"mean_audio_rtf": [0.35]
"mean_audio_rtf": [
0.35
]
}
},
{
@@ -129,15 +213,29 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
"num_prompts": [100],
"request_rate": [2.0],
"num_prompts": [
100
],
"request_rate": [
2.0
],
"seed_tts_locale": "en",
"extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
"extra_body": {
"voice": "Vivian",
"language": "English",
"task_type": "CustomVoice"
},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [3000],
"median_audio_rtf": [0.25],
"audio_throughput": [4.0]
"median_audio_ttfp_ms": [
100
],
"median_audio_rtf": [
0.20
],
"audio_throughput": [
10
]
}
},
{
@@ -147,14 +245,25 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_design",
"num_prompts": [20],
"max_concurrency": [1],
"num_prompts": [
20
],
"max_concurrency": [
1
],
"seed_tts_locale": "en",
"extra_body": {"task_type": "VoiceDesign", "language": "English"},
"extra_body": {
"task_type": "VoiceDesign",
"language": "English"
},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [150],
"median_audio_rtf": [0.15]
"median_audio_ttfp_ms": [
100
],
"median_audio_rtf": [
0.15
]
}
},
{
@@ -164,15 +273,38 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_design",
"num_prompts": [80, 128, 128],
"max_concurrency": [8, 16, 64],
"num_prompts": [
80,
128,
128
],
"max_concurrency": [
8,
16,
64
],
"seed_tts_locale": "en",
"extra_body": {"task_type": "VoiceDesign", "language": "English"},
"extra_body": {
"task_type": "VoiceDesign",
"language": "English"
},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [1500, 30000, 3500],
"median_audio_rtf": [0.35, 10.0, 0.70],
"audio_throughput": [25.0, 1.0, 80.0]
"median_audio_ttfp_ms": [
100,
1000,
5000
],
"median_audio_rtf": [
0.25,
0.40,
1.20
],
"audio_throughput": [
30,
40,
40
]
}
},
{
@@ -182,15 +314,28 @@
"backend": "openai-audio-speech",
"endpoint": "/v1/audio/speech",
"dataset_path": "benchmarks/build_dataset/seed_tts_design",
"num_prompts": [100],
"request_rate": [2.0],
"num_prompts": [
100
],
"request_rate": [
2.0
],
"seed_tts_locale": "en",
"extra_body": {"task_type": "VoiceDesign", "language": "English"},
"extra_body": {
"task_type": "VoiceDesign",
"language": "English"
},
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
"baseline": {
"median_audio_ttfp_ms": [3500],
"median_audio_rtf": [0.30],
"audio_throughput": [4.0]
"median_audio_ttfp_ms": [
100
],
"median_audio_rtf": [
0.20
],
"audio_throughput": [
9
]
}
}
]
tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
+9 -9
View File
@@ -30,9 +30,9 @@
}
],
"baseline": {
"throughput_qps": 0.034,
"latency_mean": 26.0,
"peak_memory_mb_mean": 80000
"throughput_qps": 0.04,
"latency_mean": 30,
"peak_memory_mb_mean": 75000
}
}
]
@@ -72,9 +72,9 @@
}
],
"baseline": {
"throughput_qps": 0.042,
"latency_mean": 21.6,
"peak_memory_mb_mean": 55300
"throughput_qps": 0.06,
"latency_mean": 20,
"peak_memory_mb_mean": 48000
}
},
{
@@ -97,9 +97,9 @@
}
],
"baseline": {
"throughput_qps": 0.0085,
"latency_mean": 101.6,
"peak_memory_mb_mean": 65200
"throughput_qps": 0.01,
"latency_mean": 90,
"peak_memory_mb_mean": 56000
}
}
]
tests/dfx/reliability/invalid_param_test/conftest.py
+23
View File
@@ -0,0 +1,23 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared env for ``http_invalid`` (real :func:`omni_server`)."""
from __future__ import annotations
import io
import os
import pytest
from PIL import Image
# Match ``tests/e2e/online_serving/*`` module-level env for subprocess serve.
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
os.environ.setdefault("VLLM_TEST_CLEAN_GPU_MEMORY", "0")
@pytest.fixture
def tiny_png_bytes() -> bytes:
img = Image.new("RGB", (32, 32), color="gray")
buf = io.BytesIO()
img.save(buf, format="PNG")
return buf.getvalue()
tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py
+118
View File
@@ -0,0 +1,118 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""``POST /v1/audio/generate`` validation (live Stable Audio Open server)."""
from __future__ import annotations
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
# https://github.com/vllm-project/vllm-omni/issues/3649 — invalid inputs still HTTP 200 / timeout / missing validation.
_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
_PARAMS = [
pytest.param(
OmniServerParams(
model="stabilityai/stable-audio-open-1.0",
server_args=[
"--trust-remote-code",
"--enforce-eager",
"--model-class-name",
"StableAudioPipeline",
],
),
id="stable_audio_open",
marks=hardware_marks(res={"cuda": "L4"}),
),
]
@pytest.mark.parametrize(
"overrides, err_message",
[
pytest.param({"input": ""}, ("cannot be empty", "input"), id="input_empty", marks=_SKIP_ISSUE_3649),
pytest.param(
{"input": " "}, ("cannot be empty", "input"), id="input_whitespace_only", marks=_SKIP_ISSUE_3649
),
pytest.param({"input": 123}, ("input", "a valid string"), id="input_wrong_type"),
pytest.param(
{"response_format": "mpeg"},
("response_format", "literal_error", "Input should be"),
id="response_format_invalid",
),
pytest.param({"speed": 0.1}, ("speed", "greater_than_equal", "0.25"), id="speed_too_low"),
pytest.param({"speed": 99.0}, ("speed", "less_than_equal", "4"), id="speed_too_high"),
pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"),
pytest.param(
{"stream_format": "sse"},
("sse", "stream_format", "value_error", "not a supported stream_format"),
id="stream_format_sse_blocked",
),
pytest.param(
{"audio_length": -1.0}, ("audio_end_in_s", "higher than", "audio_start_in_s"), id="audio_length_negative"
),
pytest.param({"audio_length": 0}, "audio_length", id="audio_length_zero", marks=_SKIP_ISSUE_3649),
# ``OmniOpenAIServingAudioGenerate`` only forwards ``audio_start`` when ``audio_length`` is set.
pytest.param(
{"audio_start": -0.5, "audio_length": 5.0},
("audio_start_in_s", "between", "0", "512"),
id="audio_start_negative",
),
pytest.param(
{"negative_prompt": ["noise"]}, ("negative_prompt", "a valid string"), id="negative_prompt_wrong_type"
),
pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649),
pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649),
pytest.param(
{"num_inference_steps": 0}, "num_inference_steps", id="num_inference_steps_zero", marks=_SKIP_ISSUE_3649
),
pytest.param(
{"num_inference_steps": -1}, ("number of steps", "non-negative"), id="num_inference_steps_negative"
),
pytest.param(
{"num_inference_steps": 6000},
"num_inference_steps",
id="num_inference_steps_above_max",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
{"guidance_scale": 1001.0}, "guidance_scale", id="guidance_scale_above_max", marks=_SKIP_ISSUE_3649
),
pytest.param(
{"audio_length": 86401.0}, ("Requested audio length", "exceeds maximum"), id="audio_length_above_max"
),
pytest.param(
{"audio_start": 86401.0, "audio_length": 1.0},
("audio_start_in_s", "between", "0", "512"),
id="audio_start_above_max",
),
],
)
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_audio_generate_invalid_field_values(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
overrides: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
body = {"model": omni_server.model, "input": "ambient electronic pad"}
body.update(overrides)
openai_client.send_audio_generate_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_audio_generate_missing_input(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
openai_client.send_audio_generate_http_request(
{
"json": {"model": omni_server.model},
"timeout": 120,
"err_code": 400,
"err_message": ("missing", "input", "field required"),
}
)
tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py
+862
View File
@@ -0,0 +1,862 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""HTTP + WebSocket validation for Qwen3-TTS Base: ``/v1/audio/speech``, ``/v1/audio/speech/stream``, batch, voices."""
from __future__ import annotations
import io
import json
import uuid
import wave
from typing import Any
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.media import load_test_audio_data_url
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
from tests.helpers.stage_config import get_deploy_config_path
# Speech / voice-upload numeric caps for these tests only (not exported from vllm_omni).
# Align where possible with serving_speech defaults: _TTS_MAX_INSTRUCTIONS_LENGTH (500),
# _TTS_MAX_NEW_TOKENS_MAX (4096). Voice form max lengths are not public constants in prod;
# values below keep tests self-contained and deterministic.
_SPEECH_API_MAX_INSTRUCTIONS_CHARS = 500
_SPEECH_API_MAX_NEW_TOKENS = 4096
_VOICE_UPLOAD_MAX_CONSENT_LEN = 1024
_VOICE_UPLOAD_MAX_REF_TEXT_CHARS = 8192
_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS = 2048
REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
# Qwen3-TTS 0.6B expects 1024-dim speaker embeddings (see serving_speech / talker hidden_size).
_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM = 1024
# Pop-only marker for ``test_speech_invalid_field_values``: omit ``ref_audio`` / ``ref_text`` so Base
# ``speaker_embedding`` cases mirror standalone embedding-only JSON (no ref_audio clash).
_SPEECH_INVALID_EMBEDDING_ONLY_BODY = "__speech_invalid_embedding_only_body__"
# Batch-only: merged items inherit batch ``ref_audio`` / ``ref_text`` because ``_pick`` treats item ``None``
# as "unset". Strip clone fields from the batch envelope before asserting CustomVoice preset errors.
_SPEECH_BATCH_NO_CLONE_FIELDS = "__speech_batch_no_clone_fields__"
_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
pytestmark = [pytest.mark.slow, pytest.mark.tts]
_L4_SPEECH_HW = hardware_marks(res={"cuda": "L4"})
_SPEECH_SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
_QWEN3_TTS_SPEECH = [
pytest.param(
OmniServerParams(
model="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
stage_config_path=get_deploy_config_path("qwen3_tts.yaml"),
server_args=_SPEECH_SERVER_ARGS,
),
id="qwen3_tts",
marks=_L4_SPEECH_HW,
),
]
def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None:
if loc == "batch":
body.update(overrides)
return
items = body["items"]
assert isinstance(items, list) and len(items) >= 1
first = items[0]
assert isinstance(first, dict)
body["items"] = [{**first, **overrides}]
def _pcm_wav_mono_bytes(*, duration_s: float = 1.1, sample_rate: int = 44100) -> bytes:
buf = io.BytesIO()
with wave.open(buf, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2)
wf.setframerate(sample_rate)
nframes = max(1, int(sample_rate * duration_s))
wf.writeframes(b"\x00\x00" * nframes)
return buf.getvalue()
# ─────────────────────────────────────────────────────────────────────────────
# POST /v1/audio/speech (JSON)
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
openai_client.send_audio_speech_http_request(
{
"raw_body": "{",
"timeout": 120,
"err_code": 400,
"err_message": ("json", "decode", "invalid"),
}
)
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
openai_client.send_audio_speech_http_request(
{
"json": {"model": omni_server.model},
"timeout": 120,
"err_code": 400,
"err_message": ("input", "field required", "missing"),
}
)
@pytest.mark.parametrize(
"overrides, err_message",
[
pytest.param({"input": ""}, ("input", "empty"), id="input_empty"),
pytest.param({"input": " "}, ("input", "empty"), id="input_whitespace_only"),
pytest.param({"voice": ""}, "Invalid voice", id="voice_empty", marks=_SKIP_ISSUE_3649),
pytest.param(
{"instructions": 123}, ("instructions", "string_type", "valid string"), id="instructions_wrong_type"
),
pytest.param(
{"response_format": "mpeg"}, ("response_format", "literal_error", "wav"), id="response_format_invalid"
),
pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"),
pytest.param(
{"stream_format": "sse"}, ("stream_format", "value_error", "audio"), id="stream_format_sse_blocked"
),
pytest.param({"stream": "wrong_type"}, ("stream", "bool_parsing", "validation error"), id="stream_wrong_type"),
pytest.param(
{
"voice": "vivian",
"task_type": "CustomVoice",
"ref_audio": None,
"ref_text": None,
},
("CustomVoice", "does not support"),
id="customvoice_task_not_supported",
),
pytest.param(
{"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
),
pytest.param({"language": ""}, ("language", "invalid language", "chinese"), id="language_empty"),
pytest.param({"ref_audio": "ftp://example.com/a.wav"}, ("ref_audio", "URL"), id="ref_audio_bad_scheme"),
pytest.param(
{"ref_audio": "not_a_valid_uri"},
("ref_audio", "url"),
id="ref_audio_invalid_uri",
),
pytest.param({"ref_text": ["x"]}, ("ref_text", "string_type", "valid string"), id="ref_text_wrong_type"),
pytest.param(
{"x_vector_only_mode": "wrong"},
("x_vector_only_mode", "valid boolean"),
id="x_vector_only_mode_string_not_boolean",
),
pytest.param(
{"speaker_embedding": [0.01] * _QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM},
("mutually exclusive", "speaker_embedding", "ref_audio"),
id="speaker_embedding_with_ref_audio_mutually_exclusive",
),
pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
pytest.param(
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
("max_new_tokens", "exceed 4096"),
id="max_new_tokens_above_max",
),
pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
pytest.param(
{"initial_codec_chunk_frames": -1},
("initial_codec_chunk_frames", "greater_than_equal", "0"),
id="initial_codec_chunk_frames_negative",
),
pytest.param(
{"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
("instructions", "too long"),
id="instructions_exceed_api_limit",
),
# Base clone: whitespace-only ref_text fails serving validation (not pydantic).
pytest.param(
{
"task_type": "Base",
"ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
"ref_text": " ",
},
("ref_text", "base task", "non-empty"),
id="ref_text_whitespace_only_base_clone",
),
pytest.param(
{
_SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
"task_type": "Base",
"speaker_embedding": [],
},
("speaker_embedding", "non-empty"),
id="base_speaker_embedding_empty",
),
pytest.param(
{
_SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
"task_type": "Base",
"x_vector_only_mode": True,
"speaker_embedding": [0.01] * 512,
},
("speaker_embedding", "dimensions", str(_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM)),
id="base_speaker_embedding_wrong_dim_short",
),
pytest.param(
{
_SPEECH_INVALID_EMBEDDING_ONLY_BODY: True,
"task_type": "Base",
"x_vector_only_mode": True,
"speaker_embedding": [0.01] * 2048,
},
("speaker_embedding", "2048", "dimensions", str(_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM)),
id="base_speaker_embedding_wrong_dim_long",
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_invalid_field_values(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
overrides: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
ov = dict(overrides)
embedding_only = bool(ov.pop(_SPEECH_INVALID_EMBEDDING_ONLY_BODY, False))
body = {
"model": omni_server.model,
"input": "Hello.",
"voice": "clone",
"language": "English",
"response_format": "wav",
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
if embedding_only:
body.pop("ref_audio", None)
body.pop("ref_text", None)
body.update(ov)
openai_client.send_audio_speech_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
# ─────────────────────────────────────────────────────────────────────────────
# POST /v1/audio/speech/batch (JSON)
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_batch_empty_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
openai_client.send_audio_speech_batch_http_request(
{"json": {"items": []}, "timeout": 120, "err_code": (400, 422), "err_message": ("items", "least 1 item")}
)
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
openai_client.send_audio_speech_batch_http_request(
{
"json": {"model": omni_server.model},
"timeout": 120,
"err_code": (400, 422),
"err_message": ("items", "field required", "missing"),
}
)
# Align field-level checks with ``POST /v1/audio/speech`` where ``BatchSpeechRequest`` /
# ``SpeechBatchItem`` expose the same knobs. Not on batch schema (single-speech only):
# ``stream``, ``stream_format``, ``seed``, ``speaker_embedding`` — see dedicated speech tests.
@pytest.mark.parametrize(
"loc, overrides, err_message",
[
pytest.param("item", {"input": ""}, ("input", "empty"), id="item_input_empty", marks=_SKIP_ISSUE_3649),
pytest.param(
"item",
{"input": " "},
("input", "empty"),
id="item_input_whitespace_only",
marks=_SKIP_ISSUE_3649,
),
pytest.param("batch", {"voice": ""}, "Invalid voice", id="batch_voice_empty", marks=_SKIP_ISSUE_3649),
pytest.param("item", {"voice": ""}, "Invalid voice", id="item_voice_empty", marks=_SKIP_ISSUE_3649),
pytest.param(
"batch",
{
"voice": "nonexistent_voice_xyz",
"task_type": "CustomVoice",
"ref_audio": None,
"ref_text": None,
},
"Invalid voice",
id="batch_voice_unknown_preset",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{
_SPEECH_BATCH_NO_CLONE_FIELDS: True,
"voice": "nonexistent_voice_xyz",
"task_type": "CustomVoice",
},
"Invalid voice",
id="item_voice_unknown_preset",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"batch",
{"instructions": 123},
("instructions", "string_type", "valid string"),
id="batch_instructions_wrong_type",
),
pytest.param(
"item",
{"instructions": 123},
("instructions", "string_type", "valid string"),
id="item_instructions_wrong_type",
),
pytest.param(
"batch",
{"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
("instructions", "too long"),
id="batch_instructions_exceed_api_limit",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{"instructions": "x" * (_SPEECH_API_MAX_INSTRUCTIONS_CHARS + 1)},
("instructions", "too long"),
id="item_instructions_exceed_api_limit",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"batch",
{"response_format": "mpeg"},
("response_format", "literal_error", "wav"),
id="batch_response_format_invalid",
),
pytest.param(
"item",
{"response_format": "mpeg"},
("response_format", "literal_error", "wav"),
id="item_response_format_invalid",
),
pytest.param("batch", {"speed": 0.1}, ("speed", "greater_than_equal", "0.25"), id="batch_speed_too_low"),
pytest.param("item", {"speed": 10.0}, ("speed", "less_than_equal", "4"), id="item_speed_too_high"),
pytest.param(
"batch",
{"task_type": "InvalidEnum"},
("task_type", "literal_error", "CustomVoice"),
id="batch_task_type_invalid",
),
pytest.param(
"item",
{"task_type": "InvalidEnum"},
("task_type", "literal_error", "CustomVoice"),
id="item_task_type_invalid",
),
pytest.param(
"batch",
{"language": ""},
("language", "invalid language", "chinese"),
id="batch_language_empty",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{"language": ""},
("language", "invalid language", "chinese"),
id="item_language_empty",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"batch",
{"ref_audio": "ftp://example.com/a.wav"},
("ref_audio", "URL"),
id="batch_ref_audio_bad_scheme",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{"ref_audio": "ftp://example.com/a.wav"},
("ref_audio", "URL"),
id="item_ref_audio_bad_scheme",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"batch",
{"ref_audio": "not_a_valid_uri"},
("ref_audio", "url"),
id="batch_ref_audio_invalid_uri",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{"ref_audio": "not_a_valid_uri"},
("ref_audio", "url"),
id="item_ref_audio_invalid_uri",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"batch",
{"ref_text": ["x"]},
("ref_text", "string_type", "valid string"),
id="batch_ref_text_wrong_type",
),
pytest.param(
"item",
{"ref_text": ["x"]},
("ref_text", "string_type", "valid string"),
id="item_ref_text_wrong_type",
),
pytest.param(
"batch",
{"x_vector_only_mode": "wrong_type"},
("x_vector_only_mode", "bool_parsing", "validation error"),
id="batch_x_vector_only_mode_wrong_type",
),
pytest.param(
"item",
{"x_vector_only_mode": "wrong_type"},
("x_vector_only_mode", "bool_parsing", "validation error"),
id="item_x_vector_only_mode_wrong_type",
),
pytest.param(
"batch",
{"max_new_tokens": 0},
("max_new_tokens", "least 1"),
id="batch_max_new_tokens_below_min",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{"max_new_tokens": 0},
("max_new_tokens", "least 1"),
id="item_max_new_tokens_below_min",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"batch",
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
("max_new_tokens", "exceed 4096"),
id="batch_max_new_tokens_above_max",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
("max_new_tokens", "exceed 4096"),
id="item_max_new_tokens_above_max",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"batch",
{"initial_codec_chunk_frames": -1},
("initial_codec_chunk_frames", "greater_than_equal", "0"),
id="batch_initial_codec_chunk_frames_negative",
),
pytest.param(
"item",
{"initial_codec_chunk_frames": -1},
("initial_codec_chunk_frames", "greater_than_equal", "0"),
id="item_initial_codec_chunk_frames_negative",
),
# Base clone: whitespace-only ref_text (same serving check as single speech).
pytest.param(
"batch",
{
"task_type": "Base",
"ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
"ref_text": " ",
},
("ref_text", "base task", "non-empty"),
id="batch_ref_text_whitespace_only_base_clone",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"item",
{
"task_type": "Base",
"ref_audio": "data:audio/wav;base64,UklGRiQAAABXQVZFZm10IBAAAAABAAEARKwAAIhYAQACABAAZGF0YQAAAAA=",
"ref_text": " ",
},
("ref_text", "base task", "non-empty"),
id="item_ref_text_whitespace_only_base_clone",
marks=_SKIP_ISSUE_3649,
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_batch_invalid_field_values(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
loc: str,
overrides: dict[str, object],
err_message: str | tuple[str, ...],
) -> None:
ov = dict(overrides)
strip_clone = bool(ov.pop(_SPEECH_BATCH_NO_CLONE_FIELDS, False))
body = {
"model": omni_server.model,
"voice": "clone",
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
"language": "English",
"response_format": "wav",
"items": [{"input": "Hello."}],
}
if strip_clone:
body.pop("ref_audio", None)
body.pop("ref_text", None)
_apply_batch_overrides(body, loc=loc, overrides=ov)
openai_client.send_audio_speech_batch_http_request(
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
)
# ─────────────────────────────────────────────────────────────────────────────
# WebSocket /v1/audio/speech/stream
# ─────────────────────────────────────────────────────────────────────────────
# ``test_speech_stream_invalid_requests``: build ``session.config`` JSON with ``omni_server.model``.
_SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES = object()
@pytest.mark.parametrize(
"send_frames_spec, err_message",
[
pytest.param("{not-json", "Invalid JSON", id="invalid_json"),
pytest.param(
json.dumps({"type": "input.text", "text": "hello"}),
"Expected session.config",
id="wrong_first_message_type",
),
pytest.param(
_SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES,
("invalid session config", "response_format", "pcm"),
id="session_config_stream_audio_requires_pcm",
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_speech_stream_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
send_frames_spec: Any,
err_message: str | tuple[str, ...],
) -> None:
"""Bad first WebSocket text frames: non-JSON, wrong ``type``, or invalid ``session.config``."""
if send_frames_spec is _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES:
send_frames = json.dumps(
{
"type": "session.config",
"model": omni_server.model,
"response_format": "wav",
"stream_audio": True,
}
)
else:
send_frames = send_frames_spec
assert isinstance(send_frames, str)
openai_client.send_audio_speech_stream_ws_request(
{
"send_frames": send_frames,
"timeout": 120,
"ws_max_size": None,
"ws_json_type": "error",
"err_message": err_message,
}
)
# ─────────────────────────────────────────────────────────────────────────────
# POST /v1/audio/voices (multipart + embedding-only) and DELETE /v1/audio/voices/{name}
# ─────────────────────────────────────────────────────────────────────────────
# Templates for ``test_voices_create_invalid_requests``: resolved in ``_finalize_voices_form_data``.
_VF_LONG_CONSENT = object()
_VF_LONG_REF_TEXT = object()
_VF_LONG_SPEAKER_DESC = object()
_VF_EMBEDDING_NAN_JSON = object()
def _finalize_voices_form_data(template: dict[str, Any], uuid_hex: str) -> dict[str, Any]:
out: dict[str, Any] = {}
for key, val in template.items():
if val is _VF_LONG_CONSENT:
out[key] = "x" * (_VOICE_UPLOAD_MAX_CONSENT_LEN + 1)
elif val is _VF_LONG_REF_TEXT:
out[key] = "y" * (_VOICE_UPLOAD_MAX_REF_TEXT_CHARS + 1)
elif val is _VF_LONG_SPEAKER_DESC:
out[key] = "z" * (_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS + 1)
elif val is _VF_EMBEDDING_NAN_JSON:
out[key] = json.dumps([0.1] * 1023 + [float("nan")])
elif isinstance(val, str):
out[key] = val.replace("{uuid}", uuid_hex)
else:
out[key] = val
return out
def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
"""``None`` => no multipart file (embedding-only or neither)."""
if kind is None:
return None
if kind == "wav_ok":
return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(), "audio/wav")}
if kind == "wav_short":
return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(duration_s=0.3), "audio/wav")}
if kind == "wav_bad_body":
return {"audio_sample": ("clip.wav", b"not valid wav content", "audio/wav")}
if kind == "wav_pdf_type":
return {"audio_sample": ("clip.wav", _pcm_wav_mono_bytes(), "application/pdf")}
raise AssertionError(f"unknown multipart kind {kind!r}")
@pytest.mark.parametrize(
"multipart_kind, form_template, err_message",
[
pytest.param(
None,
{"consent": "c1", "name": "v1"},
("audio_sample", "speaker_embedding", "provided"),
id="neither_audio_nor_embedding",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
{"name": "v_miss_consent_{uuid}"},
("missing", "consent", "field required"),
id="missing_consent",
),
pytest.param(
"wav_ok",
{"consent": "consent_ok"},
("missing", "name", "field required"),
id="missing_name",
),
pytest.param(
"wav_ok",
{"consent": "", "name": "v_empty_consent_{uuid}"},
("missing", "consent", "field required"),
id="empty_consent",
),
pytest.param(
"wav_ok",
{"consent": " ", "name": "v_ws_consent_{uuid}"},
"consent",
id="whitespace_consent",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
{"consent": "bad/consent", "name": "v_bad_cons_{uuid}"},
"consent",
id="consent_path_sep",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
{"consent": _VF_LONG_CONSENT, "name": "v_long_cons_{uuid}"},
("consent", "too long", "Failed to save"),
id="consent_too_long",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
{"consent": "consent_ok", "name": ""},
("missing", "name", "field required"),
id="empty_name",
),
pytest.param(
"wav_ok",
{"consent": "consent_ok", "name": "evil/name"},
("voice name", "invalid voice name"),
id="invalid_name_path_sep",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
{"consent": "consent_ok", "name": "v_long_ref_{uuid}", "ref_text": _VF_LONG_REF_TEXT},
"ref_text",
id="ref_text_too_long",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
{
"consent": "consent_ok",
"name": "v_long_desc_{uuid}",
"speaker_description": _VF_LONG_SPEAKER_DESC,
},
"speaker_description",
id="speaker_description_too_long",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_pdf_type",
{"consent": "consent_ok", "name": "v_bad_mime_{uuid}"},
("mime", "unsupported", "audio/mp4"),
id="audio_unsupported_mime",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_short",
{"consent": "consent_ok", "name": "v_short_{uuid}"},
("too short", "reference audio"),
id="audio_too_short",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_bad_body",
{"consent": "consent_ok", "name": "v_bad_wav_{uuid}"},
("decode", "audio", "Format not recognised"),
id="audio_decode_error",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
None,
{
"speaker_embedding": "not valid json [[[",
"consent": "consent_emb",
"name": "v_bad_json_{uuid}",
},
("speaker_embedding", "valid JSON"),
id="speaker_embedding_invalid_json",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
None,
{
"speaker_embedding": json.dumps([]),
"consent": "consent_emb",
"name": "v_empty_emb_{uuid}",
},
("speaker_embedding", "non-empty"),
id="speaker_embedding_empty",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
None,
{
"speaker_embedding": _VF_EMBEDDING_NAN_JSON,
"consent": "consent_emb",
"name": "v_nan_emb_{uuid}",
},
("speaker_embedding", "finite"),
id="speaker_embedding_nan",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
None,
{
"speaker_embedding": json.dumps([0.01] * 512),
"consent": "consent_emb",
"name": "v_wrong_dim_{uuid}",
},
("expected", "dimensions", "1024"),
id="speaker_embedding_wrong_dim",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"wav_ok",
{
"consent": "consent_emb",
"name": "v_both_{uuid}",
"speaker_embedding": json.dumps([0.01] * 1024),
},
("mutually exclusive", "audio_sample", "speaker_embedding"),
id="audio_and_embedding_mutually_exclusive",
marks=_SKIP_ISSUE_3649,
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_voices_create_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
multipart_kind: str | None,
form_template: dict[str, Any],
err_message: str | tuple[str, ...],
) -> None:
"""Invalid ``POST /v1/audio/voices`` bodies (multipart or embedding-only)."""
uid = uuid.uuid4().hex[:8]
cfg: dict[str, Any] = {
"data": _finalize_voices_form_data(form_template, uid),
"timeout": 120,
"err_code": 400,
"err_message": err_message,
}
files = _voices_upload_multipart_files(multipart_kind)
if files is not None:
cfg["files"] = files
openai_client.send_audio_voices_create_http_request(cfg)
@pytest.mark.parametrize(
"voice_name, err_code, err_message",
[
pytest.param(
"missing-voice-404",
404,
"not found",
id="not_found",
),
pytest.param(
" ",
404,
"not found",
id="whitespace_only",
),
pytest.param(
"evil/name",
404,
"not found",
id="path_separator",
),
pytest.param(
".",
405,
("Method Not Allowed", "detail"),
id="dot_single",
),
pytest.param(
"..",
404,
"not found",
id="dot_dot",
),
pytest.param(
"bad\x00voice",
404,
"not found",
id="nul_byte",
),
pytest.param(
"",
405,
"Method Not Allowed",
id="empty_name_path",
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
def test_voices_delete_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
voice_name: str,
err_code: int | tuple[int, ...],
err_message: str | tuple[str, ...],
) -> None:
"""DELETE ``/v1/audio/voices/{name}``: missing voice (404), odd segments (often 404), ``.``paths (405)."""
openai_client.send_audio_voices_delete_http_request(
{
"name": voice_name,
"timeout": 120,
"err_code": err_code,
"err_message": err_message,
}
)
tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py
+90
View File
@@ -0,0 +1,90 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""``POST /v1/images/edits`` multipart validation (live ``Qwen/Qwen-Image-Edit``)."""
from __future__ import annotations
from typing import Any
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
_PARAMS = [
pytest.param(
OmniServerParams(model="Qwen/Qwen-Image-Edit"),
id="qwen_image_edit",
marks=hardware_marks(res={"cuda": "H100"}),
),
]
# Populated in ``test_images_edits_invalid_requests`` from ``omni_server.model``.
_IMAGE_EDITS_SERVER_MODEL = object()
def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> dict[str, Any]:
return {k: (server_model if v is _IMAGE_EDITS_SERVER_MODEL else v) for k, v in template.items()}
@pytest.mark.parametrize(
"include_image, data_template, err_message, err_code",
[
pytest.param(
False,
{"prompt": "make it brighter", "model": _IMAGE_EDITS_SERVER_MODEL},
("image", "required", "Unprocessable Entity"),
422,
id="missing_image",
),
pytest.param(
True,
{"prompt": "edit", "model": "wrong-model-id"},
("model", "mismatch"),
400,
id="model_mismatch",
),
pytest.param(
True,
{"prompt": "edit", "response_format": "url"},
("response_format", "b64_json", "supported"),
400,
id="unsupported_response_format",
),
pytest.param(
True,
{"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "output_compression": "101"},
("less_than_equal", "output_compression", "100"),
400,
id="output_compression_above_max",
),
pytest.param(
True,
{"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"},
("Invalid layers", "3", "1"),
400,
id="invalid_layers",
),
],
)
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_images_edits_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
tiny_png_bytes: bytes,
include_image: bool,
data_template: dict[str, Any],
err_message: str | tuple[str, ...],
err_code: int | tuple[int, ...],
) -> None:
cfg: dict[str, Any] = {
"data": _finalize_edits_form_data(data_template, omni_server.model),
"timeout": 300,
"err_code": err_code,
"err_message": err_message,
}
if include_image:
cfg["files"] = {"image": ("x.png", tiny_png_bytes, "image/png")}
openai_client.send_images_edits_http_request(cfg)
tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py
+142
View File
@@ -0,0 +1,142 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""``POST /v1/images/generations`` validation (live ``Qwen/Qwen-Image``)."""
from __future__ import annotations
from typing import Any
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
_PARAMS = [
pytest.param(
OmniServerParams(model="Qwen/Qwen-Image"),
id="qwen_image",
marks=hardware_marks(res={"cuda": "H100"}),
),
]
# Full replacement body for ``test_images_generations_invalid_requests`` (no minimal prompt/model merge).
_IMAGES_GEN_MISSING_PROMPT_JSON = object()
def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
return {"prompt": "a simple red apple icon", "model": omni_server.model}
# ─────────────────────────────────────────────────────────────────────────────
# POST /v1/images/generations (JSON)
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"body_spec, err_message",
[
pytest.param(_IMAGES_GEN_MISSING_PROMPT_JSON, ("prompt", "Field required", "Missing"), id="missing_prompt"),
pytest.param(
{"prompt": "cat", "model": "not-the-loaded-checkpoint"},
("model", "mismatch"),
id="model_mismatch",
),
pytest.param({"prompt": ""}, ("prompt", "empty"), id="prompt_empty", marks=_SKIP_ISSUE_3649),
pytest.param({"prompt": " "}, ("prompt", "empty"), id="prompt_whitespace_only", marks=_SKIP_ISSUE_3649),
pytest.param({"prompt": 123}, ("prompt", "string_type", "valid string"), id="prompt_wrong_type"),
pytest.param({"n": 0}, ("n", "greater_than_equal", "1"), id="n_below_min"),
pytest.param({"n": 11}, ("n", "less_than_equal", "10"), id="n_above_max"),
pytest.param({"size": "1024"}, ("size", "value_error", "WIDTHxHEIGHT"), id="size_missing_height_token"),
pytest.param({"size": "not-a-size"}, ("size", "value_error", "WIDTHxHEIGHT"), id="size_missing_separator"),
pytest.param({"size": "1024xabc"}, ("Invalid size format", "integers"), id="size_non_integer_height"),
pytest.param({"size": "1024x"}, ("Invalid size format", "integers"), id="size_incomplete_height"),
pytest.param({"size": "0x1024"}, ("Invalid size", "positive integers"), id="size_non_positive_width"),
pytest.param(
{"response_format": "url"},
("response_format", "value_error", "b64_json"),
id="response_format_not_b64_json",
),
pytest.param(
{"use_system_prompt": "bogus"},
("use_system_prompt", "value_error", "dynamic"),
id="use_system_prompt_invalid",
),
pytest.param(
{"num_inference_steps": 0},
("num_inference_steps", "'greater_than_equal", "1"),
id="num_inference_steps_below_min",
),
pytest.param(
{"num_inference_steps": 201},
("num_inference_steps", "less_than_equal", "200"),
id="num_inference_steps_above_max",
),
pytest.param(
{"guidance_scale": -1.0}, ("guidance_scale", "greater_than_equal", "0"), id="guidance_scale_negative"
),
pytest.param(
{"guidance_scale": 21.0}, ("guidance_scale", "less_than_equal", "20"), id="guidance_scale_above_max"
),
pytest.param(
{"true_cfg_scale": -1.0}, ("true_cfg_scale", "greater_than_equal", "0"), id="true_cfg_scale_negative"
),
pytest.param(
{"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max"
),
pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"),
pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"),
pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649),
pytest.param(
{"seed": 2**32},
("seed", "less_than_equal", "4294967295"),
id="seed_above_uint32",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
{"output_format": "gif"},
("output_format", "value_error", "b64_json"),
id="output_format_invalid",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
{"vae_use_slicing": "wrong_type"},
("vae_use_slicing", "bool_parsing", "validation error"),
id="vae_use_slicing_wrong_type",
),
pytest.param({"user": 123}, ("user", "string_type", "valid string"), id="user_wrong_type"),
pytest.param(
{"negative_prompt": ["noise"]},
("negative_prompt", "string_type", "valid string"),
id="negative_prompt_wrong_type",
),
pytest.param(
{"generator_device": 1},
("generator_device", "string_type", "valid string"),
id="generator_device_wrong_type",
),
pytest.param(
{"lora": {"foo": "bar"}}, ("lora", "both name and path", "required"), id="lora_missing_required_fields"
),
],
)
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
def test_images_generations_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
body_spec: Any,
err_message: str | tuple[str, ...],
) -> None:
"""Invalid JSON bodies for ``POST /v1/images/generations`` (missing prompt, model mismatch, bad fields)."""
if body_spec is _IMAGES_GEN_MISSING_PROMPT_JSON:
body: dict[str, object] = {"size": "1024x1024"}
else:
assert isinstance(body_spec, dict)
body = _minimal_images_gen_json(omni_server)
body.update(body_spec)
openai_client.send_images_generations_http_request(
{"json": body, "timeout": 300, "err_code": 400, "err_message": err_message}
)
tests/dfx/reliability/invalid_param_test/test_invalid_omni_chat.py
+216
View File
@@ -0,0 +1,216 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
from __future__ import annotations
import json
from typing import Any
import pytest
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
from tests.helpers.stage_config import get_deploy_config_path
pytestmark = [pytest.mark.slow, pytest.mark.omni]
_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
def _minimal_chat_json(omni_server: OmniServer) -> dict[str, object]:
"""Minimal valid chat body; individual tests override one offending field."""
return {
"model": omni_server.model,
"messages": [{"role": "user", "content": "hello"}],
"stream": False,
}
_QWEN3_OMNI_SERVER = [
pytest.param(
OmniServerParams(
model="Qwen/Qwen3-Omni-30B-A3B-Instruct",
stage_config_path=get_deploy_config_path("qwen3_omni_moe.yaml"),
server_args=["--async-chunk"],
),
id="qwen3_omni",
),
]
def _chat_completions_request_without_expectations(omni_server: OmniServer, case_id: str) -> dict[str, Any]:
"""Build ``send_chat_completions_http_request`` kwargs excluding ``err_code`` / ``err_message``."""
if case_id == "malformed_json":
return {"raw_body": "{not-json", "timeout": 120}
if case_id == "missing_messages":
return {"json": {"model": omni_server.model, "stream": False}, "timeout": 120}
if case_id == "model_mismatch":
return {
"json": {
"model": "this-model-is-not-served-on-this-instance",
"messages": [{"role": "user", "content": "hello"}],
"stream": False,
},
"timeout": 120,
}
body = _minimal_chat_json(omni_server)
if case_id == "stream_wrong_type":
body["stream"] = "not-a-boolean" # type: ignore[assignment]
elif case_id == "temperature_wrong_type":
body["temperature"] = "hot" # type: ignore[assignment]
elif case_id == "max_tokens_wrong_type":
body["max_tokens"] = "not-an-int" # type: ignore[assignment] # type: ignore[assignment]
elif case_id == "modalities_list_bad":
body["modalities"] = [123] # type: ignore[assignment]
elif case_id == "response_format_json_schema_incomplete":
body["response_format"] = {"type": "json_schema"}
elif case_id == "logprobs_wrong_type":
body["logprobs"] = "yes" # type: ignore[assignment]
elif case_id == "logprobs_top_without_enabled":
body["logprobs"] = False
body["top_logprobs"] = 5
elif case_id == "speaker_unknown":
body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
else:
raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
return {"json": body, "timeout": 120}
# ─────────────────────────────────────────────────────────────────────────────
# POST /v1/chat/completions
# ─────────────────────────────────────────────────────────────────────────────
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize(
"case_id, err_code, err_message",
[
pytest.param(
"malformed_json",
400,
("json_invalid", "decode error"),
id="malformed_json",
),
pytest.param("missing_messages", 400, ("messages", "missing", "Field required"), id="missing_messages"),
pytest.param(
"model_mismatch",
404,
("model", "does not exist", "NotFoundError"),
id="model_mismatch",
),
pytest.param("stream_wrong_type", 400, ("stream", "bool_parsing", "not-a-boolean"), id="invalid_stream_type"),
pytest.param(
"temperature_wrong_type",
400,
("temperature", "float_parsing", "a valid number"),
id="invalid_temperature_type",
),
pytest.param(
"max_tokens_wrong_type", 400, ("max_tokens", "int_parsing", "integer"), id="invalid_max_tokens_type"
),
pytest.param(
"modalities_list_bad",
400,
("modalities", "value_error", ""),
id="modalities_list_bad_element",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"response_format_json_schema_incomplete",
400,
("response_format", "value_error", "json_schema"),
id="invalid_response_format_json_schema",
),
pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
pytest.param(
"logprobs_top_without_enabled",
400,
("top_logprobs", "logprobs", "True"),
id="logprobs_top_without_logprobs_true",
),
pytest.param(
"speaker_unknown",
400,
("Invalid speaker", "Supported"),
id="speaker_unknown_preset",
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
def test_chat_completions_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
case_id: str,
err_code: int | tuple[int, ...],
err_message: str | tuple[str, ...],
) -> None:
cfg = _chat_completions_request_without_expectations(omni_server, case_id)
cfg["err_code"] = err_code
cfg["err_message"] = err_message
openai_client.send_chat_completions_http_request(cfg)[0]
# ─────────────────────────────────────────────────────────────────────────────
# WS /v1/video/chat/stream
# ─────────────────────────────────────────────────────────────────────────────
# First-frame JSON with ``omni_server.model`` for ``test_video_chat_stream_invalid_requests``.
_VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR = object()
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize(
"send_frames_spec, err_message",
[
pytest.param("{not-json", "Invalid JSON", id="invalid_json"),
pytest.param(
_VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR,
"Invalid session config",
id="session_modalities_scalar",
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
def test_video_chat_stream_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
send_frames_spec: Any,
err_message: str,
) -> None:
if send_frames_spec is _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR:
send_frames = json.dumps({"type": "session.config", "model": omni_server.model, "modalities": "text"})
else:
send_frames = send_frames_spec
assert isinstance(send_frames, str)
openai_client.send_video_chat_stream_ws_request(
{
"send_frames": send_frames,
"timeout": 120,
"ws_max_size": None,
"ws_json_type": "error",
"err_message": err_message,
}
)
# ─────────────────────────────────────────────────────────────────────────────
# WS /v1/realtime
# ─────────────────────────────────────────────────────────────────────────────
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
def test_realtime_rejected_when_async_chunk_enabled(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
) -> None:
openai_client.send_realtime_ws_request(
{
"timeout": 120,
"ws_max_size": None,
"ws_json_type": "error",
"ws_error_code": "unsupported",
"err_message": "async_chunk",
}
)
tests/dfx/reliability/invalid_param_test/test_invalid_server_control.py
+95
View File
@@ -0,0 +1,95 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Omni sleep/wakeup JSON validation (live ``Qwen/Qwen-Image`` process)."""
from __future__ import annotations
from typing import Any, Literal
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni]
_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649")
_QWEN_IMAGE = [
pytest.param(
OmniServerParams(model="Qwen/Qwen-Image"),
id="qwen_image",
marks=hardware_marks(res={"cuda": "H100"}),
),
]
# ─────────────────────────────────────────────────────────────────────────────
# POST omni sleep / omni wakeup (invalid JSON)
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"endpoint, json_body, err_message",
[
pytest.param("sleep", {"level": 2}, ("stage_ids", "Field required", "Missing"), id="sleep_missing_stage_ids"),
pytest.param(
"sleep",
{"stage_ids": [], "level": 2},
("stage_ids", "Field required", "Missing"),
id="sleep_empty_stage_ids",
marks=_SKIP_ISSUE_3649,
),
pytest.param(
"sleep",
{"stage_ids": "not-a-list", "level": 2},
("stage_ids", "list_type", "a valid list"),
id="sleep_invalid_stage_ids_type",
),
pytest.param(
"sleep",
{"stage_ids": [0], "level": "high"},
("level", "int_parsing", "a valid integer"),
id="sleep_invalid_level_type",
),
pytest.param(
"sleep",
{"stage_ids": [0], "level": -1},
("level", "greater"),
id="sleep_negative_level",
marks=_SKIP_ISSUE_3649,
),
pytest.param("wakeup", {}, ("stage_ids", "missing", "Field required"), id="wakeup_missing_stage_ids"),
pytest.param(
"wakeup",
{"stage_ids": None},
("stage_ids", "list_type", "a valid list"),
id="wakeup_invalid_stage_ids_type",
),
pytest.param(
"wakeup",
{"stage_ids": []},
("stage_ids", "length"),
id="wakeup_empty_stage_ids",
marks=_SKIP_ISSUE_3649,
),
],
)
@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
def test_omni_sleep_wakeup_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
endpoint: Literal["sleep", "wakeup"],
json_body: dict[str, Any],
err_message: str | tuple[str, ...],
) -> None:
cfg: dict[str, Any] = {
"json": json_body,
"timeout": 120,
"err_code": 400,
"err_message": err_message,
}
if endpoint == "sleep":
openai_client.send_omni_sleep_http_request(cfg)
else:
openai_client.send_omni_wakeup_http_request(cfg)
tests/dfx/reliability/invalid_param_test/test_invalid_video_generation.py
+171
View File
@@ -0,0 +1,171 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""HTTP video job routes (live ``Wan-AI/Wan2.2-T2V-A14B-Diffusers``).
WebSocket ``/v1/video/chat/stream`` and ``/v1/realtime`` invalid-input checks live in
``test_invalid_omni_chat.py`` with the same Qwen3-Omni server module fixture (Wan cannot serve those paths).
"""
from __future__ import annotations
from typing import Any
import pytest
from tests.helpers.mark import hardware_marks
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
_WAN_T2V = [
pytest.param(
OmniServerParams(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers"),
id="wan22_t2v",
marks=hardware_marks(res={"cuda": "H100"}),
),
]
def _minimal_video_form_data() -> dict[str, str]:
return {"prompt": "a cat walking on grass"}
# ─────────────────────────────────────────────────────────────────────────────
# POST /v1/videos (async)
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"form_extra, err_message",
[
pytest.param(None, ("prompt", "Field required", "Missing"), id="missing_prompt"),
pytest.param({"seconds": "0"}, ("seconds", "string_pattern_mismatch", "1"), id="seconds_zero"),
pytest.param({"size": "bad"}, ("size", "string_pattern_mismatch"), id="size_bad"),
pytest.param(
{"num_inference_steps": "0"},
("num_inference_steps", "greater_than_equal", "1"),
id="num_inference_steps_zero",
),
pytest.param(
{"guidance_scale": "25"}, ("guidance_scale", "less_than_equal", "20"), id="guidance_scale_above_max"
),
],
)
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_post_videos_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
form_extra: dict[str, str] | None,
err_message: str | tuple[str, ...],
) -> None:
data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
openai_client.send_videos_create_http_request(
{"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
)
# ─────────────────────────────────────────────────────────────────────────────
# POST /v1/videos/sync
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"form_extra, err_message",
[
pytest.param(None, ("prompt", "Field required", "Missing"), id="missing_prompt"),
pytest.param({"seconds": "0"}, ("seconds", "string_pattern_mismatch", "1"), id="seconds_zero"),
pytest.param({"size": "bad"}, ("size", "string_pattern_mismatch"), id="size_bad"),
pytest.param(
{"num_inference_steps": "201"},
("num_inference_steps", "less_than_equal", "200"),
id="num_inference_steps_above_max",
),
],
)
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_post_videos_sync_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
form_extra: dict[str, str] | None,
err_message: str | tuple[str, ...],
) -> None:
data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
openai_client.send_videos_sync_http_request(
{"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
)
# ─────────────────────────────────────────────────────────────────────────────
# GET /v1/videos (list)
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"params, err_message",
[
pytest.param({"limit": -1}, ("limit", "greater_than_equal", "0"), id="invalid_limit_negative"),
pytest.param({"limit": 101}, ("limit", "less_than_equal", "100"), id="limit_above_max"),
pytest.param({"order": "newest"}, ("order", "literal_error", "asc", "desc"), id="invalid_order"),
],
)
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_get_videos_list_invalid_requests(
omni_server: OmniServer,
openai_client: OpenAIClientHandler,
params: dict[str, Any],
err_message: str | tuple[str, ...],
) -> None:
openai_client.send_videos_list_http_request(
{"params": params, "timeout": 120, "err_code": 400, "err_message": err_message}
)
# ─────────────────────────────────────────────────────────────────────────────
# GET /v1/videos/{video_id}
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
openai_client.send_video_retrieve_http_request(
{
"video_id": "does-not-exist",
"timeout": 120,
"err_code": 404,
"err_message": ("not found", "video"),
}
)
# ─────────────────────────────────────────────────────────────────────────────
# DELETE /v1/videos/{video_id}
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
openai_client.send_video_delete_http_request(
{
"video_id": "does-not-exist",
"timeout": 120,
"err_code": 404,
"err_message": ("not found", "video"),
}
)
# ─────────────────────────────────────────────────────────────────────────────
# GET /v1/videos/{video_id}/content
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
def test_get_video_content_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
openai_client.send_video_content_http_request(
{
"video_id": "does-not-exist",
"timeout": 120,
"err_code": 404,
"err_message": ("not found", "video"),
}
)
tests/diffusion/attention/test_attention_config.py
+99
View File
@@ -420,3 +420,102 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
assert attn.use_ring is True
assert attn.ring_runner is not None
assert attn.ring_runner.attn_backend_pref == "TORCH_SDPA"
class TestDiffusionKvCacheQuantization:
@staticmethod
def _install_attention_init_stubs(monkeypatch):
class _FakeAttentionImpl:
def __init__(self, **kwargs):
self.kwargs = kwargs
def forward(self, query, key, value, attn_metadata=None):
return query
class _FakeBackend:
@staticmethod
def get_name() -> str:
return "FAKE_BACKEND"
@staticmethod
def get_impl_cls():
return _FakeAttentionImpl
@staticmethod
def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
return True
class _FakeRingParallelAttention:
def __init__(self, sp_group, attn_backend_pref=None):
self.sp_group = sp_group
self.attn_backend_pref = attn_backend_pref
monkeypatch.setattr(
layer_mod,
"get_attn_backend_for_role",
lambda role, head_size, attention_config=None, role_category=None: (_FakeBackend, None),
)
monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
monkeypatch.setattr(layer_mod, "get_sp_group", lambda: SimpleNamespace(ring_group="ring-group"))
monkeypatch.setattr(layer_mod, "RingParallelAttention", _FakeRingParallelAttention)
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
def test_diffusion_kv_cache_dtype_none_does_not_trigger_ring_quantization_error(self, monkeypatch):
self._install_attention_init_stubs(monkeypatch)
od_config = SimpleNamespace(
diffusion_attention_config=AttentionConfig(),
parallel_config=SimpleNamespace(ring_degree=2),
diffusion_kv_cache_dtype=None,
diffusion_kv_cache_skip_step_indices=None,
diffusion_kv_cache_skip_layer_indices=None,
)
with set_current_diffusion_config(od_config):
attn = Attention(
num_heads=4,
head_size=64,
causal=False,
softmax_scale=1.0,
)
assert attn._kv_cache_dtype is None
def test_diffusion_kv_cache_dtype_auto_does_not_trigger_ring_quantization_error(self, monkeypatch):
self._install_attention_init_stubs(monkeypatch)
od_config = SimpleNamespace(
diffusion_attention_config=AttentionConfig(),
parallel_config=SimpleNamespace(ring_degree=2),
diffusion_kv_cache_dtype="auto",
diffusion_kv_cache_skip_step_indices=None,
diffusion_kv_cache_skip_layer_indices=None,
)
with set_current_diffusion_config(od_config):
attn = Attention(
num_heads=4,
head_size=64,
causal=False,
softmax_scale=1.0,
)
assert attn._kv_cache_dtype is None
def test_diffusion_kv_cache_dtype_fp8_raises_with_ring_attention(self, monkeypatch):
self._install_attention_init_stubs(monkeypatch)
od_config = SimpleNamespace(
diffusion_attention_config=AttentionConfig(),
parallel_config=SimpleNamespace(ring_degree=2),
diffusion_kv_cache_dtype="fp8",
diffusion_kv_cache_skip_step_indices=None,
diffusion_kv_cache_skip_layer_indices=None,
)
with set_current_diffusion_config(od_config):
with pytest.raises(ValueError, match="KV quantization is not compatible with ring attention"):
Attention(
num_heads=4,
head_size=64,
causal=False,
softmax_scale=1.0,
)
tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py
+79
View File
@@ -0,0 +1,79 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import (
DistributedAutoencoderKLHunyuan,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _DummyDistributedAutoencoderKLHunyuan(DistributedAutoencoderKLHunyuan):
def __init__(self):
torch.nn.Module.__init__(self)
self.tile_latent_min_size = 8
self.tile_sample_min_size = 8
self.tile_overlap_factor = 0.25
self.use_spatial_tiling = False
@property
def dtype(self) -> torch.dtype:
return torch.float32
def decoder(self, x: torch.Tensor) -> torch.Tensor:
return x + 10
def encoder(self, x: torch.Tensor) -> torch.Tensor:
return x + 20
def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
return b
def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
return b
def test_hunyuan_vae_use_tiling_aliases_spatial_tiling():
# Verify use_tiling property maps to use_spatial_tiling.
vae = _DummyDistributedAutoencoderKLHunyuan()
assert not vae.use_tiling
vae.use_tiling = True
assert vae.use_spatial_tiling
def test_hunyuan_vae_decode_tiles_round_trip():
# Validate decode tile split/exec/merge returns expected reconstructed tensor.
vae = _DummyDistributedAutoencoderKLHunyuan()
z = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12)
tile_tasks, grid_spec = vae.tile_split(z)
decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks}
output = vae.tile_merge(decoded_tiles, grid_spec)
assert grid_spec.split_dims == (3, 4)
assert grid_spec.grid_shape == (2, 2)
assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
assert len(tile_tasks) == 4
assert torch.equal(output, z + 10)
def test_hunyuan_vae_encode_tiles_round_trip():
# Validate encode tile split/exec/merge returns expected latent tensor.
vae = _DummyDistributedAutoencoderKLHunyuan()
x = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12)
tile_tasks, grid_spec = vae.encode_tile_split(x)
encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks}
output = vae.encode_tile_merge(encoded_tiles, grid_spec)
assert grid_spec.split_dims == (3, 4)
assert grid_spec.grid_shape == (2, 2)
assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
assert len(tile_tasks) == 4
assert torch.equal(output, x + 20)
tests/diffusion/distributed/test_pipeline_parallel.py
+802
View File
@@ -0,0 +1,802 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
import socket
import pytest
import torch
from vllm.model_executor.models.utils import PPMissingLayer, make_empty_intermediate_tensors_factory, make_layers
from vllm.sequence import IntermediateTensors
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
destroy_distributed_env,
get_classifier_free_guidance_rank,
get_pp_group,
init_distributed_environment,
initialize_model_parallel,
)
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.parallel]
def _find_free_port() -> str:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return str(s.getsockname()[1])
def update_environment_variables(envs_dict: dict[str, str]) -> None:
for k, v in envs_dict.items():
os.environ[k] = v
# ---------------------------------------------------------------------------
# Shared stubs used by both unit and distributed tests
# ---------------------------------------------------------------------------
class FakeWork:
"""Drop-in for torch.distributed.Work that records whether wait() was called."""
def __init__(self):
self.waited = False
def wait(self):
self.waited = True
class SimpleScheduler:
"""Minimal diffusion-step scheduler: latents -= 0.1 * noise_pred."""
def step(self, noise_pred: torch.Tensor, t, latents: torch.Tensor, return_dict: bool = False):
return (latents - 0.1 * noise_pred,)
class FakeVAE:
def __init__(self, distributed_enabled: bool = False):
self.calls = 0
self.distributed_enabled = distributed_enabled
def decode(self, z: torch.Tensor):
"""Original decode docstring."""
self.calls += 1
return (z + 1,)
def is_distributed_enabled(self) -> bool:
return self.distributed_enabled
class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
"""Minimal pipeline used to exercise PipelineParallelMixin.
Uses vLLM's ``make_layers`` for layer partitioning — the same utility used
by real DiT models — so the PP layer-split logic is exercised faithfully.
Each layer's weights are seeded by ``seed + layer_index`` so that layer ``i``
is initialized identically on every rank regardless of which ranks are active,
allowing the distributed output to be compared against the single-GPU baseline.
Args:
num_layers: Total number of Linear layers.
dim: Input / hidden dimension.
seed: Base RNG seed; layer ``i`` uses ``seed + i``.
device: Target device for layer weights (default: CPU).
dtype: Target dtype for layer weights (default: float32).
"""
def __init__(
self,
num_layers: int = 4,
dim: int = 64,
seed: int = 42,
device: torch.device | None = None,
dtype: torch.dtype = torch.float32,
):
self.start_layer, self.end_layer, self.layers = make_layers(
num_layers,
lambda prefix: torch.nn.Linear(dim, dim, bias=False),
prefix="layers",
)
for i, layer in enumerate(self.layers):
if not isinstance(layer, PPMissingLayer):
torch.manual_seed(seed + i)
torch.nn.init.normal_(layer.weight, mean=0.0, std=0.02)
self.layers.to(device=device, dtype=dtype)
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], dim)
self.scheduler = SimpleScheduler()
def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.Tensor | IntermediateTensors:
"""Layer-split forward pass.
* First PP rank: uses ``x`` from caller kwargs.
* Later PP ranks: overrides ``x`` with ``intermediate_tensors["hidden_states"]``
(which transparently waits for the async receive).
* Non-last PP ranks return ``IntermediateTensors``; the last rank
returns the plain noise-prediction tensor.
"""
if intermediate_tensors is not None:
x = intermediate_tensors["hidden_states"]
for i in range(self.start_layer, self.end_layer):
x = self.layers[i](x)
pp_group = get_pp_group()
if not pp_group.is_last_rank:
return IntermediateTensors({"hidden_states": x})
return x
# ---------------------------------------------------------------------------
# 1. AsyncLatents – unit tests (no distributed env required)
# ---------------------------------------------------------------------------
class TestAsyncLatents:
"""Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
pytestmark = [pytest.mark.cpu]
def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
def test_resolve_returns_wrapped_tensor(self):
t = torch.randn(2, 4)
al = self._make(t)
assert al._resolve() is t
def test_attribute_access_resolves(self):
t = torch.randn(2, 4)
al = self._make(t)
assert al.shape == t.shape
assert al.dtype == t.dtype
def test_torch_function_protocol(self):
"""torch ops that receive an AsyncLatents should see the underlying tensor."""
t = torch.randn(2, 4)
al = self._make(t)
mask = torch.ones_like(t)
result = mask * al # triggers __torch_function__
torch.testing.assert_close(result, mask * t)
def test_torch_function_with_list_arg(self):
"""__torch_function__ must unwrap AsyncLatents inside list/tuple args."""
t = torch.randn(2, 4)
al = self._make(t)
result = torch.cat([al, al], dim=0)
torch.testing.assert_close(result, torch.cat([t, t], dim=0))
def test_torch_tensor_conversion(self):
"""torch.as_tensor on an AsyncLatents must share storage with the underlying tensor (no copy)."""
t = torch.randn(2, 4)
al = self._make(t)
result = torch.as_tensor(al)
assert result.data_ptr() == t.data_ptr(), "torch.as_tensor copied the data instead of sharing storage"
def test_handles_are_waited_before_resolve(self):
t = torch.randn(2, 4)
h1, h2 = FakeWork(), FakeWork()
al = self._make(t, handles=[h1, h2])
_ = al.shape # trigger resolution
assert h1.waited and h2.waited, "Not all handles were waited on"
def test_postproc_callbacks_invoked_on_resolve(self):
t = torch.randn(2, 4)
log: list[int] = []
al = self._make(t, postproc=[lambda: log.append(1), lambda: log.append(2)])
_ = al.shape
assert log == [1, 2], f"postproc not called in order: {log}"
def test_idempotent_resolve(self):
"""handle.wait() must not be called twice if _resolve() is called twice."""
t = torch.randn(2, 4)
h = FakeWork()
al = self._make(t, handles=[h])
_ = al.shape # first resolve
h.waited = False # reset sentinel
_ = al.dtype # second resolve
assert not h.waited, "handle.wait() was called a second time"
# ---------------------------------------------------------------------------
# 2. _sync_pp_send / diffuse wrapper – unit tests (no distributed env required)
# ---------------------------------------------------------------------------
class TestSyncPPSend:
"""Verifies PipelineParallelMixin's internal PP-send flush."""
pytestmark = [pytest.mark.cpu]
@staticmethod
def _make_pipeline() -> PipelineParallelMixin:
# Instantiate a bare mixin — no layers, no distributed env needed.
# _sync_pp_send only touches _pp_send_work, so this is sufficient.
class _BarePP(PipelineParallelMixin, CFGParallelMixin):
pass
return _BarePP()
def test_noop_when_work_list_empty(self):
pipeline = self._make_pipeline()
pipeline._sync_pp_send()
assert pipeline._pp_send_work == []
def test_waits_all_pending_handles(self):
pipeline = self._make_pipeline()
works = [FakeWork(), FakeWork(), FakeWork()]
pipeline._pp_send_work = works
pipeline._sync_pp_send()
assert all(w.waited for w in works), "Some handles were not waited on"
def test_clears_work_list_after_sync(self):
pipeline = self._make_pipeline()
pipeline._pp_send_work = [FakeWork()]
pipeline._sync_pp_send()
assert pipeline._pp_send_work == []
class TestDiffuseWrapper:
"""Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
pytestmark = [pytest.mark.cpu]
def test_diffuse_flushes_pending_sends_on_success(self):
work = FakeWork()
class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
def diffuse(self):
self._pp_send_work = [work]
return "done"
pipeline = _DiffusePP()
assert pipeline.diffuse() == "done"
assert work.waited
assert pipeline._pp_send_work == []
def test_diffuse_flushes_pending_sends_on_exception(self):
work = FakeWork()
class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
def diffuse(self):
self._pp_send_work = [work]
raise RuntimeError("boom")
pipeline = _DiffusePP()
with pytest.raises(RuntimeError, match="boom"):
pipeline.diffuse()
assert work.waited
assert pipeline._pp_send_work == []
def test_diffuse_wrapper_preserves_metadata(self):
class _DiffusePP(PipelineParallelMixin, CFGParallelMixin):
def diffuse(self):
"""Original diffuse docstring."""
return "done"
assert _DiffusePP.diffuse.__name__ == "diffuse"
assert _DiffusePP.diffuse.__doc__ == "Original diffuse docstring."
class TestVaeDecodeGuard:
pytestmark = [pytest.mark.cpu]
@staticmethod
def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
class _DecodePP(PipelineParallelMixin, CFGParallelMixin):
def __init__(self):
self.vae = FakeVAE(distributed_enabled=distributed_enabled)
return _DecodePP()
@staticmethod
def _set_rank(monkeypatch, world_size: int, first_stage: bool) -> None:
monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: world_size)
monkeypatch.setattr(pp_module, "is_pipeline_first_stage", lambda: first_stage)
def test_calls_original_decode_when_pp_disabled(self, monkeypatch):
self._set_rank(monkeypatch, world_size=1, first_stage=True)
pipeline = self._make_pipeline()
z = torch.ones(2, 3)
output = pipeline.vae.decode(z)[0]
assert pipeline.vae.calls == 1
torch.testing.assert_close(output, z + 1)
def test_calls_original_decode_on_first_stage(self, monkeypatch):
self._set_rank(monkeypatch, world_size=2, first_stage=True)
pipeline = self._make_pipeline()
z = torch.ones(2, 3)
output = pipeline.vae.decode(z)[0]
assert pipeline.vae.calls == 1
torch.testing.assert_close(output, z + 1)
def test_skips_decode_on_non_first_stage(self, monkeypatch):
self._set_rank(monkeypatch, world_size=2, first_stage=False)
pipeline = self._make_pipeline()
z = torch.ones(2, 3)
output = pipeline.vae.decode(z)
assert pipeline.vae.calls == 0
assert output == (None,)
def test_calls_original_decode_when_distributed_vae_enabled(self, monkeypatch):
self._set_rank(monkeypatch, world_size=2, first_stage=False)
pipeline = self._make_pipeline(distributed_enabled=True)
z = torch.ones(2, 3)
output = pipeline.vae.decode(z)[0]
assert pipeline.vae.calls == 1
torch.testing.assert_close(output, z + 1)
def test_decode_wrapper_preserves_metadata(self):
pipeline = self._make_pipeline()
assert pipeline.vae.decode.__name__ == "decode"
assert pipeline.vae.decode.__doc__ == "Original decode docstring."
@pytest.mark.cpu
def test_pipeline_parallel_requires_cfg_mixin():
with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
class _MissingCFG(PipelineParallelMixin):
pass
@pytest.mark.cpu
def test_pipeline_parallel_requires_mro_before_cfg_mixin():
with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
class _WrongOrder(CFGParallelMixin, PipelineParallelMixin):
pass
# ---------------------------------------------------------------------------
# Distributed test helpers
# ---------------------------------------------------------------------------
def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
"""Initialise the distributed environment for a spawned worker."""
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
current_omni_platform.set_device(device)
update_environment_variables(
{
"RANK": str(local_rank),
"LOCAL_RANK": str(local_rank),
"WORLD_SIZE": str(world_size),
"MASTER_ADDR": "localhost",
"MASTER_PORT": master_port,
}
)
init_distributed_environment()
return device
def make_pipeline_and_inputs(
test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
) -> tuple["MockPipelineParallel", dict, dict | None]:
"""Create a MockPipelineParallel and seeded inputs from a test_config dict.
Must be called after ``initialize_model_parallel`` so that ``make_layers``
can read the PP group to determine this rank's layer slice.
Returns ``(pipeline, positive_kwargs, negative_kwargs)``.
``negative_kwargs`` is ``None`` when ``do_true_cfg=False``.
"""
pipeline = MockPipelineParallel(
num_layers=test_config["num_layers"],
dim=test_config["dim"],
seed=test_config["model_seed"],
device=device,
dtype=dtype,
)
torch.manual_seed(test_config["input_seed"])
if torch.cuda.is_available():
torch.cuda.manual_seed_all(test_config["input_seed"])
pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
negative_kwargs = None
if do_true_cfg:
torch.manual_seed(test_config["input_seed"] + 1)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
negative_kwargs = {"x": neg_x}
return pipeline, pos_x, negative_kwargs
# ---------------------------------------------------------------------------
# 3. isend_tensor_dict / irecv_tensor_dict (2 GPUs)
# ---------------------------------------------------------------------------
def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
device = init_dist(local_rank, world_size, master_port)
initialize_model_parallel(pipeline_parallel_size=world_size)
pp_group = get_pp_group()
if pp_group.is_first_rank:
torch.manual_seed(77)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(77)
tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
handles = pp_group.isend_tensor_dict({"t": tensor})
for h in handles:
h.wait()
result_queue.put(("sent", tensor.cpu()))
else:
received = AsyncIntermediateTensors(*pp_group.irecv_tensor_dict())
result_queue.put(("received", received["t"].cpu()))
if torch.distributed.is_initialized():
torch.distributed.barrier()
destroy_distributed_env()
@pytest.mark.gpu
@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
@pytest.mark.parametrize("pp_size", [2])
def test_isend_irecv_tensor_dict(pp_size: int):
"""isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
q = manager.Queue()
port = _find_free_port()
torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
results = {label: tensor for label, tensor in [q.get(), q.get()]}
torch.testing.assert_close(
results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
)
# ---------------------------------------------------------------------------
# 4. predict_noise_maybe_with_cfg
# ---------------------------------------------------------------------------
_baseline_cache: dict[tuple, torch.Tensor] = {}
def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
"""Compute expected single-GPU output using the same MockPipelineParallel.
Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
Results are cached so identical configs are only computed once.
"""
key = (
test_config["num_layers"],
test_config["dim"],
test_config["batch_size"],
test_config["model_seed"],
test_config["input_seed"],
test_config["cfg_scale"],
dtype,
do_true_cfg,
)
if key in _baseline_cache:
return _baseline_cache[key]
device = init_dist(0, 1, _find_free_port())
initialize_model_parallel(pipeline_parallel_size=1)
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
test_config, dtype, device, do_true_cfg=do_true_cfg
)
with torch.inference_mode():
noise_pred = pipeline.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=test_config["cfg_scale"],
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
destroy_distributed_env()
_baseline_cache[key] = noise_pred.cpu()
return _baseline_cache[key]
def predict_noise_worker(
local_rank: int,
world_size: int,
master_port: str,
pp_size: int,
cfg_size: int,
do_true_cfg: bool,
dtype: torch.dtype,
test_config: dict,
result_queue,
):
"""Generic predict-noise worker parameterized by PP and CFG topology."""
device = init_dist(local_rank, world_size, master_port)
initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
pp_group = get_pp_group()
cfg_rank = get_classifier_free_guidance_rank()
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
test_config, dtype, device, do_true_cfg=do_true_cfg
)
with torch.inference_mode():
noise_pred = pipeline.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=test_config["cfg_scale"],
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
# This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
# Flush the non-last PP rank's async send before barrier / process teardown.
pipeline._sync_pp_send()
if pp_group.is_last_rank:
assert noise_pred is not None
if cfg_rank == 0:
result_queue.put(noise_pred.cpu())
else:
assert noise_pred is None
if torch.distributed.is_initialized():
torch.distributed.barrier()
destroy_distributed_env()
@pytest.mark.gpu
@pytest.mark.parametrize(
"pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
[
pytest.param(
2,
1,
False,
torch.float32,
4,
100,
1e-5,
1e-5,
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
id="pp2-no_cfg-float32",
),
pytest.param(
2,
1,
False,
torch.bfloat16,
4,
100,
1e-2,
1e-2,
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
id="pp2-no_cfg-bfloat16",
),
pytest.param(
2,
1,
True,
torch.bfloat16,
4,
100,
1e-2,
1e-2,
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
id="pp2-seq_cfg-bfloat16",
),
pytest.param(
2,
2,
True,
torch.bfloat16,
4,
100,
1e-2,
1e-2,
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
id="pp2-cfg2-bfloat16",
),
pytest.param(
3,
1,
False,
torch.bfloat16,
6,
100,
1e-2,
1e-2,
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
id="pp3-no_cfg-bfloat16",
),
],
)
def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
"""predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
test_config = {
"num_layers": num_layers,
"dim": 64,
"batch_size": 2,
"cfg_scale": 7.5,
"model_seed": 42,
"input_seed": input_seed,
}
baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
pp_q = manager.Queue()
world_size = pp_size * cfg_size
port = _find_free_port()
torch.multiprocessing.spawn(
predict_noise_worker,
args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
nprocs=world_size,
)
pp_out = pp_q.get()
assert baseline_out.shape == pp_out.shape
torch.testing.assert_close(
pp_out,
baseline_out,
rtol=rtol,
atol=atol,
msg=f"PP={pp_size} cfg={cfg_size} {'with' if do_true_cfg else 'no'} CFG output differs from baseline ({dtype=})",
)
# ---------------------------------------------------------------------------
# 5. scheduler_step_maybe_with_cfg
# ---------------------------------------------------------------------------
def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
"""Single-GPU reference: predict_noise + scheduler_step."""
device = init_dist(0, 1, _find_free_port())
initialize_model_parallel(pipeline_parallel_size=1)
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
test_config, torch.float32, device, do_true_cfg=do_true_cfg
)
latents = positive_kwargs["x"]
t = torch.tensor(500, device=device)
with torch.inference_mode():
noise_pred = pipeline.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=test_config["cfg_scale"],
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
result = pipeline.scheduler_step_maybe_with_cfg(
noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
)
destroy_distributed_env()
return result.cpu()
def scheduler_step_worker(
local_rank: int,
world_size: int,
master_port: str,
pp_size: int,
cfg_size: int,
do_true_cfg: bool,
test_config: dict,
result_queue,
):
device = init_dist(local_rank, world_size, master_port)
initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
pp_group = get_pp_group()
cfg_rank = get_classifier_free_guidance_rank()
pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
test_config, torch.float32, device, do_true_cfg=do_true_cfg
)
latents = positive_kwargs["x"]
t = torch.tensor(500, device=device)
with torch.inference_mode():
noise_pred = pipeline.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=test_config["cfg_scale"],
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
latents = pipeline.scheduler_step_maybe_with_cfg(
noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
)
# This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
# Flush the last PP rank's async latent send before barrier / process teardown.
pipeline._sync_pp_send()
if pp_group.is_first_rank and cfg_rank == 0:
resolved = latents.contiguous()
result_queue.put(resolved.cpu())
if torch.distributed.is_initialized():
torch.distributed.barrier()
destroy_distributed_env()
@pytest.mark.gpu
@pytest.mark.parametrize(
"pp_size, cfg_size, do_true_cfg, input_seed",
[
pytest.param(
2,
1,
False,
300,
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
id="pp2-no_cfg",
),
pytest.param(
2,
2,
True,
600,
marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
id="pp2-cfg2-true_cfg",
),
],
)
def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
"""Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
test_config = {
"num_layers": 4,
"dim": 64,
"batch_size": 2,
"cfg_scale": 7.5,
"model_seed": 42,
"input_seed": input_seed,
}
baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
mp_context = torch.multiprocessing.get_context("spawn")
manager = mp_context.Manager()
q = manager.Queue()
port = _find_free_port()
world_size = pp_size * cfg_size
torch.multiprocessing.spawn(
scheduler_step_worker,
args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
nprocs=world_size,
)
result = q.get()
torch.testing.assert_close(
result,
baseline,
rtol=0,
atol=0,
msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
)
tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py
+119
View File
@@ -0,0 +1,119 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock
import pytest
from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import (
Flux2KleinPipeline,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
class _StopAfterCheckInputsError(Exception):
pass
def _make_pipeline():
pipeline = object.__new__(Flux2KleinPipeline)
pipeline.vae_scale_factor = 8
pipeline.is_distilled = True
pipeline._guidance_scale = 0.0
return pipeline
def _make_minimal_request(
prompt="valid prompt",
*,
num_inference_steps=None,
):
"""Build a minimal OmniDiffusionRequest-like object that forward() reads from."""
params = OmniDiffusionSamplingParams(
height=512,
width=512,
num_inference_steps=num_inference_steps,
seed=42,
)
req = MagicMock()
req.sampling_params = params
req.prompts = [prompt]
req.multi_modal_data = {}
return req
def _capture_num_inference_steps(pipe, steps):
"""Run pipe.forward() and capture the num_inference_steps passed to check_inputs."""
captured = {}
def fake_check_inputs(**kwargs):
captured["num_inference_steps"] = kwargs.get("num_inference_steps")
raise _StopAfterCheckInputsError
original = pipe.check_inputs
pipe.check_inputs = fake_check_inputs
try:
req = _make_minimal_request(num_inference_steps=steps)
pipe.forward(req)
except _StopAfterCheckInputsError:
pass
finally:
pipe.check_inputs = original
return captured.get("num_inference_steps")
# --- Forward resolution tests (the actual regression) ---
def test_forward_preserves_zero_not_default():
"""#3703: num_inference_steps=0 must reach validation as 0, not as default 50."""
pipe = _make_pipeline()
captured = _capture_num_inference_steps(pipe, 0)
assert captured == 0, f"Expected 0, got {captured}"
def test_forward_preserves_negative():
pipe = _make_pipeline()
captured = _capture_num_inference_steps(pipe, -1)
assert captured == -1
def test_forward_none_uses_default():
"""None means unset — forward() should substitute the pipeline default."""
pipe = _make_pipeline()
captured = _capture_num_inference_steps(pipe, None)
assert captured == 50
def test_forward_preserves_positive():
pipe = _make_pipeline()
captured = _capture_num_inference_steps(pipe, 9)
assert captured == 9
# --- check_inputs validation tests ---
@pytest.mark.parametrize("steps", [0, -1])
def test_check_inputs_rejects_non_positive(steps):
pipe = _make_pipeline()
with pytest.raises(ValueError):
pipe.check_inputs(
prompt="valid prompt",
height=512,
width=512,
num_inference_steps=steps,
)
def test_check_inputs_accepts_none():
pipe = _make_pipeline()
pipe.check_inputs(prompt="valid prompt", height=512, width=512, num_inference_steps=None)
@pytest.mark.parametrize("steps", [1, 2])
def test_check_inputs_accepts_positive(steps):
pipe = _make_pipeline()
pipe.check_inputs(prompt="valid prompt", height=512, width=512, num_inference_steps=steps)
tests/diffusion/models/flux2/test_flux2_klein_validation.py
+80
View File
@@ -0,0 +1,80 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import (
Flux2KleinPipeline,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _make_pipeline():
pipeline = object.__new__(Flux2KleinPipeline)
pipeline.vae_scale_factor = 8
pipeline.is_distilled = True
pipeline._guidance_scale = 0.0
return pipeline
def _check(pipeline, prompt, prompt_embeds=None):
pipeline.check_inputs(
prompt=prompt,
height=512,
width=512,
prompt_embeds=prompt_embeds,
num_inference_steps=4,
guidance_scale=0.0,
)
@pytest.mark.parametrize(
"prompt",
[
"",
" ",
],
)
def test_rejects_empty_or_whitespace_prompt(prompt):
pipe = _make_pipeline()
with pytest.raises(ValueError, match="empty or whitespace-only"):
_check(pipe, prompt)
@pytest.mark.parametrize(
"prompt",
[
["valid prompt", ""],
["valid prompt", " "],
[" "],
],
)
def test_rejects_list_with_empty_or_whitespace_element(prompt):
pipe = _make_pipeline()
with pytest.raises(ValueError, match="empty or whitespace-only"):
_check(pipe, prompt)
@pytest.mark.parametrize(
"prompt",
[
"valid prompt",
["valid prompt", "another valid prompt"],
],
)
def test_accepts_valid_prompts(prompt):
pipe = _make_pipeline()
_check(pipe, prompt)
def test_rejects_none_without_embeds():
pipe = _make_pipeline()
with pytest.raises(ValueError):
_check(pipe, None)
def test_accepts_none_with_prompt_embeds():
pipe = _make_pipeline()
_check(pipe, None, prompt_embeds=torch.randn(1, 256, 4096))
tests/diffusion/models/wan2_2/test_create_transformer_quant.py
+383
View File
@@ -0,0 +1,383 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Regression tests for transformer quant-config auto-detection.
The loader path at pipeline_wan2_2.py carries two quantization contracts:
create_transformer_from_config (~L137)
- Reads quantization_config from config.json (injected by the merge scripts)
- Auto-detects the quant method when no CLI quant_config is provided
- Rejects method mismatches (CLI vs disk)
- Upgrades online → offline when disk marks is_checkpoint_*_serialized=True
- Rebuilds when the active ignored_layers differs from the disk value
Wan22Pipeline._create_transformer (~L456)
- Passes od_config.quantization_config (set by CLI or externally) to
create_transformer_from_config for each transformer independently.
- When od_config.quantization_config is None, each transformer auto-detects
from its own config.json; od_config is NOT modified by this call.
All tests are pure-CPU and do not load model weights.
"""
from __future__ import annotations
from types import SimpleNamespace
import pytest
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import create_transformer_from_config
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
# Minimum config that create_transformer_from_config accepts without raising.
_MIN_CFG: dict = {"patch_size": [1, 2, 2]}
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_fake_transformer():
"""Return (FakeTransformer class, captured list).
Each _FakeTransformer.__init__ call appends its **kwargs to captured,
letting tests inspect what quant_config was passed per transformer.
"""
captured: list[dict] = []
class _FakeTransformer:
def __init__(self, **kwargs):
captured.append(kwargs)
return _FakeTransformer, captured
class _FakePipeline:
"""Minimal stand-in exposing only what _create_transformer needs from self."""
def __init__(self, od_config: SimpleNamespace) -> None:
self.od_config = od_config
# Bind the real unbound method so the tests exercise production code.
_create_transformer = wan22_module.Wan22Pipeline._create_transformer
# ---------------------------------------------------------------------------
# create_transformer_from_config — auto-detection
# ---------------------------------------------------------------------------
def test_create_transformer_detects_mxfp8_serialized_from_config_json(monkeypatch):
"""When config.json carries MXFP8 quant and no CLI quant_config is provided,
the transformer must receive a DiffusionMXFP8Config with serialized=True."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
config = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp8",
"is_checkpoint_mxfp8_serialized": True,
},
}
create_transformer_from_config(config)
qc = captured[0].get("quant_config")
assert isinstance(qc, DiffusionMXFP8Config)
assert qc.is_checkpoint_mxfp8_serialized is True
def test_create_transformer_detects_mxfp4_dualscale_from_config_json(monkeypatch):
"""config.json with mxfp4_dualscale + ignored_layers must produce
a DiffusionMXFP4DualScaleMixedConfig with the correct ignored_layers and
serialized flag."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
config = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp4_dualscale",
"ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"],
"is_checkpoint_serialized": True,
},
}
create_transformer_from_config(config)
qc = captured[0].get("quant_config")
assert isinstance(qc, DiffusionMXFP4DualScaleMixedConfig)
assert set(qc.ignored_layers) == {"blocks.0.attn1.to_q", "blocks.0.attn1.to_k"}
assert qc.is_checkpoint_serialized is True
def test_create_transformer_without_quantization_config_passes_no_quant(monkeypatch):
"""A plain BF16 config.json (no quantization_config key) must result in no
quant_config being passed to WanTransformer3DModel."""
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
create_transformer_from_config(_MIN_CFG)
assert "quant_config" not in captured[0]
# ---------------------------------------------------------------------------
# create_transformer_from_config — method-mismatch guard
# ---------------------------------------------------------------------------
def test_create_transformer_rejects_method_mismatch(monkeypatch):
"""Passing a CLI quant_config whose get_name() differs from the config.json
quant_method must raise ValueError immediately (prevents silent weight corruption).
fp8 (vLLM built-in, get_name()=='fp8') vs disk 'mxfp8' triggers the guard.
These are distinct methods; using the same method for both would not trigger it.
"""
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
FakeTransformer, _ = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
fp8_cli = Fp8Config(is_checkpoint_fp8_serialized=True, activation_scheme="dynamic")
config = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp8",
"is_checkpoint_mxfp8_serialized": True,
},
}
with pytest.raises(ValueError, match="quant_method"):
create_transformer_from_config(config, quant_config=fp8_cli)
# ---------------------------------------------------------------------------
# create_transformer_from_config — online → offline upgrade
# ---------------------------------------------------------------------------
def test_create_transformer_upgrades_to_serialized_when_disk_marks_it(monkeypatch):
"""CLI passes online (is_checkpoint_mxfp8_serialized=False) but config.json
marks is_checkpoint_mxfp8_serialized=True → must switch to offline (serialized)
so that pre-quantized FP8 tensors are loaded correctly."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
online_cli = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False)
config = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp8",
"is_checkpoint_mxfp8_serialized": True,
},
}
create_transformer_from_config(config, quant_config=online_cli)
qc = captured[0].get("quant_config")
assert isinstance(qc, DiffusionMXFP8Config)
assert qc.is_checkpoint_mxfp8_serialized is True
# ---------------------------------------------------------------------------
# create_transformer_from_config — ignored_layers rebuild
# ---------------------------------------------------------------------------
def test_create_transformer_rebuilds_when_ignored_layers_differ(monkeypatch):
"""When the active quant_config has different ignored_layers than config.json,
the config must be rebuilt from disk so per-layer BF16 routing is authoritative
for this specific transformer."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
stale = DiffusionMXFP4DualScaleMixedConfig(
is_checkpoint_serialized=True,
ignored_layers=["blocks.0.attn1.to_q"],
)
config = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp4_dualscale",
"ignored_layers": ["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"],
"is_checkpoint_serialized": True,
},
}
create_transformer_from_config(config, quant_config=stale)
qc = captured[0].get("quant_config")
assert isinstance(qc, DiffusionMXFP4DualScaleMixedConfig)
assert set(qc.ignored_layers) == {"blocks.0.attn1.to_q", "blocks.1.attn1.to_q"}
def test_create_transformer_does_not_rebuild_when_ignored_layers_match(monkeypatch):
"""When the active quant_config already has the same ignored_layers,
the same instance must be passed through unchanged (no unnecessary rebuild)."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
matching = DiffusionMXFP4DualScaleMixedConfig(
is_checkpoint_serialized=True,
ignored_layers=["blocks.0.attn1.to_q"],
)
config = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp4_dualscale",
"ignored_layers": ["blocks.0.attn1.to_q"],
"is_checkpoint_serialized": True,
},
}
create_transformer_from_config(config, quant_config=matching)
assert captured[0].get("quant_config") is matching
# ---------------------------------------------------------------------------
# Wan22Pipeline._create_transformer — od_config passthrough
# ---------------------------------------------------------------------------
def test_pipeline_create_transformer_auto_detects_from_config_json(monkeypatch):
"""When od_config.quantization_config is None, _create_transformer must
auto-detect the quant method from config.json and pass it to the transformer.
od_config itself must remain unchanged (None)."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
od_config = SimpleNamespace(quantization_config=None)
pipeline = _FakePipeline(od_config)
config = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp8",
"is_checkpoint_mxfp8_serialized": True,
},
}
pipeline._create_transformer(config)
qc = captured[0].get("quant_config")
assert isinstance(qc, DiffusionMXFP8Config)
assert qc.is_checkpoint_mxfp8_serialized is True
# od_config must NOT be modified — each transformer auto-detects independently.
assert od_config.quantization_config is None
def test_pipeline_create_transformer_does_not_overwrite_existing_od_config(monkeypatch):
"""If od_config.quantization_config is already set (e.g. via CLI), _create_transformer
must pass it through to the transformer unchanged and leave od_config unmodified."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
FakeTransformer, _ = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
existing = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
od_config = SimpleNamespace(quantization_config=existing)
pipeline = _FakePipeline(od_config)
config = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp8",
"is_checkpoint_mxfp8_serialized": True,
},
}
pipeline._create_transformer(config)
assert od_config.quantization_config is existing
# ---------------------------------------------------------------------------
# Wan22Pipeline._create_transformer — cascade contracts
# ---------------------------------------------------------------------------
def test_pipeline_cascade_both_transformers_get_mxfp8_serialized_config(monkeypatch):
"""Cascade model (transformer + transformer_2) with MXFP8 checkpoint:
- Each transformer auto-detects from its own config.json independently.
Both must receive is_checkpoint_mxfp8_serialized=True."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
od_config = SimpleNamespace(quantization_config=None)
pipeline = _FakePipeline(od_config)
mxfp8_qc = {"quant_method": "mxfp8", "is_checkpoint_mxfp8_serialized": True}
pipeline._create_transformer({**_MIN_CFG, "quantization_config": mxfp8_qc})
pipeline._create_transformer({**_MIN_CFG, "quantization_config": mxfp8_qc})
assert len(captured) == 2
for i, kwargs in enumerate(captured):
qc = kwargs.get("quant_config")
assert isinstance(qc, DiffusionMXFP8Config), f"transformer[{i}]: expected DiffusionMXFP8Config, got {type(qc)}"
assert qc.is_checkpoint_mxfp8_serialized is True, f"transformer[{i}]: expected serialized=True"
def test_pipeline_cascade_mxfp4_dualscale_each_transformer_gets_correct_ignored_layers(monkeypatch):
"""Cascade with mxfp4_dualscale where transformer and transformer_2 have
different ignored_layers in their config.json.
Expected outcome:
transformer → ignored_layers=["blocks.0.attn1.to_q"] (auto-detected from its config.json)
transformer_2 → ignored_layers=["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"]
(auto-detected from its own config.json independently)
od_config → quantization_config remains None (not modified by _create_transformer)
"""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
FakeTransformer, captured = _make_fake_transformer()
monkeypatch.setattr(wan22_module, "WanTransformer3DModel", FakeTransformer)
od_config = SimpleNamespace(quantization_config=None)
pipeline = _FakePipeline(od_config)
cfg1 = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp4_dualscale",
"ignored_layers": ["blocks.0.attn1.to_q"],
"is_checkpoint_serialized": True,
},
}
cfg2 = {
**_MIN_CFG,
"quantization_config": {
"quant_method": "mxfp4_dualscale",
"ignored_layers": ["blocks.0.attn1.to_q", "blocks.1.attn1.to_q"],
"is_checkpoint_serialized": True,
},
}
pipeline._create_transformer(cfg1)
pipeline._create_transformer(cfg2)
assert len(captured) == 2
qc1 = captured[0].get("quant_config")
qc2 = captured[1].get("quant_config")
assert isinstance(qc1, DiffusionMXFP4DualScaleMixedConfig)
assert isinstance(qc2, DiffusionMXFP4DualScaleMixedConfig)
assert set(qc1.ignored_layers) == {"blocks.0.attn1.to_q"}, f"transformer expected 1 layer, got {qc1.ignored_layers}"
assert set(qc2.ignored_layers) == {
"blocks.0.attn1.to_q",
"blocks.1.attn1.to_q",
}, f"transformer_2 expected 2 layers, got {qc2.ignored_layers}"
# od_config must remain unchanged — _create_transformer does not modify it.
assert od_config.quantization_config is None
tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py
+301
View File
@@ -0,0 +1,301 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for Wan2.2 quant_config propagation through transformer creation.
Tests cover:
- create_transformer_from_config passes quant_config and prefix
- create_vace_transformer_from_config passes quant_config and prefix
- set_tf_model_config propagates quant_config to OmniDiffusionConfig
- patch_wan_rms_norm safely iterates sys.modules with concurrent modifications
- I2V transformer_2 quant_config is built from config dict
"""
import sys
from types import SimpleNamespace
import pytest
from pytest_mock import MockerFixture
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace as wan22_vace_module
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
create_transformer_from_config,
)
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
create_vace_transformer_from_config,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
# ---------------------------------------------------------------------------
# create_transformer_from_config: quant_config / prefix forwarding
# ---------------------------------------------------------------------------
class TestCreateTransformerQuant:
"""Verify quant_config and prefix are forwarded to WanTransformer3DModel."""
def test_quant_config_passed_through(self, mocker: MockerFixture):
captured = {}
class FakeTransformer:
def __init__(self, **kwargs):
captured.update(kwargs)
mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer)
fake_qc = mocker.MagicMock()
create_transformer_from_config(
{"patch_size": [1, 2, 2], "num_layers": 2},
quant_config=fake_qc,
)
assert captured.get("quant_config") is fake_qc
def test_prefix_passed_through(self, mocker: MockerFixture):
captured = {}
class FakeTransformer:
def __init__(self, **kwargs):
captured.update(kwargs)
mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer)
create_transformer_from_config(
{"patch_size": [1, 2, 2]},
prefix="model.transformer.",
)
assert captured.get("prefix") == "model.transformer."
def test_quant_config_none_by_default(self, mocker: MockerFixture):
captured = {}
class FakeTransformer:
def __init__(self, **kwargs):
captured.update(kwargs)
mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer)
create_transformer_from_config({"patch_size": [1, 2, 2]})
# When quant_config is None and prefix is "", they are not added
assert "quant_config" not in captured or captured["quant_config"] is None
def test_quant_config_and_prefix_together(self, mocker: MockerFixture):
captured = {}
class FakeTransformer:
def __init__(self, **kwargs):
captured.update(kwargs)
mocker.patch.object(wan22_module, "WanTransformer3DModel", FakeTransformer)
fake_qc = mocker.MagicMock()
create_transformer_from_config(
{"patch_size": [1, 2, 2], "num_attention_heads": 4},
quant_config=fake_qc,
prefix="blocks.",
)
assert captured["quant_config"] is fake_qc
assert captured["prefix"] == "blocks."
# ---------------------------------------------------------------------------
# create_vace_transformer_from_config: quant_config / prefix forwarding
# ---------------------------------------------------------------------------
class TestCreateVaceTransformerQuant:
"""Verify quant_config and prefix are forwarded to WanVACETransformer3DModel."""
def test_quant_config_passed_through(self, mocker: MockerFixture):
captured = {}
class FakeVACETransformer:
def __init__(self, **kwargs):
captured.update(kwargs)
mocker.patch.object(wan22_vace_module, "WanVACETransformer3DModel", FakeVACETransformer)
fake_qc = mocker.MagicMock()
create_vace_transformer_from_config(
{"patch_size": [1, 2, 2], "num_layers": 2},
quant_config=fake_qc,
)
assert captured.get("quant_config") is fake_qc
def test_prefix_passed_through(self, mocker: MockerFixture):
captured = {}
class FakeVACETransformer:
def __init__(self, **kwargs):
captured.update(kwargs)
mocker.patch.object(wan22_vace_module, "WanVACETransformer3DModel", FakeVACETransformer)
create_vace_transformer_from_config(
{"patch_size": [1, 2, 2]},
prefix="vace.",
)
assert captured.get("prefix") == "vace."
# ---------------------------------------------------------------------------
# set_tf_model_config: propagation of quant_config
# ---------------------------------------------------------------------------
class TestSetTfModelConfig:
"""Test that set_tf_model_config propagates quant_config correctly."""
def _make_od_config(self):
"""Create a minimal OmniDiffusionConfig-like object for testing."""
from vllm_omni.diffusion.data import OmniDiffusionConfig
cfg = object.__new__(OmniDiffusionConfig)
cfg.quantization_config = None
cfg.tf_model_config = None
return cfg
def test_propagates_quant_config_when_none(self, mocker: MockerFixture):
cfg = self._make_od_config()
fake_qc = mocker.MagicMock()
tf_config = SimpleNamespace(quant_config=fake_qc, quant_method="auto-round")
cfg.set_tf_model_config(tf_config)
assert cfg.tf_model_config is tf_config
assert cfg.quantization_config is fake_qc
def test_does_not_overwrite_existing_quantization_config(self, mocker: MockerFixture):
cfg = self._make_od_config()
existing_qc = mocker.MagicMock()
cfg.quantization_config = existing_qc
tf_config = SimpleNamespace(quant_config=mocker.MagicMock())
cfg.set_tf_model_config(tf_config)
assert cfg.tf_model_config is tf_config
assert cfg.quantization_config is existing_qc # not overwritten
def test_no_propagation_when_tf_quant_config_is_none(self, mocker: MockerFixture):
cfg = self._make_od_config()
tf_config = SimpleNamespace(quant_config=None)
cfg.set_tf_model_config(tf_config)
assert cfg.tf_model_config is tf_config
assert cfg.quantization_config is None
# ---------------------------------------------------------------------------
# patch_wan_rms_norm: sys.modules snapshot safety
# ---------------------------------------------------------------------------
class TestPatchWanRmsNorm:
"""Test that patch_wan_rms_norm doesn't raise on concurrent module registration."""
def test_patches_modules_with_wan_rms_norm(self):
from vllm_omni.diffusion.layers.norm import RMSNormVAE
from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm
# Create a fake module that has WanRMS_norm
fake_module = SimpleNamespace(WanRMS_norm=lambda x: x)
sys.modules["_test_fake_wan_module"] = fake_module
try:
patch_wan_rms_norm()
assert fake_module.WanRMS_norm is RMSNormVAE
finally:
del sys.modules["_test_fake_wan_module"]
def test_no_error_when_modules_change_during_iteration(self):
"""Regression test: list() snapshot prevents RuntimeError."""
from vllm_omni.diffusion.models.wan2_2.patch_diffusers import patch_wan_rms_norm
# Simulate a module being added during iteration by a side effect
original_items = sys.modules.items
def items_with_side_effect():
# This would cause RuntimeError without list() snapshot
result = list(original_items())
# Add a new module to simulate concurrent modification
sys.modules["_test_dynamic_module"] = SimpleNamespace()
return result
try:
# The function uses list(sys.modules.items()) so it takes a snapshot
# Just verify it doesn't raise
patch_wan_rms_norm()
finally:
sys.modules.pop("_test_dynamic_module", None)
# ---------------------------------------------------------------------------
# I2V transformer_2 quant_config extraction
# ---------------------------------------------------------------------------
class TestI2VTransformer2QuantConfig:
"""Test the transformer_2 quant_config build logic from pipeline_wan2_2_i2v."""
def test_transformer_2_quant_config_built_from_dict(self):
"""When transformer_2 config has quantization_config dict, build_quant_config is called."""
from vllm_omni.quantization.factory import build_quant_config
t2_config = {
"patch_size": [1, 2, 2],
"num_layers": 2,
"quantization_config": {
"quant_method": "auto-round",
"bits": 4,
"group_size": 128,
"sym": True,
"packing_format": "auto_round:auto_gptq",
},
}
# Replicate the logic from pipeline_wan2_2_i2v.py
t2_quant = t2_config.get("quantization_config")
if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
method = t2_quant["quant_method"]
kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"}
t2_quant = build_quant_config(method, **kwargs)
else:
t2_quant = None
from vllm.model_executor.layers.quantization.inc import INCConfig
assert isinstance(t2_quant, INCConfig)
assert t2_quant.weight_bits == 4
assert t2_quant.group_size == 128
def test_transformer_2_quant_config_none_when_missing(self):
"""When transformer_2 config has no quantization_config, result is None."""
t2_config = {
"patch_size": [1, 2, 2],
"num_layers": 2,
}
t2_quant = t2_config.get("quantization_config")
if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
pass # won't enter
else:
t2_quant = None
assert t2_quant is None
def test_transformer_2_quant_config_none_when_dict_lacks_method(self):
"""When quantization_config is a dict but missing quant_method, result is None."""
t2_config = {
"patch_size": [1, 2, 2],
"quantization_config": {"bits": 4}, # no quant_method key
}
t2_quant = t2_config.get("quantization_config")
if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
pass
else:
t2_quant = None
assert t2_quant is None
tests/diffusion/models/wan2_2/test_wan22_ti2v_pipeline.py
-97
View File
@@ -1,97 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from types import SimpleNamespace
import pytest
import torch
from PIL import Image
from torch import nn
from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import (
Wan22TI2VPipeline,
get_wan22_ti2v_pre_process_func,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def _make_ti2v_pipeline() -> Wan22TI2VPipeline:
pipeline = object.__new__(Wan22TI2VPipeline)
nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
pipeline.vae = StubVAE(z_dim=4)
pipeline.vae_scale_factor_temporal = 4
pipeline.vae_scale_factor_spatial = 8
pipeline.progress_bar = noop_progress_bar
return pipeline
def test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None:
preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace())
request = SimpleNamespace(
prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}],
sampling_params=SimpleNamespace(height=None, width=None),
)
result = preprocess(request)
assert result.sampling_params.height == 672
assert result.sampling_params.width == 1344
assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672)
def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None:
pipeline = _make_ti2v_pipeline()
latents = torch.zeros(1, 4, 2, 4, 4)
calls = []
def fake_predict_noise_maybe_with_cfg(**kwargs):
calls.append(kwargs)
return torch.ones_like(latents)
pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg # type: ignore[method-assign]
pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise # type: ignore[method-assign]
result = pipeline.diffuse(
latents=latents,
timesteps=torch.tensor([7]),
prompt_embeds=torch.zeros(1, 2, 3),
negative_prompt_embeds=torch.zeros(1, 2, 3),
guidance_scale=3.0,
dtype=torch.float32,
attention_kwargs={"a": "b"},
num_latent_frames=2,
latent_height=4,
latent_width=4,
)
positive = calls[0]["positive_kwargs"]
assert calls[0]["do_true_cfg"] is True
assert positive["timestep"].shape == (1, 8)
torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype))
torch.testing.assert_close(positive["hidden_states"], latents)
torch.testing.assert_close(result, torch.ones_like(latents))
def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None:
pipeline = _make_ti2v_pipeline()
latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents(
image=torch.zeros(1, 3, 16, 16),
batch_size=1,
num_channels_latents=4,
height=16,
width=16,
num_frames=5,
dtype=torch.float32,
device=torch.device("cpu"),
generator=None,
latents=torch.zeros(1, 4, 2, 2, 2),
)
torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2))
assert latent_condition.shape == (1, 4, 1, 2, 2)
assert first_frame_mask[:, :, 0].sum() == 0
assert first_frame_mask[:, :, 1].sum() == 4
tests/diffusion/quantization/test_mxfp4_config.py
+539
View File
@@ -0,0 +1,539 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for MXFP4 quantization configs and the MXFP4 DualScale + BF16 mixed config."""
import pytest
import torch
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _patch_tp_state(monkeypatch):
"""Patch TP rank/world_size so ModelWeightParameter can be instantiated on CPU
without an initialized distributed group. Returns TP=1 rank=0 for all tests."""
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_rank", lambda: 0)
monkeypatch.setattr("vllm.model_executor.parameter.get_tensor_model_parallel_world_size", lambda: 1)
# ---------------------------------------------------------------------------
# DiffusionMXFP4Config
# ---------------------------------------------------------------------------
def test_mxfp4_config_get_name():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
assert DiffusionMXFP4Config.get_name() == "mxfp4"
def test_mxfp4_config_from_config_defaults():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = DiffusionMXFP4Config.from_config({})
assert cfg.is_checkpoint_mxfp4_serialized is False
assert cfg.ignored_layers == []
def test_mxfp4_config_from_config_serialized():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = DiffusionMXFP4Config.from_config({"is_checkpoint_mxfp4_serialized": True})
assert cfg.is_checkpoint_mxfp4_serialized is True
def test_mxfp4_config_from_config_ignored_layers():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = DiffusionMXFP4Config.from_config({"ignored_layers": ["proj_out"]})
assert cfg.ignored_layers == ["proj_out"]
def test_mxfp4_config_from_config_modules_to_not_convert_fallback():
"""modules_to_not_convert must be accepted as an alias for ignored_layers."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = DiffusionMXFP4Config.from_config({"modules_to_not_convert": ["proj_out"]})
assert cfg.ignored_layers == ["proj_out"]
# ---------------------------------------------------------------------------
# build_quant_config integration
# ---------------------------------------------------------------------------
def test_build_quant_config_mxfp4_string():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = build_quant_config("mxfp4")
assert isinstance(cfg, DiffusionMXFP4Config)
assert cfg.get_name() == "mxfp4"
assert cfg.is_checkpoint_mxfp4_serialized is False
def test_build_quant_config_mxfp4_dict():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config
cfg = build_quant_config({"method": "mxfp4", "is_checkpoint_mxfp4_serialized": True})
assert isinstance(cfg, DiffusionMXFP4Config)
assert cfg.is_checkpoint_mxfp4_serialized is True
def test_build_quant_config_mxfp4_dualscale_string():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = build_quant_config("mxfp4_dualscale")
assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
assert cfg.is_checkpoint_serialized is False
assert cfg.num_bf16_fallback_layers == 5
assert cfg.ignored_layers == []
def test_build_quant_config_mxfp4_dualscale_dict_offline():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = build_quant_config(
{
"method": "mxfp4_dualscale",
"is_checkpoint_serialized": True,
"ignored_layers": ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"],
}
)
assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
assert cfg.is_checkpoint_serialized is True
assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"]
def test_build_quant_config_mxfp4_dualscale_dict_online_custom_fallback():
from vllm_omni.quantization import build_quant_config
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = build_quant_config({"method": "mxfp4_dualscale", "num_bf16_fallback_layers": 10})
assert isinstance(cfg, DiffusionMXFP4DualScaleMixedConfig)
assert cfg.num_bf16_fallback_layers == 10
# ---------------------------------------------------------------------------
# Block-index dispatch (_parse_block_idx)
# ---------------------------------------------------------------------------
def test_parse_block_idx_valid():
from vllm_omni.quantization.mxfp4_config import _parse_block_idx
assert _parse_block_idx("blocks.0.attn1.to_q") == 0
assert _parse_block_idx("blocks.5.ffn.net.0.proj") == 5
assert _parse_block_idx("blocks.40.norm1.weight") == 40
def test_parse_block_idx_non_block_prefixes():
"""Prefixes that do not start with 'blocks.N.' must return None."""
from vllm_omni.quantization.mxfp4_config import _parse_block_idx
assert _parse_block_idx("condition_embedder.time_embedder.linear_1") is None
assert _parse_block_idx("proj_out.weight") is None
assert _parse_block_idx("model.layers.0.self_attn.q_proj") is None
assert _parse_block_idx("scale_shift_table") is None
# ---------------------------------------------------------------------------
# SUPPORTED_QUANTIZATION_METHODS
# ---------------------------------------------------------------------------
def test_supported_methods_include_mxfp4_variants():
from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
assert "mxfp4" in SUPPORTED_QUANTIZATION_METHODS
assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS
assert "mxfp4_dualscale" in SUPPORTED_QUANTIZATION_METHODS
# ---------------------------------------------------------------------------
# DiffusionMXFP4DualScaleMixedConfig — config roundtrips
# ---------------------------------------------------------------------------
def test_mixed_dualscale_config_get_name():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
assert DiffusionMXFP4DualScaleMixedConfig.get_name() == "mxfp4_dualscale"
def test_mixed_dualscale_config_no_args_defaults():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = DiffusionMXFP4DualScaleMixedConfig()
assert cfg.is_checkpoint_serialized is False
assert cfg.ignored_layers == []
assert cfg.num_bf16_fallback_layers == 5
def test_mixed_dualscale_config_from_config_offline():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = DiffusionMXFP4DualScaleMixedConfig.from_config(
{
"quant_method": "mxfp4_dualscale",
"is_checkpoint_serialized": True,
"ignored_layers": ["blocks.0.attn1.to_q", "proj_out"],
}
)
assert cfg.is_checkpoint_serialized is True
assert cfg.ignored_layers == ["blocks.0.attn1.to_q", "proj_out"]
assert cfg.num_bf16_fallback_layers == 5 # default
def test_mixed_dualscale_config_from_config_online_custom_fallback():
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = DiffusionMXFP4DualScaleMixedConfig.from_config({"num_bf16_fallback_layers": 10})
assert cfg.is_checkpoint_serialized is False
assert cfg.num_bf16_fallback_layers == 10
def test_mixed_dualscale_config_from_config_modules_to_not_convert_fallback():
"""modules_to_not_convert must be accepted as an alias for ignored_layers."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = DiffusionMXFP4DualScaleMixedConfig.from_config(
{"is_checkpoint_serialized": True, "modules_to_not_convert": ["proj_out"]}
)
assert cfg.ignored_layers == ["proj_out"]
# ---------------------------------------------------------------------------
# DiffusionMXFP4DualScaleMixedConfig — get_quant_method dispatch
# ---------------------------------------------------------------------------
def test_mixed_dualscale_offline_ignored_layer_returns_unquantized(
mocker,
monkeypatch: pytest.MonkeyPatch,
):
"""Offline: a prefix in ignored_layers must return UnquantizedLinearMethod."""
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = DiffusionMXFP4DualScaleMixedConfig(
is_checkpoint_serialized=True,
ignored_layers=["blocks.0.attn1.to_q"],
)
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
method = cfg.get_quant_method(layer, "blocks.0.attn1.to_q")
assert isinstance(method, UnquantizedLinearMethod)
def test_mixed_dualscale_offline_non_ignored_returns_mxfp4(
mocker,
monkeypatch: pytest.MonkeyPatch,
):
"""Offline: a prefix NOT in ignored_layers must return NPUMxfp4DualScaleLinearMethod."""
from vllm.model_executor.layers.linear import LinearBase
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
cfg = DiffusionMXFP4DualScaleMixedConfig(
is_checkpoint_serialized=True,
ignored_layers=["blocks.0.attn1.to_q"],
)
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
method = cfg.get_quant_method(layer, "blocks.1.attn1.to_q")
assert isinstance(method, NPUMxfp4DualScaleLinearMethod)
def test_mixed_dualscale_online_fallback_block_returns_unquantized(
mocker,
monkeypatch: pytest.MonkeyPatch,
):
"""Online: blocks < num_bf16_fallback_layers must return UnquantizedLinearMethod."""
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
assert isinstance(cfg.get_quant_method(layer, "blocks.0.attn1.to_q"), UnquantizedLinearMethod)
assert isinstance(cfg.get_quant_method(layer, "blocks.4.ffn.net.0.proj"), UnquantizedLinearMethod)
def test_mixed_dualscale_online_quantized_block_returns_mxfp4(
mocker,
monkeypatch: pytest.MonkeyPatch,
):
"""Online: blocks >= num_bf16_fallback_layers must return NPUMxfp4DualScaleOnlineLinearMethod."""
from vllm.model_executor.layers.linear import LinearBase
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization.mxfp4_config import (
DiffusionMXFP4DualScaleMixedConfig,
NPUMxfp4DualScaleOnlineLinearMethod,
)
cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
assert isinstance(cfg.get_quant_method(layer, "blocks.5.attn1.to_q"), NPUMxfp4DualScaleOnlineLinearMethod)
assert isinstance(cfg.get_quant_method(layer, "blocks.40.ffn.net.0.proj"), NPUMxfp4DualScaleOnlineLinearMethod)
def test_mixed_dualscale_online_non_block_prefix_returns_mxfp4(
mocker,
monkeypatch: pytest.MonkeyPatch,
):
"""Online: layers outside 'blocks.N.*' (condition_embedder etc.) always use MXFP4 online."""
from vllm.model_executor.layers.linear import LinearBase
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization.mxfp4_config import (
DiffusionMXFP4DualScaleMixedConfig,
NPUMxfp4DualScaleOnlineLinearMethod,
)
cfg = DiffusionMXFP4DualScaleMixedConfig(is_checkpoint_serialized=False, num_bf16_fallback_layers=5)
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
method = cfg.get_quant_method(layer, "condition_embedder.time_embedder.linear_1")
assert isinstance(method, NPUMxfp4DualScaleOnlineLinearMethod)
def test_mixed_dualscale_online_ignored_layers_override(
mocker,
monkeypatch: pytest.MonkeyPatch,
):
"""Online: explicit ignored_layers must return UnquantizedLinearMethod regardless of block index.
A layer that is NOT in the leading-block range (block 10 >= num_bf16_fallback_layers=5)
but IS listed in ignored_layers must still fall back to BF16. This lets power users
pin specific interleaved layers to BF16 during online quantization without needing an
offline checkpoint.
"""
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = DiffusionMXFP4DualScaleMixedConfig(
is_checkpoint_serialized=False,
num_bf16_fallback_layers=5,
ignored_layers=["blocks.10.attn1.to_q"],
)
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
# block 10 is above the leading-block threshold but is in ignored_layers → BF16
assert isinstance(cfg.get_quant_method(layer, "blocks.10.attn1.to_q"), UnquantizedLinearMethod)
def test_mixed_dualscale_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch):
"""Non-LinearBase layers (norms, embeddings) must return None → no quantization."""
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig
cfg = DiffusionMXFP4DualScaleMixedConfig()
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
norm_layer = torch.nn.LayerNorm(64)
assert cfg.get_quant_method(norm_layer, "blocks.0.norm1") is None
# ---------------------------------------------------------------------------
# TP=2 create_weights: parameter shapes and input_dim/output_dim
#
# Two scenarios mirror real Wan2.2 A14B linear layer types:
# Column-parallel (to_q, ffn.net_0): output is sharded (N/TP), input is full (K).
# Row-parallel (to_out, ffn.net_2): input is sharded (K/TP), output is full (N).
#
# Tests verify:
# 1. Registered parameter shapes are correct for each partition configuration.
# 2. input_dim/output_dim attributes are set so RowParallelLinear.weight_loader
# can shard scale tensors correctly (the fix for the TP>1 shape-mismatch bug).
# 3. Simulated loader slicing: slicing the full checkpoint tensor along the
# declared input_dim produces the exact shape stored in the parameter —
# proving the dim declaration is consistent with the allocation.
# ---------------------------------------------------------------------------
# K must be divisible by 32 (fine groups) and 512 (coarse groups).
_TP2_K, _TP2_N, _TP2 = 1024, 512, 2
class _FakeLayer(torch.nn.Module):
"""Bare nn.Module that accepts register_parameter without a real weight_loader."""
def _create_weights(method, *, input_size_per_partition, output_partition_sizes):
layer = _FakeLayer()
method.create_weights(
layer=layer,
input_size_per_partition=input_size_per_partition,
output_partition_sizes=output_partition_sizes,
input_size=_TP2_K,
output_size=_TP2_N,
params_dtype=torch.bfloat16,
)
return layer
def _shard(tensor, param, rank, tp, dim_attr):
"""Slice `tensor` along the dimension given by `param.<dim_attr>` for `rank`."""
dim = getattr(param, dim_attr)
if dim is None:
return tensor # not sharded along this axis
shard_size = param.shape[dim]
slices = [slice(None)] * tensor.ndim
slices[dim] = slice(rank * shard_size, (rank + 1) * shard_size)
return tensor[tuple(slices)]
# ── DualScale method ─────────────────────────────────────────────────────────
def test_dualscale_column_parallel_tp2_shapes():
"""Column-parallel TP=2: output halved, fine/coarse groups stay full, mul_scale full."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K)
assert layer.weight_scale.shape == (_TP2_N // _TP2, _TP2_K // 32)
assert layer.weight_dual_scale.shape == (_TP2_N // _TP2, _TP2_K // 512, 1)
assert layer.mul_scale.shape == (_TP2_K,)
def test_dualscale_row_parallel_tp2_shapes():
"""Row-parallel TP=2: input halved, fine/coarse groups halved, mul_scale halved."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32)
assert layer.weight_dual_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 512, 1)
assert layer.mul_scale.shape == (_TP2_K // _TP2,)
def test_dualscale_scale_parameter_input_dims():
"""weight_scale/weight_dual_scale must have input_dim=1; mul_scale must have input_dim=0.
RowParallelLinear.weight_loader only shards a parameter when input_dim is set.
Without these, loading a full checkpoint tensor into a per-rank shape causes a
shape mismatch for TP>1 row-parallel layers (to_out, ffn.net_2).
"""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N])
assert layer.weight_scale.input_dim == 1
assert layer.weight_scale.output_dim == 0
assert layer.weight_dual_scale.input_dim == 1
assert layer.weight_dual_scale.output_dim == 0
assert layer.mul_scale.input_dim == 0
assert layer.mul_scale.output_dim is None
def test_dualscale_row_parallel_tp2_loader_simulation():
"""Slicing full checkpoint tensors along input_dim must match row-parallel parameter shapes.
Simulates what RowParallelLinear.weight_loader does: for each scale parameter,
take the slice at rank*shard_size:(rank+1)*shard_size along input_dim.
The resulting shape must equal the per-rank parameter shape allocated by create_weights.
"""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
# Full checkpoint tensors (what the loader reads from disk).
ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1)
ckpt_mul_scale = torch.zeros(_TP2_K)
for rank in range(_TP2):
assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape
assert (
_shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "input_dim").shape
== layer.weight_dual_scale.shape
)
assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "input_dim").shape == layer.mul_scale.shape
def test_dualscale_column_parallel_tp2_loader_simulation():
"""Slicing full checkpoint tensors along output_dim must match column-parallel parameter shapes.
For column-parallel layers, the loader shards along output_dim (rows).
mul_scale has output_dim=None → not sharded (full tensor, same for all ranks).
"""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4DualScaleMixedConfig, NPUMxfp4DualScaleLinearMethod
method = NPUMxfp4DualScaleLinearMethod(DiffusionMXFP4DualScaleMixedConfig())
layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2])
ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
ckpt_weight_dual_scale = torch.zeros(_TP2_N, _TP2_K // 512, 1)
ckpt_mul_scale = torch.zeros(_TP2_K)
for rank in range(_TP2):
assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "output_dim").shape == layer.weight_scale.shape
assert (
_shard(ckpt_weight_dual_scale, layer.weight_dual_scale, rank, _TP2, "output_dim").shape
== layer.weight_dual_scale.shape
)
# mul_scale: output_dim=None → no sharding → full tensor fits the column-parallel parameter
assert _shard(ckpt_mul_scale, layer.mul_scale, rank, _TP2, "output_dim").shape == layer.mul_scale.shape
# ── Single-scale method ───────────────────────────────────────────────────────
def test_single_scale_row_parallel_tp2_shapes():
"""Row-parallel TP=2: input halved → weight_scale groups halved."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2)
assert layer.weight_scale.shape == (_TP2_N, (_TP2_K // _TP2) // 32)
def test_single_scale_scale_parameter_input_dims():
"""Single-scale weight_scale must have input_dim=1 for RowParallel TP sharding."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N])
assert layer.weight_scale.input_dim == 1
assert layer.weight_scale.output_dim == 0
def test_single_scale_row_parallel_tp2_loader_simulation():
"""Slicing full checkpoint weight_scale along input_dim matches row-parallel parameter shape."""
from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, NPUMxfp4LinearMethod
method = NPUMxfp4LinearMethod(DiffusionMXFP4Config())
layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N])
ckpt_weight_scale = torch.zeros(_TP2_N, _TP2_K // 32)
for rank in range(_TP2):
assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape
tests/diffusion/quantization/test_mxfp4_key_remap.py
+399
View File
@@ -0,0 +1,399 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for merge_mxfp4_dualscale_checkpoint.py key-remapping helpers.
These are pure-Python unit tests that exercise the transformation functions
without loading any actual checkpoint files or requiring NPU hardware.
"""
import pytest
import torch
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
# ---------------------------------------------------------------------------
# SUPPORTED_MODEL_TYPES
# ---------------------------------------------------------------------------
def test_supported_model_types_includes_a14b():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES
assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES
assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES
def test_supported_model_types_excludes_ti2v_5b():
"""TI2V-5B is explicitly NOT supported under W4A4 quantization."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import SUPPORTED_MODEL_TYPES
assert "Wan2.2-TI2V-5B" not in SUPPORTED_MODEL_TYPES
# ---------------------------------------------------------------------------
# _apply_rename_dict
# ---------------------------------------------------------------------------
def test_apply_rename_dict_self_attn_qkvo():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
assert _apply_rename_dict("blocks.0.self_attn.q.weight") == "blocks.0.attn1.to_q.weight"
assert _apply_rename_dict("blocks.0.self_attn.k.weight") == "blocks.0.attn1.to_k.weight"
assert _apply_rename_dict("blocks.0.self_attn.v.weight") == "blocks.0.attn1.to_v.weight"
assert _apply_rename_dict("blocks.0.self_attn.o.weight") == "blocks.0.attn1.to_out.0.weight"
def test_apply_rename_dict_ffn():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
assert _apply_rename_dict("blocks.1.ffn.0.weight") == "blocks.1.ffn.net.0.proj.weight"
assert _apply_rename_dict("blocks.1.ffn.2.weight") == "blocks.1.ffn.net.2.weight"
def test_apply_rename_dict_norm_order_swap():
"""norm2↔norm3 swap: quant tool uses norm1/norm3/norm2 order,
Diffusers uses norm1/norm2/norm3."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
assert _apply_rename_dict("blocks.0.norm2.weight") == "blocks.0.norm3.weight"
assert _apply_rename_dict("blocks.0.norm3.weight") == "blocks.0.norm2.weight"
def test_apply_rename_dict_cross_attn():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
assert _apply_rename_dict("blocks.0.cross_attn.q.weight") == "blocks.0.attn2.to_q.weight"
assert _apply_rename_dict("blocks.0.cross_attn.k.weight") == "blocks.0.attn2.to_k.weight"
assert _apply_rename_dict("blocks.0.cross_attn.v.weight") == "blocks.0.attn2.to_v.weight"
assert _apply_rename_dict("blocks.0.cross_attn.o.weight") == "blocks.0.attn2.to_out.0.weight"
def test_apply_rename_dict_head_and_modulation():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _apply_rename_dict
assert _apply_rename_dict("head.head.weight") == "proj_out.weight"
# ---------------------------------------------------------------------------
# _strip_mxfp4_wrapper
# ---------------------------------------------------------------------------
def test_strip_mxfp4_wrapper_linear_weight():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight") == "blocks.0.attn1.to_q.weight"
def test_strip_mxfp4_wrapper_linear_weight_scale():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_scale") == "blocks.0.attn1.to_q.weight_scale"
def test_strip_mxfp4_wrapper_linear_weight_dual_scale():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
assert (
_strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.weight_dual_scale") == "blocks.0.attn1.to_q.weight_dual_scale"
)
def test_strip_mxfp4_wrapper_linear_bias():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.linear.bias") == "blocks.0.attn1.to_q.bias"
def test_strip_mxfp4_wrapper_div_mul_scale():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.div.mul_scale") == "blocks.0.attn1.to_q.mul_scale"
def test_strip_mxfp4_wrapper_noop_for_plain_weight():
"""MXFP8 / FLOAT tensors have no wrapper — must be returned unchanged."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _strip_mxfp4_wrapper
assert _strip_mxfp4_wrapper("blocks.0.attn1.to_q.weight") == "blocks.0.attn1.to_q.weight"
assert _strip_mxfp4_wrapper("blocks.0.norm_q.weight") == "blocks.0.norm_q.weight"
assert _strip_mxfp4_wrapper("condition_embedder.time_embedder.linear_1.weight") == (
"condition_embedder.time_embedder.linear_1.weight"
)
# ---------------------------------------------------------------------------
# _classify_blocks
# ---------------------------------------------------------------------------
def test_classify_blocks_mixed_mxfp8_and_mxfp4():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
quant_meta = {
"blocks.0.attn1.to_q.weight": "W8A8_MXFP8",
"blocks.1.attn1.to_q.weight": "W8A8_MXFP8",
"blocks.2.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
"blocks.3.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
"condition_embedder.time_embedder.linear_1.weight": "FLOAT",
}
block_types = _classify_blocks(quant_meta)
assert block_types[0] == "mxfp8"
assert block_types[1] == "mxfp8"
assert block_types[2] == "mxfp4_dualscale"
assert block_types[3] == "mxfp4_dualscale"
# Non-block key must not produce an entry
assert None not in block_types
def test_classify_blocks_float_entries_skipped():
"""FLOAT-typed tensors must not contribute to block classification."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
quant_meta = {
"blocks.0.bias": "FLOAT",
"blocks.1.attn1.to_q.weight": "W4A4_MXFP4_DUALSCALE",
}
block_types = _classify_blocks(quant_meta)
assert 0 not in block_types
assert block_types[1] == "mxfp4_dualscale"
def test_classify_blocks_empty():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _classify_blocks
assert _classify_blocks({}) == {}
# ---------------------------------------------------------------------------
# _is_mxfp4_tensor
# ---------------------------------------------------------------------------
def test_is_mxfp4_tensor_quantized_weight():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight", "W4A4_MXFP4_DUALSCALE") is True
assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_scale", "W4A4_MXFP4_DUALSCALE") is True
assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.weight_dual_scale", "W4A4_MXFP4_DUALSCALE") is True
def test_is_mxfp4_tensor_companion_bias_and_mul_scale():
"""Companion tensors (bias, mul_scale) are FLOAT but belong to MXFP4 layers
and must be included in the merged checkpoint."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
# bias inside .linear. wrapper → must be included
assert _is_mxfp4_tensor("blocks.0.attn1.to_q.linear.bias", "FLOAT") is True
# mul_scale inside .div. wrapper → must be included
assert _is_mxfp4_tensor("blocks.0.attn1.to_q.div.mul_scale", "FLOAT") is True
def test_is_mxfp4_tensor_bf16_fallback_weight():
"""BF16 fallback linear layers have plain .weight keys (no wrapper) → must NOT be included."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
assert _is_mxfp4_tensor("blocks.0.attn1.to_q.weight", "FLOAT") is False
assert _is_mxfp4_tensor("condition_embedder.time_embedder.linear_1.weight", "FLOAT") is False
def test_is_mxfp4_tensor_norm_weight():
"""Norm layers are always FLOAT with plain .weight → not MXFP4 related."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _is_mxfp4_tensor
assert _is_mxfp4_tensor("blocks.0.norm1.weight", "FLOAT") is False
# ---------------------------------------------------------------------------
# _collect_ignored_layers
# ---------------------------------------------------------------------------
def test_collect_ignored_layers_basic():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
# block 0: all three of Q/K/V are BF16 → fused as to_qkv in ignored_layers
# block 1: to_qkv is MXFP4 → not in ignored_layers
merged = {
"blocks.0.attn1.to_q.weight": torch.zeros(4, 4),
"blocks.0.attn1.to_k.weight": torch.zeros(4, 4),
"blocks.0.attn1.to_v.weight": torch.zeros(4, 4),
"blocks.1.attn1.to_q.weight": torch.zeros(4, 4),
"blocks.1.attn1.to_q.weight_scale": torch.zeros(4, 1),
}
mxfp4_prefixes = {"blocks.1.attn1.to_q"}
ignored = _collect_ignored_layers(merged, mxfp4_prefixes)
assert "blocks.0.attn1.to_qkv" in ignored # fused vllm-omni name
assert "blocks.0.attn1.to_q" not in ignored # diffusers name must not appear
assert "blocks.0.attn1.to_k" not in ignored
assert "blocks.1.attn1.to_qkv" not in ignored # MXFP4, not ignored
def test_collect_ignored_layers_empty_mxfp4():
"""When no layers are MXFP4 (all BF16), all .weight prefixes become ignored_layers."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
# Use non-attn layers to avoid the Q/K/V completeness requirement.
merged = {
"blocks.0.ffn.net.0.proj.weight": torch.zeros(4, 4),
"proj_out.weight": torch.zeros(4, 4),
}
ignored = _collect_ignored_layers(merged, set())
assert sorted(ignored) == ["blocks.0.ffn.net_0.proj", "proj_out"]
def test_collect_ignored_layers_returns_sorted():
"""ignored_layers must be sorted for deterministic config.json output."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _collect_ignored_layers
# Use FFN layers (no Q/K/V fusion concern) to isolate the sort guarantee.
merged = {
"blocks.2.ffn.net.0.proj.weight": torch.zeros(1),
"blocks.0.ffn.net.0.proj.weight": torch.zeros(1),
"blocks.1.ffn.net.0.proj.weight": torch.zeros(1),
}
ignored = _collect_ignored_layers(merged, set())
assert ignored == sorted(ignored)
# ---------------------------------------------------------------------------
# _build_quant_config (new mxfp4_dualscale format)
# ---------------------------------------------------------------------------
def test_build_quant_config_new_format():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config
config = _build_quant_config(["blocks.0.attn1.to_q", "proj_out"])
assert config["quant_method"] == "mxfp4_dualscale"
assert config["is_checkpoint_serialized"] is True
assert config["ignored_layers"] == ["blocks.0.attn1.to_q", "proj_out"]
assert "num_mxfp8_blocks" not in config
def test_build_quant_config_empty_ignored_layers():
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _build_quant_config
config = _build_quant_config([])
assert config["ignored_layers"] == []
# ---------------------------------------------------------------------------
# _diffusers_to_vllm_ignored
# ---------------------------------------------------------------------------
def test_remap_self_attn_qkv_fusion():
"""All three of attn1.to_q/k/v in BF16 → fused attn1.to_qkv."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_k", "blocks.0.attn1.to_q", "blocks.0.attn1.to_v"])
assert result == ["blocks.0.attn1.to_qkv"]
def test_remap_self_attn_qkv_partial_raises():
"""Partial Q/K/V in ignored_layers is invalid — must raise ValueError.
Self-attention Q/K/V are fused into a single to_qkv layer at runtime;
partial precision (some BF16, some MXFP4) cannot be expressed and must
be caught early in the merge script.
"""
import pytest
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
with pytest.raises(ValueError, match="Partial BF16 fallback"):
_diffusers_to_vllm_ignored(["blocks.0.attn1.to_q", "blocks.0.attn1.to_k"])
def test_remap_cross_attn_kept_separate():
"""Cross-attention (attn2) to_q/k/v must NOT be fused."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
result = _diffusers_to_vllm_ignored(["blocks.0.attn2.to_k", "blocks.0.attn2.to_q", "blocks.0.attn2.to_v"])
assert "blocks.0.attn2.to_qkv" not in result
assert "blocks.0.attn2.to_q" in result
assert "blocks.0.attn2.to_k" in result
assert "blocks.0.attn2.to_v" in result
def test_remap_ffn_dot_to_underscore():
"""ffn.net.0.proj → ffn.net_0.proj; ffn.net.2 → ffn.net_2."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
result = _diffusers_to_vllm_ignored(["blocks.0.ffn.net.0.proj", "blocks.0.ffn.net.2"])
assert "blocks.0.ffn.net_0.proj" in result
assert "blocks.0.ffn.net_2" in result
assert "blocks.0.ffn.net.0.proj" not in result
assert "blocks.0.ffn.net.2" not in result
def test_remap_to_out_strips_index():
"""attn1.to_out.0 and attn2.to_out.0 must have the .0 suffix removed."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
result = _diffusers_to_vllm_ignored(["blocks.0.attn1.to_out.0", "blocks.0.attn2.to_out.0"])
assert "blocks.0.attn1.to_out" in result
assert "blocks.0.attn2.to_out" in result
assert "blocks.0.attn1.to_out.0" not in result
assert "blocks.0.attn2.to_out.0" not in result
def test_remap_noop_for_non_block_layers():
"""Layers outside 'blocks.N.*' (condition_embedder, proj_out) pass through unchanged."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
inputs = [
"condition_embedder.time_embedder.linear_1",
"condition_embedder.text_embedder.linear_1",
"proj_out",
]
result = _diffusers_to_vllm_ignored(inputs)
assert result == sorted(inputs)
def test_remap_mixed_block_layers():
"""Typical mixed block: QKV fused, to_out stripped, FFN renamed, cross-attn separate."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
inputs = [
"blocks.0.attn1.to_k",
"blocks.0.attn1.to_out.0",
"blocks.0.attn1.to_q",
"blocks.0.attn1.to_v",
"blocks.0.attn2.to_k",
"blocks.0.attn2.to_out.0",
"blocks.0.attn2.to_q",
"blocks.0.attn2.to_v",
"blocks.0.ffn.net.0.proj",
"blocks.0.ffn.net.2",
]
result = _diffusers_to_vllm_ignored(inputs)
assert "blocks.0.attn1.to_qkv" in result
assert "blocks.0.attn1.to_out" in result
assert "blocks.0.attn2.to_q" in result
assert "blocks.0.attn2.to_k" in result
assert "blocks.0.attn2.to_v" in result
assert "blocks.0.attn2.to_out" in result
assert "blocks.0.ffn.net_0.proj" in result
assert "blocks.0.ffn.net_2" in result
# Old names must be gone
assert "blocks.0.attn1.to_q" not in result
assert "blocks.0.ffn.net.0.proj" not in result
def test_remap_returns_sorted():
"""Output of _diffusers_to_vllm_ignored must always be sorted."""
from vllm_omni.quantization.tools.merge_mxfp4_dualscale_checkpoint import _diffusers_to_vllm_ignored
inputs = [
"blocks.2.ffn.net.0.proj",
"blocks.0.attn1.to_q",
"blocks.0.attn1.to_k",
"blocks.0.attn1.to_v",
"blocks.1.attn2.to_out.0",
]
result = _diffusers_to_vllm_ignored(inputs)
assert result == sorted(result)
tests/diffusion/quantization/test_mxfp8_config.py
+308
View File
@@ -0,0 +1,308 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for MXFP8 quantization config and linear method dispatch.
Coverage:
- DiffusionMXFP8Config.from_config roundtrips (CPU, no NPU required)
- get_quant_method dispatch (mocked platform)
- MXFPLinearMethodBase.apply() reshape skeleton (CPU)
- Weight / scale shape-transform arithmetic from process_weights_after_loading (CPU)
- build_quant_config integration
- MXFP8_QUANT_CONFIG structure as the auto-detection contract
"""
import pytest
import torch
from pytest_mock import MockerFixture
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
from vllm_omni.platforms import current_omni_platform
from vllm_omni.quantization import build_quant_config
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
npu_available = pytest.mark.skipif(not current_omni_platform.is_npu(), reason="NPU platform not available")
# ---------------------------------------------------------------------------
# DiffusionMXFP8Config — from_config roundtrips
# ---------------------------------------------------------------------------
def test_mxfp8_config_get_name():
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
assert DiffusionMXFP8Config.get_name() == "mxfp8"
def test_mxfp8_config_from_config_defaults():
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
cfg = DiffusionMXFP8Config.from_config({})
assert cfg.is_checkpoint_mxfp8_serialized is False
assert cfg.ignored_layers == []
def test_mxfp8_config_from_config_serialized():
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
cfg = DiffusionMXFP8Config.from_config({"is_checkpoint_mxfp8_serialized": True})
assert cfg.is_checkpoint_mxfp8_serialized is True
def test_mxfp8_config_from_config_ignored_layers():
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
cfg = DiffusionMXFP8Config.from_config({"ignored_layers": ["proj_out"]})
assert cfg.ignored_layers == ["proj_out"]
def test_mxfp8_config_from_config_modules_to_not_convert_fallback():
"""modules_to_not_convert must be accepted as an alias for ignored_layers."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
cfg = DiffusionMXFP8Config.from_config({"modules_to_not_convert": ["proj_out"]})
assert cfg.ignored_layers == ["proj_out"]
# ---------------------------------------------------------------------------
# build_quant_config integration
# ---------------------------------------------------------------------------
def test_build_quant_config_mxfp8_string():
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
cfg = build_quant_config("mxfp8")
assert isinstance(cfg, DiffusionMXFP8Config)
assert cfg.get_name() == "mxfp8"
assert cfg.is_checkpoint_mxfp8_serialized is False
def test_build_quant_config_mxfp8_dict():
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
cfg = build_quant_config({"method": "mxfp8", "is_checkpoint_mxfp8_serialized": True})
assert isinstance(cfg, DiffusionMXFP8Config)
assert cfg.is_checkpoint_mxfp8_serialized is True
def test_build_quant_config_mxfp8_config_json_format():
"""Verify that the exact quantization_config injected by merge_mxfp8_checkpoint.py
is accepted by build_quant_config and selects the offline (serialized) path.
This is the critical auto-detection contract: TransformerConfig.from_dict()
reads quant_method + is_checkpoint_mxfp8_serialized to pick NPUMxfp8LinearMethod.
"""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
cfg = build_quant_config(MXFP8_QUANT_CONFIG)
assert isinstance(cfg, DiffusionMXFP8Config)
assert cfg.is_checkpoint_mxfp8_serialized is True
def test_mxfp8_quant_config_structure():
"""MXFP8_QUANT_CONFIG must contain exactly the keys that auto-detection reads."""
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
assert MXFP8_QUANT_CONFIG.get("quant_method") == "mxfp8"
assert MXFP8_QUANT_CONFIG.get("is_checkpoint_mxfp8_serialized") is True
# ---------------------------------------------------------------------------
# get_quant_method dispatch
# ---------------------------------------------------------------------------
def test_get_quant_method_npu_offline(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""Offline (serialized) path on NPU must return NPUMxfp8LinearMethod."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8LinearMethod
config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
assert isinstance(method, NPUMxfp8LinearMethod)
def test_get_quant_method_npu_online(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""Online (BF16 checkpoint) path on NPU must return NPUMxfp8OnlineLinearMethod."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8OnlineLinearMethod
config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False)
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
assert isinstance(method, NPUMxfp8OnlineLinearMethod)
def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""Non-NPU platform must raise NotImplementedError (CUDA not yet supported)."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
config = DiffusionMXFP8Config()
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
with pytest.raises(NotImplementedError):
config.get_quant_method(layer, "blocks.0.attn1.to_q")
def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""A prefix in ignored_layers must bypass quantization → UnquantizedLinearMethod."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
config = DiffusionMXFP8Config(ignored_layers=["proj_out"])
layer = mocker.Mock(spec=LinearBase)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
method = config.get_quant_method(layer, "proj_out")
assert isinstance(method, UnquantizedLinearMethod)
def test_get_quant_method_non_linear_returns_none(monkeypatch: pytest.MonkeyPatch):
"""Non-LinearBase layers (norms, embeddings) must get None → no quantization."""
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
config = DiffusionMXFP8Config()
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
norm_layer = torch.nn.LayerNorm(64)
assert config.get_quant_method(norm_layer, "blocks.0.norm1") is None
# ---------------------------------------------------------------------------
# MXFPLinearMethodBase.apply() — reshape skeleton (CPU, no NPU)
# ---------------------------------------------------------------------------
def test_apply_reshape_skeleton():
"""apply() must flatten batch dims → _apply_inner → restore original leading dims."""
from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase
OUT_FEATURES = 4
class _StubMethod(MXFPLinearMethodBase):
def create_weights(
self,
layer,
input_size_per_partition,
output_partition_sizes,
input_size,
output_size,
params_dtype,
**extra_weight_attrs,
):
pass
def _quantize_activation(self, x):
return x, None
def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype):
return torch.zeros(x_q.shape[0], OUT_FEATURES, dtype=ori_dtype)
method = _StubMethod()
x = torch.randn(2, 3, 8) # (batch=2, seq=3, K=8)
out = method.apply(None, x)
assert out.shape == (2, 3, OUT_FEATURES)
def test_apply_reshape_with_bias():
"""apply() must pass bias through to _apply_inner unchanged."""
from vllm_omni.quantization.mxfp8_config import MXFPLinearMethodBase
received_bias = []
class _StubMethod(MXFPLinearMethodBase):
def create_weights(
self,
layer,
input_size_per_partition,
output_partition_sizes,
input_size,
output_size,
params_dtype,
**extra_weight_attrs,
):
pass
def _quantize_activation(self, x):
return x, None
def _quant_matmul(self, x_q, x_scale, layer, bias, ori_dtype):
received_bias.append(bias)
return torch.zeros(x_q.shape[0], 4, dtype=ori_dtype)
method = _StubMethod()
bias = torch.zeros(4)
method.apply(None, torch.randn(2, 8), bias=bias)
assert received_bias[0] is bias
# ---------------------------------------------------------------------------
# process_weights_after_loading shape arithmetic (pure torch, no NPU ops)
#
# These tests replicate the CPU-safe portions of process_weights_after_loading
# to guard the key layout contract: (N,K) weight → (K,N) and (N,S) scale →
# (S/2,N,2). They do NOT call NPU ops; they test only the math.
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"n, k",
[(64, 128), (32, 64), (16, 96)],
ids=["64x128", "32x64", "16x96"],
)
def test_weight_transpose_contract(n: int, k: int):
"""Weight must be transposed from (N, K) to (K, N) and be contiguous."""
w = torch.zeros(n, k, dtype=torch.uint8)
w = w.transpose(0, 1).contiguous()
assert w.shape == (k, n)
assert w.is_contiguous()
@pytest.mark.parametrize(
"n, k_groups, expected_groups",
[
(64, 4, 2), # even K_groups — no padding
(64, 3, 2), # odd K_groups — padded to 4
(32, 1, 1), # odd K_groups — padded to 2
],
ids=["even", "odd-to-4", "odd-to-2"],
)
def test_weight_scale_reshape_contract(n: int, k_groups: int, expected_groups: int):
"""Scale must be reshaped from (N, K_groups) to (K_groups_even//2, N, 2).
Odd K_groups must be padded to even before the reshape.
"""
s = torch.zeros(n, k_groups, dtype=torch.uint8)
if k_groups % 2 == 1:
s = torch.cat([s, torch.zeros(n, 1, dtype=s.dtype)], dim=1)
k_groups += 1
s = s.reshape(n, k_groups // 2, 2).transpose(0, 1).contiguous()
assert s.shape == (expected_groups, n, 2)
assert s.is_contiguous()
def test_num_groups_formula():
"""K_groups formula: ceil(K / 32) — spot-check boundary values."""
assert (31 + 31) // 32 == 1 # K=31 → 1 group
assert (32 + 31) // 32 == 1 # K=32 → 1 group
assert (33 + 31) // 32 == 2 # K=33 → 2 groups
assert (128 + 31) // 32 == 4 # K=128 → 4 groups (even)
assert (96 + 31) // 32 == 3 # K=96 → 3 groups (odd → needs padding)
# ---------------------------------------------------------------------------
# SUPPORTED_QUANTIZATION_METHODS
# ---------------------------------------------------------------------------
def test_supported_methods_include_mxfp8():
from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
assert "mxfp8" in SUPPORTED_QUANTIZATION_METHODS
tests/diffusion/quantization/test_mxfp8_key_remap.py
+214
View File
@@ -0,0 +1,214 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for merge_mxfp8_checkpoint.py key-remapping helpers and model metadata.
These are pure-Python unit tests that exercise the transformation functions
without loading any actual checkpoint files or requiring NPU hardware.
Key contracts verified:
- SUPPORTED_MODEL_TYPES includes Wan2.2-TI2V-5B (MXFP8 supports it; MXFP4 does not)
- MXFP8_QUANT_CONFIG structure matches what TransformerConfig.from_dict() reads
- _remap_keys correctly translates msModelSlim naming → Diffusers naming
- _get_transformer_dirs routes cascade vs single-transformer models
- _get_quant_subdir maps high/low noise subdirs for cascade models
"""
import pathlib
import pytest
import torch
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
# ---------------------------------------------------------------------------
# SUPPORTED_MODEL_TYPES and CASCADE_MODEL_TYPES
# ---------------------------------------------------------------------------
def test_supported_model_types_includes_all_wan22():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES
assert "Wan2.2-T2V-A14B" in SUPPORTED_MODEL_TYPES
assert "Wan2.2-I2V-A14B" in SUPPORTED_MODEL_TYPES
def test_supported_model_types_includes_ti2v_5b():
"""MXFP8 supports TI2V-5B (contrast: MXFP4 explicitly excludes it)."""
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import SUPPORTED_MODEL_TYPES
assert "Wan2.2-TI2V-5B" in SUPPORTED_MODEL_TYPES
def test_cascade_model_types_excludes_ti2v():
"""TI2V-5B is a single-transformer model — must not be in CASCADE_MODEL_TYPES."""
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import CASCADE_MODEL_TYPES
assert "Wan2.2-TI2V-5B" not in CASCADE_MODEL_TYPES
assert "Wan2.2-T2V-A14B" in CASCADE_MODEL_TYPES
assert "Wan2.2-I2V-A14B" in CASCADE_MODEL_TYPES
# ---------------------------------------------------------------------------
# MXFP8_QUANT_CONFIG — auto-detection contract
# ---------------------------------------------------------------------------
def test_mxfp8_quant_config_has_required_keys():
"""MXFP8_QUANT_CONFIG must carry exactly the keys that auto-detection reads:
quant_method (selects DiffusionMXFP8Config) and is_checkpoint_mxfp8_serialized
(selects NPUMxfp8LinearMethod over the online path)."""
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import MXFP8_QUANT_CONFIG
assert MXFP8_QUANT_CONFIG["quant_method"] == "mxfp8"
assert MXFP8_QUANT_CONFIG["is_checkpoint_mxfp8_serialized"] is True
# ---------------------------------------------------------------------------
# _get_transformer_dirs
# ---------------------------------------------------------------------------
def test_get_transformer_dirs_cascade():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs
assert _get_transformer_dirs("Wan2.2-T2V-A14B") == ["transformer", "transformer_2"]
assert _get_transformer_dirs("Wan2.2-I2V-A14B") == ["transformer", "transformer_2"]
def test_get_transformer_dirs_single():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_transformer_dirs
assert _get_transformer_dirs("Wan2.2-TI2V-5B") == ["transformer"]
# ---------------------------------------------------------------------------
# _get_quant_subdir
# ---------------------------------------------------------------------------
def test_get_quant_subdir_cascade_high_noise():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
base = pathlib.Path("/quant")
result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer")
assert result == base / "high_noise_model"
def test_get_quant_subdir_cascade_low_noise():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
base = pathlib.Path("/quant")
result = _get_quant_subdir("Wan2.2-T2V-A14B", base, "transformer_2")
assert result == base / "low_noise_model"
def test_get_quant_subdir_non_cascade():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _get_quant_subdir
base = pathlib.Path("/quant")
result = _get_quant_subdir("Wan2.2-TI2V-5B", base, "transformer")
assert result == base
# ---------------------------------------------------------------------------
# _remap_keys — msModelSlim naming → Diffusers naming
# ---------------------------------------------------------------------------
def test_remap_keys_self_attn_q():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
state = {"blocks.0.self_attn.q.weight": torch.zeros(1)}
meta = {"blocks.0.self_attn.q.weight": "W8A8_MXFP8"}
new_state, new_meta = _remap_keys(state, meta)
assert "blocks.0.attn1.to_q.weight" in new_state
assert new_meta.get("blocks.0.attn1.to_q.weight") == "W8A8_MXFP8"
def test_remap_keys_self_attn_all_heads():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
pairs = {
"self_attn.q": "attn1.to_q",
"self_attn.k": "attn1.to_k",
"self_attn.v": "attn1.to_v",
"self_attn.o": "attn1.to_out.0",
}
for src_part, dst_part in pairs.items():
src_key = f"blocks.0.{src_part}.weight"
state = {src_key: torch.zeros(1)}
new_state, _ = _remap_keys(state, {})
expected = f"blocks.0.{dst_part}.weight"
assert expected in new_state, f"{src_key} → expected {expected}, got {list(new_state)}"
def test_remap_keys_ffn():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
state = {
"blocks.1.ffn.0.weight": torch.zeros(1),
"blocks.1.ffn.2.weight": torch.zeros(1),
}
new_state, _ = _remap_keys(state, {})
assert "blocks.1.ffn.net.0.proj.weight" in new_state
assert "blocks.1.ffn.net.2.weight" in new_state
def test_remap_keys_norm_order_swap():
"""norm2↔norm3 swap: msModelSlim uses norm1/norm3/norm2 order,
Diffusers uses norm1/norm2/norm3."""
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
state = {
"blocks.0.norm2.weight": torch.zeros(1),
"blocks.0.norm3.weight": torch.zeros(1),
}
new_state, _ = _remap_keys(state, {})
# norm2 → norm3 and norm3 → norm2
assert "blocks.0.norm3.weight" in new_state
assert "blocks.0.norm2.weight" in new_state
# Both must be present (swap, not collapse)
assert len([k for k in new_state if "norm" in k and "norm_q" not in k and "norm_k" not in k]) == 2
def test_remap_keys_head():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
state = {"head.head.weight": torch.zeros(1)}
new_state, _ = _remap_keys(state, {})
assert "proj_out.weight" in new_state
def test_remap_keys_cross_attn():
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
state = {"blocks.0.cross_attn.q.weight": torch.zeros(1)}
new_state, _ = _remap_keys(state, {})
assert "blocks.0.attn2.to_q.weight" in new_state
def test_remap_keys_meta_only_mapped_for_existing_state_keys():
"""quant_meta entries are only emitted for keys present in state_dict."""
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
state = {"blocks.0.self_attn.q.weight": torch.zeros(1)}
# meta has an extra key not in state_dict
meta = {
"blocks.0.self_attn.q.weight": "W8A8_MXFP8",
"blocks.0.self_attn.q.weight_scale": "W8A8_MXFP8",
}
_, new_meta = _remap_keys(state, meta)
assert "blocks.0.attn1.to_q.weight" in new_meta
# weight_scale was in meta but not state_dict → must be absent
assert "blocks.0.attn1.to_q.weight_scale" not in new_meta
def test_remap_keys_preserves_tensors():
"""Tensor values must survive the key rename unchanged."""
from vllm_omni.quantization.tools.merge_mxfp8_checkpoint import _remap_keys
t = torch.randn(4, 8)
state = {"blocks.0.self_attn.q.weight": t}
new_state, _ = _remap_keys(state, {})
assert torch.equal(new_state["blocks.0.attn1.to_q.weight"], t)
tests/diffusion/quantization/test_quantization_fp8.py
+98
View File
@@ -109,6 +109,85 @@ def _generate_single_stage_image(
return images, peak_mem
def _generate_single_stage_video(
model: str,
quantization: str | None = None,
height: int = 256,
width: int = 256,
num_frames: int = 25,
num_inference_steps: int = 8,
guidance_scale: float = 4.0,
seed: int = 42,
prompt: str = "A serene lakeside sunrise with mist over the water",
**extra_omni_kwargs: Any,
) -> tuple[int, float]:
"""Generate a t2v output with a single-stage diffusion model.
Returns (num_frames_produced, peak_memory_gib)
"""
omni_kwargs: dict[str, Any] = dict(extra_omni_kwargs)
if quantization:
omni_kwargs["quantization"] = quantization
with OmniRunner(model, **omni_kwargs) as runner:
torch.accelerator.reset_peak_memory_stats()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
outputs = runner.omni.generate(
{"prompt": prompt, "negative_prompt": ""},
OmniDiffusionSamplingParams(
height=height,
width=width,
num_frames=num_frames,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
generator=generator,
),
)
first = outputs[0]
# Unwrap pipeline-style outputs (multi-stage / OmniRequestOutput.request_output).
frames: Any = None
if hasattr(first, "request_output") and isinstance(first.request_output, list):
inner = first.request_output[0]
if isinstance(inner, OmniRequestOutput) and inner.images:
frames = inner.images[0]
if frames is None and hasattr(first, "images") and first.images:
frames = first.images[0]
assert frames is not None, "No video frames returned from generate()"
# LTX-2 (audio+video) may surface (video, audio) tuples or {"video": ...} dicts.
if isinstance(frames, dict):
frames = frames.get("video") or frames.get("frames")
elif isinstance(frames, tuple) and len(frames) == 2:
frames = frames[0]
assert frames is not None, "Could not extract video frames from output"
if isinstance(frames, torch.Tensor):
video = frames.detach().cpu()
if video.dim() == 5:
video = video[0]
if video.dim() == 4 and video.shape[0] in (3, 4):
video = video.permute(1, 2, 3, 0)
num_frames_produced = int(video.shape[0])
else:
import numpy as np
arr = np.asarray(frames)
if arr.ndim == 5:
arr = arr[0]
num_frames_produced = int(arr.shape[0])
peak_mem_mb = getattr(first, "peak_memory_mb", None)
if peak_mem_mb:
peak_mem = float(peak_mem_mb) / 1024.0
else:
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
return num_frames_produced, peak_mem
def _generate_bagel_image(
quantization_config: str | None = None,
num_inference_steps: int = 15,
@@ -277,6 +356,25 @@ def test_single_stage_flux_fp8_uses_less_memory():
assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
@hardware_test(res={"cuda": "H100"})
def test_single_stage_ltx2_fp8_uses_less_memory():
"""FP8 should use less peak memory than BF16 for LTX-2."""
_, mem_bf16 = _generate_single_stage_video(
model="Lightricks/LTX-2",
quantization=None,
)
torch.accelerator.empty_cache()
_, mem_fp8 = _generate_single_stage_video(
model="Lightricks/LTX-2",
quantization="fp8",
)
print(f"LTX-2 BF16 peak memory: {mem_bf16:.2f} GiB")
print(f"LTX-2 FP8 peak memory: {mem_fp8:.2f} GiB")
assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
# ─── Multi-stage model tests (BAGEL) ─────────────────────────────────────────
tests/diffusion/quantization/test_quantization_quality.py
+34 -6
View File
@@ -143,6 +143,18 @@ QUALITY_CONFIGS = [
seed=142,
num_inference_steps=20,
),
QualityTestConfig(
id="fp8_ltx2",
model="Lightricks/LTX-2",
quantization="fp8",
task="t2v",
prompt="A serene lakeside sunrise with mist over the water",
max_lpips=0.10,
height=256,
width=256,
num_frames=25,
num_inference_steps=8,
),
]
@@ -261,10 +273,19 @@ def _generate_video(omni, config: QualityTestConfig):
else:
frames = inner
elif hasattr(first, "images") and first.images:
frames = first.images
frames = first.images[0]
else:
raise ValueError("Could not extract video frames from output.")
# LTX-2 (audio+video) may surface (video, audio) tuples or {"video": ...} dicts
if isinstance(frames, dict):
frames = frames.get("video") or frames.get("frames")
elif isinstance(frames, tuple) and len(frames) == 2:
frames = frames[0]
if frames is None:
raise ValueError("Could not extract video frames from output.")
if isinstance(frames, torch.Tensor):
video = frames.detach().cpu()
if video.dim() == 5:
@@ -275,7 +296,11 @@ def _generate_video(omni, config: QualityTestConfig):
video = video.clamp(-1, 1) * 0.5 + 0.5
return video.float().numpy(), peak_mem
return np.asarray(frames), peak_mem
frames_array = np.asarray(frames)
if frames_array.ndim == 5:
# strip the leading batch dim
frames_array = frames_array[0]
return frames_array, peak_mem
def _compute_lpips(baseline, quantized, task: str) -> float:
@@ -318,8 +343,7 @@ def _compute_psnr_and_mae(baseline, quantized, task: str) -> tuple[float, float]
return psnr, mae
def _unload(omni):
del omni
def _free_gpu_memory():
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
@@ -395,7 +419,9 @@ def test_quantization_quality(config: QualityTestConfig):
# --- BF16 baseline ---
omni_bl = Omni(model=config.baseline_ref())
baseline_out, bl_mem = generate_fn(omni_bl, config)
_unload(omni_bl)
omni_bl.shutdown()
del omni_bl
_free_gpu_memory()
_maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out)
# --- Quantized ---
@@ -405,7 +431,9 @@ def test_quantization_quality(config: QualityTestConfig):
else:
omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization)
quant_out, qt_mem = generate_fn(omni_qt, config)
_unload(omni_qt)
omni_qt.shutdown()
del omni_qt
_free_gpu_memory()
_maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out)
# --- Similarity metrics ---
tests/diffusion/quantization/test_trajectory_similarity_tool.py
+126
View File
@@ -0,0 +1,126 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
import sys
from types import SimpleNamespace
import numpy as np
import pytest
import torch
from PIL import Image
from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import (
VariantRun,
_build_variant_config,
_request_peak_memory_mb,
_run_summary,
compute_tensor_metrics,
compute_uint8_image_metrics,
metric_guidance,
parse_args,
summarize_output_image_metrics,
)
def test_compute_tensor_metrics_identical_tensors():
metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3))
assert metrics["cosine_similarity"] == pytest.approx(1.0)
assert metrics["mae"] == 0.0
assert metrics["mse"] == 0.0
assert metrics["rmse"] == 0.0
assert metrics["max_abs"] == 0.0
assert metrics["l2"] == 0.0
assert metrics["relative_l2"] == 0.0
def test_compute_uint8_image_metrics_adds_psnr():
lhs = np.zeros((2, 2, 3), dtype=np.uint8)
rhs = np.zeros((2, 2, 3), dtype=np.uint8)
metrics = compute_uint8_image_metrics(lhs, rhs)
assert math.isinf(metrics["psnr_db"])
def test_summarize_output_image_metrics_stacks_pil_images():
reference = [Image.fromarray(np.zeros((2, 2, 3), dtype=np.uint8))]
candidate = [Image.fromarray(np.ones((2, 2, 3), dtype=np.uint8))]
summary = summarize_output_image_metrics(reference, candidate)
assert summary["num_images"] == 1
assert summary["image0_metrics"]["mae"] == 1.0
assert summary["all_images_metrics"]["mse"] == 1.0
def test_run_summary_reports_worker_peak_memory():
summary = _run_summary(
VariantRun(
label="candidate",
result=object(),
generation_times_s=[1.0, 3.0],
peak_memory_mb=[100.0, 150.0],
)
)
assert summary["peak_memory_mb"] == 150.0
assert summary["avg_peak_memory_mb"] == 125.0
assert summary["max_peak_memory_mb"] == 150.0
assert summary["per_run_peak_memory_mb"] == [100.0, 150.0]
def test_request_peak_memory_prefers_inner_diffusion_output():
inner = type("InnerOutput", (), {"peak_memory_mb": 321.0})()
outer = type("OuterOutput", (), {"request_output": inner, "peak_memory_mb": 123.0})()
assert _request_peak_memory_mb(outer) == 321.0
def test_metric_guidance_describes_thresholds():
guidance = metric_guidance()
assert "cosine_similarity" in guidance["descriptions"]
assert guidance["recommended_thresholds"]["output_images_or_frames_uint8"]["psnr_db"]["recommended_min"] == 20.0
assert (
guidance["recommended_thresholds"]["performance"]["max_peak_memory_ratio_candidate_over_reference"][
"recommended_max"
]
== 1.00
)
def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantization():
args = SimpleNamespace(
model="Qwen/Qwen-Image",
candidate_model="Qwen/Qwen-Image-FP8",
candidate_quantization=None,
candidate_quantization_config_json=None,
candidate_ignored_layers=None,
ignored_layers=None,
candidate_gguf_model=None,
candidate_load_format="default",
)
config = _build_variant_config(args, "candidate")
assert config.model == "Qwen/Qwen-Image-FP8"
assert config.quantization is None
assert config.quantization_config is None
def test_step_execution_defaults_to_false(monkeypatch):
monkeypatch.setattr(
sys,
"argv",
[
"compare_diffusion_trajectory_similarity.py",
"--output-json",
"result.json",
],
)
args = parse_args()
assert args.step_execution is False
tests/diffusion/test_diffusion_engine_rpc_routing.py
+3
View File
@@ -136,6 +136,7 @@ def _make_engine_with_loop(
calls and a real ``RequestScheduler``.
"""
engine = DiffusionEngine.__new__(DiffusionEngine)
engine._closed = False
engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
sched = RequestScheduler()
@@ -371,6 +372,7 @@ def test_collective_rpc_before_loop_starts_calls_executor_directly():
on the caller's thread without enqueueing.
"""
engine = DiffusionEngine.__new__(DiffusionEngine)
engine._closed = False
engine._loop_started = False
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
@@ -391,6 +393,7 @@ def test_collective_rpc_before_loop_starts_serializes_concurrent_callers():
serialize them so they cannot race on the shared executor MQ pair.
"""
engine = DiffusionEngine.__new__(DiffusionEngine)
engine._closed = False
engine._loop_started = False
engine._rpc_lock = threading.RLock()
engine._cv = threading.Condition(engine._rpc_lock)
tests/diffusion/test_diffusion_scheduler.py
+156
View File
@@ -153,6 +153,54 @@ class _StubScheduler(SchedulerInterface):
return None
class TestGetSamplingParamsKey:
"""Pure-function tests for the batch-compatibility key builder."""
@staticmethod
def _make(lora_int_id: int | None = None, lora_scale: float = 1.0) -> OmniDiffusionRequest:
from vllm_omni.lora.request import LoRARequest
sp = OmniDiffusionSamplingParams(num_inference_steps=2)
if lora_int_id is not None:
sp.lora_request = LoRARequest(
lora_name=f"adapter-{lora_int_id}",
lora_int_id=lora_int_id,
lora_path=f"/tmp/lora-{lora_int_id}",
)
sp.lora_scale = lora_scale
return OmniDiffusionRequest(
prompts=["prompt"],
sampling_params=sp,
request_ids=[f"req-{lora_int_id}-{lora_scale}"],
)
def test_distinguishes_lora_id(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
assert get_sampling_params_key(self._make(lora_int_id=1)) != get_sampling_params_key(self._make(lora_int_id=2))
def test_distinguishes_lora_scale(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
assert get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5)) != get_sampling_params_key(
self._make(lora_int_id=1, lora_scale=1.0)
)
def test_treats_no_lora_as_distinct_bucket(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
assert get_sampling_params_key(self._make(lora_int_id=None)) != get_sampling_params_key(
self._make(lora_int_id=1)
)
def test_equal_for_same_lora_identity(self) -> None:
from vllm_omni.diffusion.sched.base_scheduler import get_sampling_params_key
a = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
b = get_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
assert a == b
class TestRequestScheduler:
def setup_method(self) -> None:
self.scheduler: RequestScheduler = RequestScheduler()
@@ -728,6 +776,114 @@ class TestStepScheduler:
assert second.num_running_reqs == 1
assert second.num_waiting_reqs == 1
def test_step_batch_co_schedules_requests_sharing_lora(self) -> None:
"""Multiple requests with the same LoRA (id + scale) co-batch."""
from vllm_omni.lora.request import LoRARequest
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=3))
lora = LoRARequest(lora_name="adapter", lora_int_id=42, lora_path="/tmp/lora")
def _with_lora(req_id: str) -> OmniDiffusionRequest:
sp = OmniDiffusionSamplingParams(num_inference_steps=4)
sp.lora_request = lora
sp.lora_scale = 0.5
return _make_step_request(req_id, sampling_params=sp)
req_a = scheduler.add_request(_with_lora("a"))
req_b = scheduler.add_request(_with_lora("b"))
req_c = scheduler.add_request(_with_lora("c"))
sched_output = scheduler.schedule()
assert _new_ids(sched_output) == [req_a, req_b, req_c]
assert sched_output.num_running_reqs == 3
assert sched_output.num_waiting_reqs == 0
def test_step_batch_separates_requests_with_different_lora_ids(self) -> None:
"""Different LoRA adapters → distinct batches admitted in FIFO order."""
from vllm_omni.lora.request import LoRARequest
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=4))
lora_a = LoRARequest(lora_name="adapter-A", lora_int_id=1, lora_path="/tmp/lora-a")
lora_b = LoRARequest(lora_name="adapter-B", lora_int_id=2, lora_path="/tmp/lora-b")
def _build(req_id: str, lora: LoRARequest) -> OmniDiffusionRequest:
sp = OmniDiffusionSamplingParams(num_inference_steps=2)
sp.lora_request = lora
return _make_step_request(req_id, sampling_params=sp)
req_a1 = scheduler.add_request(_build("a1", lora_a))
req_b1 = scheduler.add_request(_build("b1", lora_b))
req_a2 = scheduler.add_request(_build("a2", lora_a))
# Strict FIFO admission: a1 starts; b1 (different LoRA) blocks the
# queue head, so a2 (compatible with a1) is *not* skipped ahead.
first = scheduler.schedule()
assert _new_ids(first) == [req_a1]
assert first.num_waiting_reqs == 2
# Drain a1 → b1 becomes head-of-line and is admitted with its LoRA.
scheduler.update_from_output(first, _make_step_output(req_a1, step_index=2, finished=True))
second = scheduler.schedule()
assert _new_ids(second) == [req_b1]
assert second.num_waiting_reqs == 1
# Drain b1 → a2 is admitted next; LoRA-A is re-activated for it.
scheduler.update_from_output(second, _make_step_output(req_b1, step_index=2, finished=True))
third = scheduler.schedule()
assert _new_ids(third) == [req_a2]
assert third.num_waiting_reqs == 0
def test_step_batch_separates_requests_with_different_lora_scale(self) -> None:
"""Same adapter id but different scales → still separate batches."""
from vllm_omni.lora.request import LoRARequest
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=4))
lora = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora")
def _build(req_id: str, scale: float) -> OmniDiffusionRequest:
sp = OmniDiffusionSamplingParams(num_inference_steps=2)
sp.lora_request = lora
sp.lora_scale = scale
return _make_step_request(req_id, sampling_params=sp)
req_full = scheduler.add_request(_build("full", 1.0))
req_half = scheduler.add_request(_build("half", 0.5))
sched_output = scheduler.schedule()
admitted = _new_ids(sched_output)
assert admitted == [req_full]
assert req_half not in admitted
assert sched_output.num_waiting_reqs == 1
def test_step_batch_separates_lora_from_no_lora(self) -> None:
"""A LoRA request and a no-LoRA request do not share a batch."""
from vllm_omni.lora.request import LoRARequest
scheduler = StepScheduler()
scheduler.initialize(SimpleNamespace(max_num_seqs=4))
lora = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora")
sp_with = OmniDiffusionSamplingParams(num_inference_steps=2)
sp_with.lora_request = lora
req_with = scheduler.add_request(_make_step_request("with", sampling_params=sp_with))
req_without = scheduler.add_request(_make_step_request("without", num_inference_steps=2))
sched_output = scheduler.schedule()
admitted = _new_ids(sched_output)
assert admitted == [req_with]
assert req_without not in admitted
assert sched_output.num_waiting_reqs == 1
def test_preempt_request_preserves_step_index(self) -> None:
request = _make_step_request("preempt", num_inference_steps=3)
req_id = self.scheduler.add_request(request)
tests/diffusion/test_diffusion_step_pipeline.py
+118 -48
View File
@@ -285,6 +285,19 @@ def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_id
)
def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
"""Scheduler output for a homogeneous batch (one NewRequestData per req)."""
new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs]
return DiffusionSchedulerOutput(
step_id=step_id,
scheduled_new_reqs=new_reqs,
scheduled_cached_reqs=CachedRequestData.make_empty(),
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
num_running_reqs=len(new_reqs),
num_waiting_reqs=0,
)
def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
return DiffusionSchedulerOutput(
step_id=step_id,
@@ -489,71 +502,128 @@ class TestRunner:
DiffusionModelRunner.load_model(runner)
class _RecordingLoRAManager:
def __init__(self) -> None:
self.calls: list[tuple[object | None, float]] = []
def set_active_adapter(self, adapter, scale: float = 1.0) -> None:
self.calls.append((adapter, scale))
def _make_step_worker(lora_manager=None, *, expected_output=None):
"""Build a bare DiffusionWorker primed for execute_stepwise tests."""
worker = object.__new__(DiffusionWorker)
worker.lora_manager = lora_manager
worker._step_lora_state = {}
output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1")
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
return worker
@pytest.mark.cpu
class TestWorker:
"""DiffusionWorker.execute_stepwise"""
def test_delegates_to_model_runner(self):
worker = object.__new__(DiffusionWorker)
expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
scheduler_output = SimpleNamespace(
scheduled_new_reqs=[
SimpleNamespace(
req=SimpleNamespace(
sampling_params=SimpleNamespace(lora_request=None),
)
)
]
)
worker.lora_manager = None
worker.model_runner = SimpleNamespace(
execute_stepwise=lambda arg: expected if arg is scheduler_output else None
)
worker = _make_step_worker(expected_output=expected)
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
assert output is expected
def test_clears_active_lora_before_stepwise_execution(self):
worker = object.__new__(DiffusionWorker)
scheduler_output = SimpleNamespace(
scheduled_new_reqs=[
SimpleNamespace(
req=SimpleNamespace(
sampling_params=SimpleNamespace(lora_request=None),
)
)
]
)
calls: list[object | None] = []
class _FakeLoRAManager:
def set_active_adapter(self, adapter):
calls.append(adapter)
worker.lora_manager = _FakeLoRAManager()
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
def test_deactivates_lora_when_request_has_no_adapter(self):
manager = _RecordingLoRAManager()
worker = _make_step_worker(lora_manager=manager)
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
DiffusionWorker.execute_stepwise(worker, scheduler_output)
assert calls == [None]
assert manager.calls == [(None, 1.0)]
def test_rejects_lora_requests_in_step_mode(self):
worker = object.__new__(DiffusionWorker)
scheduler_output = SimpleNamespace(
scheduled_new_reqs=[
SimpleNamespace(
req=SimpleNamespace(
sampling_params=SimpleNamespace(lora_request=object()),
)
)
]
def test_activates_lora_for_step_requests(self):
from vllm_omni.lora.request import LoRARequest
lora_request = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora")
request = _make_engine_request("req-1")
request.sampling_params.lora_request = lora_request
request.sampling_params.lora_scale = 0.75
manager = _RecordingLoRAManager()
worker = _make_step_worker(lora_manager=manager)
scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
DiffusionWorker.execute_stepwise(worker, scheduler_output)
assert manager.calls == [(lora_request, 0.75)]
def test_recovers_lora_for_cached_step_requests(self):
from vllm_omni.lora.request import LoRARequest
lora_request = LoRARequest(lora_name="adapter", lora_int_id=11, lora_path="/tmp/lora")
request = _make_engine_request("req-1")
request.sampling_params.lora_request = lora_request
request.sampling_params.lora_scale = 0.5
manager = _RecordingLoRAManager()
worker = _make_step_worker(lora_manager=manager)
first = _make_scheduler_output(request, sched_req_id="req-1")
second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1)
DiffusionWorker.execute_stepwise(worker, first)
DiffusionWorker.execute_stepwise(worker, second)
assert manager.calls == [(lora_request, 0.5), (lora_request, 0.5)]
def test_activates_single_lora_for_homogeneous_batch(self):
"""Multiple requests sharing the same LoRA → exactly one activation,
and every request id is registered in ``_step_lora_state``."""
from vllm_omni.lora.request import LoRARequest
lora_request = LoRARequest(lora_name="adapter", lora_int_id=9, lora_path="/tmp/lora")
reqs = []
for rid in ("req-1", "req-2", "req-3"):
r = _make_engine_request(rid)
r.sampling_params.lora_request = lora_request
r.sampling_params.lora_scale = 0.6
reqs.append(r)
manager = _RecordingLoRAManager()
worker = _make_step_worker(lora_manager=manager)
scheduler_output = _make_batch_scheduler_output(reqs)
DiffusionWorker.execute_stepwise(worker, scheduler_output)
assert manager.calls == [(lora_request, 0.6)]
assert set(worker._step_lora_state) == {"req-1", "req-2", "req-3"}
for entry in worker._step_lora_state.values():
assert entry == (lora_request, 0.6)
def test_evicts_step_lora_state_for_finished_requests(self):
from vllm_omni.lora.request import LoRARequest
lora_request = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora")
finishing = _make_engine_request("req-1")
finishing.sampling_params.lora_request = lora_request
next_request = _make_engine_request("req-2")
next_request.sampling_params.lora_request = lora_request
worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
first = _make_scheduler_output(finishing, sched_req_id="req-1")
next_batch = _make_scheduler_output(
next_request,
sched_req_id="req-2",
step_id=1,
finished_req_ids={"req-1"},
)
worker.lora_manager = None
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1"))
with pytest.raises(ValueError, match="does not support LoRA"):
DiffusionWorker.execute_stepwise(worker, scheduler_output)
DiffusionWorker.execute_stepwise(worker, first)
assert "req-1" in worker._step_lora_state
DiffusionWorker.execute_stepwise(worker, next_batch)
assert "req-1" not in worker._step_lora_state
assert worker._step_lora_state == {"req-2": (lora_request, 1.0)}
@pytest.mark.cpu
tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py
+69 -1
View File
@@ -15,8 +15,9 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c
BufferAllocator,
ManagedBuffer,
)
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import KVCacheTransferData
# All tests in this file are pure-CPU unit tests for the memory allocator.
# Most tests in this file are pure-CPU unit tests; CUDA smoke tests are skip-guarded.
pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model]
@@ -225,3 +226,70 @@ class TestManagedBuffer:
# Space should be reusable
new_offset = self.allocator.alloc(128)
assert new_offset == offset
def test_to_bytes_cpu_pool(self):
"""Smoke-test Mooncake's ManagedBuffer byte extraction on CPU pools."""
expected = bytes(range(64))
offset = self.allocator.alloc(len(expected))
buf = ManagedBuffer(self.allocator, offset, len(expected), self.pool)
self.pool[offset : offset + len(expected)] = torch.tensor(list(expected), dtype=torch.uint8)
try:
assert buf.to_bytes() == expected
finally:
buf.release()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for GPU pool smoke test")
def test_to_bytes_cuda_pool(self):
"""Smoke-test Mooncake's existing CUDA pool D2H byte path."""
expected = bytes(range(64))
allocator = BufferAllocator(total_size=4096, alignment=64)
pool = torch.zeros(4096, dtype=torch.uint8, device="cuda:0")
offset = allocator.alloc(len(expected))
buf = ManagedBuffer(allocator, offset, len(expected), pool)
pool[offset : offset + len(expected)].copy_(torch.tensor(list(expected), dtype=torch.uint8, device="cuda:0"))
try:
assert buf.to_bytes() == expected
finally:
buf.release()
def _make_kv_payload(device: str | torch.device = "cpu") -> tuple[KVCacheTransferData, torch.Tensor, torch.Tensor]:
key_tensor = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
value_tensor = key_tensor + 100
payload = KVCacheTransferData(
request_id="moon-smoke",
layer_blocks={"key_cache": [key_tensor], "value_cache": [value_tensor]},
block_ids=[1],
metadata={"seq_len": 3},
)
return payload, key_tensor, value_tensor
class TestMooncakePackedPayloadSmoke:
"""Smoke tests for Mooncake raw payload helpers without real Mooncake/RDMA."""
def test_load_header_from_cpu_tensor(self):
payload, _, _ = _make_kv_payload()
raw = payload.to_bytes()
packed = torch.tensor(list(raw), dtype=torch.uint8)
header, data_start = KVCacheTransferData._load_header_from_tensor(packed)
assert header["rid"] == "moon-smoke"
assert data_start == 4 + int.from_bytes(raw[:4], "big")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for packed CUDA smoke test")
def test_from_bytes_device_cuda_round_trip(self):
payload, key_tensor, value_tensor = _make_kv_payload(device="cuda:0")
packed = payload.to_gpu_tensor()
header, data_start = KVCacheTransferData._load_header_from_tensor(packed)
decoded = KVCacheTransferData.from_bytes_device(packed)
assert header["rid"] == "moon-smoke"
assert data_start == 4 + int.from_bytes(payload.to_bytes()[:4], "big")
assert decoded["layer_blocks"]["key_cache"][0].is_cuda
assert decoded["layer_blocks"]["value_cache"][0].is_cuda
assert torch.equal(decoded["layer_blocks"]["key_cache"][0].cpu(), key_tensor.cpu())
assert torch.equal(decoded["layer_blocks"]["value_cache"][0].cpu(), value_tensor.cpu())
tests/distributed/omni_coordinator/test_load_balancer.py
+50 -50
View File
@@ -6,18 +6,18 @@ from time import time
import pytest
from vllm_omni.distributed.omni_coordinator import (
InstanceInfo,
LeastQueueLengthBalancer,
RandomBalancer,
ReplicaInfo,
ReplicaStatus,
RoundRobinBalancer,
StageStatus,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_load_balancer_select_returns_valid_index():
"""Verify RandomBalancer.select() returns a valid index for instances."""
"""Verify RandomBalancer.select() returns a valid index for replicas."""
# Task structure mirrors async_omni; RandomBalancer ignores task contents.
task: dict = {
"request_id": "test",
@@ -26,30 +26,30 @@ def test_load_balancer_select_returns_valid_index():
}
now = time()
instances = [
InstanceInfo(
replicas = [
ReplicaInfo(
input_addr="tcp://host:10001",
output_addr="tcp://host:10001-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=0,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10002",
output_addr="tcp://host:10002-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=1,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10003",
output_addr="tcp://host:10003-out",
stage_id=1,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=2,
last_heartbeat=now,
registered_at=now,
@@ -58,38 +58,38 @@ def test_load_balancer_select_returns_valid_index():
balancer = RandomBalancer()
index = balancer.select(task, instances)
index = balancer.select(task, replicas)
assert isinstance(index, int)
assert 0 <= index < len(instances)
assert 0 <= index < len(replicas)
def test_round_robin_balancer_cycles_instances():
def test_round_robin_balancer_cycles_replicas():
now = time()
instances = [
InstanceInfo(
replicas = [
ReplicaInfo(
input_addr="tcp://host:10001",
output_addr="tcp://host:10001-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=2,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10002",
output_addr="tcp://host:10002-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=1,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10003",
output_addr="tcp://host:10003-out",
stage_id=1,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=0,
last_heartbeat=now,
registered_at=now,
@@ -97,35 +97,35 @@ def test_round_robin_balancer_cycles_instances():
]
balancer = RoundRobinBalancer()
results = [balancer.select({}, instances) for _ in range(5)]
results = [balancer.select({}, replicas) for _ in range(5)]
# Default start_index=0 => 0,1,2,0,1
assert results == [0, 1, 2, 0, 1]
def test_round_robin_balancer_empty_instances_raises():
with pytest.raises(ValueError, match="instances must not be empty"):
def test_round_robin_balancer_empty_replicas_raises():
with pytest.raises(ValueError, match="replicas must not be empty"):
RoundRobinBalancer().select({}, [])
def test_round_robin_balancer_after_large_index_and_shorter_list():
"""Large start_index % len(instances) then counter wraps with shorter list."""
"""Large start_index % len(replicas) then counter wraps with shorter list."""
now = time()
two = [
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10001",
output_addr="tcp://host:10001-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=0,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10002",
output_addr="tcp://host:10002-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=0,
last_heartbeat=now,
registered_at=now,
@@ -138,30 +138,30 @@ def test_round_robin_balancer_after_large_index_and_shorter_list():
def test_least_queue_length_balancer_picks_min_queue():
now = time()
instances = [
InstanceInfo(
replicas = [
ReplicaInfo(
input_addr="tcp://host:10001",
output_addr="tcp://host:10001-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=2,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10002",
output_addr="tcp://host:10002-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=0,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10003",
output_addr="tcp://host:10003-out",
stage_id=1,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=5,
last_heartbeat=now,
registered_at=now,
@@ -169,41 +169,41 @@ def test_least_queue_length_balancer_picks_min_queue():
]
balancer = LeastQueueLengthBalancer()
index = balancer.select({}, instances)
index = balancer.select({}, replicas)
assert index == 1
def test_least_queue_length_balancer_empty_instances_raises():
with pytest.raises(ValueError, match="instances must not be empty"):
def test_least_queue_length_balancer_empty_replicas_raises():
with pytest.raises(ValueError, match="replicas must not be empty"):
LeastQueueLengthBalancer().select({}, [])
def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
now = time()
instances = [
InstanceInfo(
replicas = [
ReplicaInfo(
input_addr="tcp://host:10001",
output_addr="tcp://host:10001-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=3,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10002",
output_addr="tcp://host:10002-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=3,
last_heartbeat=now,
registered_at=now,
),
InstanceInfo(
ReplicaInfo(
input_addr="tcp://host:10003",
output_addr="tcp://host:10003-out",
stage_id=1,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=3,
last_heartbeat=now,
registered_at=now,
@@ -214,21 +214,21 @@ def test_least_queue_length_balancer_equal_queues_uses_choice(mocker):
"vllm_omni.distributed.omni_coordinator.load_balancer.random.choice",
return_value=2,
)
assert balancer.select({}, instances) == 2
assert balancer.select({}, replicas) == 2
def test_least_queue_length_balancer_negative_queue_raises():
now = time()
instances = [
InstanceInfo(
replicas = [
ReplicaInfo(
input_addr="tcp://host:10001",
output_addr="tcp://host:10001-out",
stage_id=0,
status=StageStatus.UP,
status=ReplicaStatus.UP,
queue_length=-1,
last_heartbeat=now,
registered_at=now,
),
]
with pytest.raises(ValueError, match="queue_length must be non-negative"):
LeastQueueLengthBalancer().select({}, instances)
LeastQueueLengthBalancer().select({}, replicas)
tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py
+23 -23
View File
@@ -8,8 +8,8 @@ import pytest
import zmq
from vllm_omni.distributed.omni_coordinator import (
InstanceList,
OmniCoordClientForHub,
ReplicaList,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -32,8 +32,8 @@ def _wait_for_condition(cond, timeout: float = 2.0, interval: float = 0.01) -> b
return False
def test_hub_client_caches_instance_list_from_pub():
"""Verify OmniCoordClientForHub receives instance list updates from OmniCoordinator and caches for get_instance_list()."""
def test_hub_client_caches_replica_list_from_pub():
"""Verify OmniCoordClientForHub receives replica list updates from OmniCoordinator and caches for get_replica_list()."""
ctx, pub, endpoint = _bind_pub()
client = OmniCoordClientForHub(endpoint)
@@ -41,7 +41,7 @@ def test_hub_client_caches_instance_list_from_pub():
time.sleep(0.2)
now = time.time()
instances_payload = [
replicas_payload = [
{
"input_addr": "tcp://stage:10001",
"output_addr": "tcp://stage:10001-out",
@@ -71,37 +71,37 @@ def test_hub_client_caches_instance_list_from_pub():
},
]
payload = {"instances": instances_payload, "timestamp": now}
payload = {"replicas": replicas_payload, "timestamp": now}
pub.send(json.dumps(payload).encode("utf-8"))
assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 3)
assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 3)
inst_list = client.get_instance_list()
assert isinstance(inst_list, InstanceList)
assert len(inst_list.instances) == 3
rep_list = client.get_replica_list()
assert isinstance(rep_list, ReplicaList)
assert len(rep_list.replicas) == 3
for src, inst in zip(instances_payload, inst_list.instances, strict=True):
assert inst.input_addr == src["input_addr"]
assert inst.output_addr == src["output_addr"]
assert inst.stage_id == src["stage_id"]
assert inst.status.value == src["status"]
for src, rep in zip(replicas_payload, rep_list.replicas, strict=True):
assert rep.input_addr == src["input_addr"]
assert rep.output_addr == src["output_addr"]
assert rep.stage_id == src["stage_id"]
assert rep.status.value == src["status"]
stage0 = client.get_instances_for_stage(0)
stage1 = client.get_instances_for_stage(1)
stage0 = client.get_replicas_for_stage(0)
stage1 = client.get_replicas_for_stage(1)
assert all(inst.stage_id == 0 for inst in stage0.instances)
assert all(inst.stage_id == 1 for inst in stage1.instances)
assert all(rep.stage_id == 0 for rep in stage0.replicas)
assert all(rep.stage_id == 1 for rep in stage1.replicas)
# Send an updated list with fewer instances and verify cache refresh.
# Send an updated list with fewer replicas and verify cache refresh.
updated_payload = {
"instances": instances_payload[:2],
"replicas": replicas_payload[:2],
"timestamp": now + 1.0,
}
pub.send(json.dumps(updated_payload).encode("utf-8"))
assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 2)
updated_list = client.get_instance_list()
assert len(updated_list.instances) == 2
assert _wait_for_condition(lambda: len(client.get_replica_list().replicas) == 2)
updated_list = client.get_replica_list()
assert len(updated_list.replicas) == 2
client.close()
pub.close(0)
tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py
+6 -5
View File
@@ -9,7 +9,7 @@ import zmq
from vllm_omni.distributed.omni_coordinator import (
OmniCoordClientForStage,
StageStatus,
ReplicaStatus,
)
from vllm_omni.distributed.omni_coordinator import (
omni_coord_client_for_stage as stage_client_module,
@@ -47,7 +47,7 @@ def test_stage_client_auto_register_on_init():
event = _recv_event(router)
assert event["event_type"] == "update"
assert event["status"] == StageStatus.UP.value
assert event["status"] == ReplicaStatus.UP.value
assert event["stage_id"] == stage_id
assert event["input_addr"] == input_addr
assert event["output_addr"] == output_addr
@@ -70,13 +70,13 @@ def test_stage_client_update_info_sends_correct_event():
# Discard initial registration event.
_recv_event(router)
client.update_info(status=StageStatus.ERROR)
client.update_info(status=ReplicaStatus.ERROR)
client.update_info(queue_length=10)
first = _recv_event(router)
second = _recv_event(router)
assert first["status"] == StageStatus.ERROR.value
assert first["status"] == ReplicaStatus.ERROR.value
assert first["stage_id"] == stage_id
assert first["input_addr"] == input_addr
assert first["output_addr"] == output_addr
@@ -107,7 +107,7 @@ def test_stage_client_close_sends_down_status():
client.close()
event = _recv_event(router)
assert event["status"] == StageStatus.DOWN.value
assert event["status"] == ReplicaStatus.DOWN.value
assert event["stage_id"] == stage_id
assert event["input_addr"] == input_addr
assert event["output_addr"] == output_addr
@@ -274,6 +274,7 @@ def test_heartbeat_loop_retries_after_transient_send_failure():
client._closed = False
client._heartbeat_interval = 0.0
client._stop_event = _FakeStopEvent()
client._on_heartbeat = None
calls = {"count": 0}
tests/distributed/omni_coordinator/test_omni_coordinator.py
+34 -34
View File
@@ -11,14 +11,14 @@ from vllm.v1.utils import get_engine_client_zmq_addr
from vllm_omni.distributed.omni_coordinator import (
OmniCoordClientForStage,
OmniCoordinator,
StageStatus,
ReplicaStatus,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
"""Receive InstanceList JSON from SUB socket. Returns None on timeout."""
def _recv_replica_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
"""Receive ReplicaList JSON from SUB socket. Returns None on timeout."""
sub.setsockopt(zmq.RCVTIMEO, timeout_ms)
try:
data = sub.recv()
@@ -27,16 +27,16 @@ def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None:
return None
def _wait_for_instance_list(
def _wait_for_replica_list(
sub: zmq.Socket,
expected_count: int,
timeout: float = 3.0,
) -> dict | None:
"""Wait until received InstanceList with expected_count active instances."""
"""Wait until received ReplicaList with expected_count active replicas."""
start = time.time()
while time.time() - start < timeout:
msg = _recv_instance_list(sub, timeout_ms=500)
if msg is not None and len(msg.get("instances", [])) == expected_count:
msg = _recv_replica_list(sub, timeout_ms=500)
if msg is not None and len(msg.get("replicas", [])) == expected_count:
return msg
return None
@@ -45,7 +45,7 @@ def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None:
"""Drain queued SUB messages for a short window."""
deadline = time.time() + max_seconds
while time.time() < deadline:
_recv_instance_list(sub, timeout_ms=50)
_recv_replica_list(sub, timeout_ms=50)
def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
@@ -81,7 +81,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
)
# Wait for initial registration broadcast and clear any queued messages.
msg = _wait_for_instance_list(sub, expected_count=1)
msg = _wait_for_replica_list(sub, expected_count=1)
assert msg is not None
_drain_sub_messages(sub)
@@ -96,7 +96,7 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
deadline = time.time() + window_s
recv_count = 0
while time.time() < deadline:
if _recv_instance_list(sub, timeout_ms=100) is not None:
if _recv_replica_list(sub, timeout_ms=100) is not None:
recv_count += 1
assert recv_count < update_count // 2, (
@@ -110,8 +110,8 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
def test_omni_coordinator_registration_broadcast():
"""Verify that after multiple OmniCoordClientForStage instances register,
OmniCoordinator publishes an InstanceList containing all registered instances.
"""Verify that after multiple OmniCoordClientForStage replicas register,
OmniCoordinator publishes a ReplicaList containing all registered replicas.
"""
router_addr = get_engine_client_zmq_addr(
local_only=False,
@@ -144,12 +144,12 @@ def test_omni_coordinator_registration_broadcast():
OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
]
msg = _wait_for_instance_list(sub, expected_count=3)
assert msg is not None, "Expected InstanceList with 3 instances"
assert len(msg["instances"]) == 3
msg = _wait_for_replica_list(sub, expected_count=3)
assert msg is not None, "Expected ReplicaList with 3 replicas"
assert len(msg["replicas"]) == 3
assert isinstance(msg["timestamp"], (int, float))
input_addrs = {inst["input_addr"] for inst in msg["instances"]}
input_addrs = {rep["input_addr"] for rep in msg["replicas"]}
assert "tcp://stage:10001" in input_addrs
assert "tcp://stage:10002" in input_addrs
assert "tcp://stage:10003" in input_addrs
@@ -162,7 +162,7 @@ def test_omni_coordinator_registration_broadcast():
def test_omni_coordinator_heartbeat_timeout_handling():
"""Verify that when a stage instance stops sending heartbeats,
"""Verify that when a stage replica stops sending heartbeats,
OmniCoordinator marks it as unhealthy and excludes it from the active list.
"""
router_addr = get_engine_client_zmq_addr(
@@ -201,23 +201,23 @@ def test_omni_coordinator_heartbeat_timeout_handling():
"output_addr": "tcp://stage:c-out",
"stage_id": 0,
"event_type": "update",
"status": StageStatus.UP.value,
"status": ReplicaStatus.UP.value,
"queue_length": 0,
}
dealer_c.send(json.dumps(reg_event).encode("utf-8"))
msg = _wait_for_instance_list(sub, expected_count=3)
assert msg is not None, "Expected initial 3 instances"
assert len(msg["instances"]) == 3
msg = _wait_for_replica_list(sub, expected_count=3)
assert msg is not None, "Expected initial 3 replicas"
assert len(msg["replicas"]) == 3
# Wait for heartbeat timeout (timeout=5s, check interval ~2.5s).
time.sleep(8.0)
# Receive the update (C should be ERROR and excluded from active list).
msg_after_timeout = _wait_for_instance_list(sub, expected_count=2, timeout=5.0)
assert msg_after_timeout is not None, "Expected InstanceList with 2 instances after timeout"
instances = msg_after_timeout.get("instances", [])
input_addrs = {inst["input_addr"] for inst in instances}
msg_after_timeout = _wait_for_replica_list(sub, expected_count=2, timeout=5.0)
assert msg_after_timeout is not None, "Expected ReplicaList with 2 replicas after timeout"
replicas = msg_after_timeout.get("replicas", [])
input_addrs = {rep["input_addr"] for rep in replicas}
assert "tcp://stage:a" in input_addrs
assert "tcp://stage:b" in input_addrs
@@ -232,8 +232,8 @@ def test_omni_coordinator_heartbeat_timeout_handling():
sub_ctx.term()
def test_omni_coordinator_instance_shutdown_handling():
"""Verify that when a stage instance sends status='down',
def test_omni_coordinator_replica_shutdown_handling():
"""Verify that when a stage replica sends status='down',
OmniCoordinator removes it from the active list and broadcasts an updated list.
"""
router_addr = get_engine_client_zmq_addr(
@@ -261,18 +261,18 @@ def test_omni_coordinator_instance_shutdown_handling():
client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
msg = _wait_for_instance_list(sub, expected_count=1)
msg = _wait_for_replica_list(sub, expected_count=1)
assert msg is not None
assert len(msg["instances"]) == 1
assert msg["instances"][0]["input_addr"] == "tcp://stage:shutdown"
assert len(msg["replicas"]) == 1
assert msg["replicas"][0]["input_addr"] == "tcp://stage:shutdown"
# Send down status (simulating graceful shutdown).
client.update_info(status=StageStatus.DOWN)
client.update_info(status=ReplicaStatus.DOWN)
# Receive updated list (should have 0 active instances).
msg = _wait_for_instance_list(sub, expected_count=0)
# Receive updated list (should have 0 active replicas).
msg = _wait_for_replica_list(sub, expected_count=0)
assert msg is not None
assert len(msg["instances"]) == 0
assert len(msg["replicas"]) == 0
client.close()
coordinator.close()
tests/e2e/accuracy/test_hunyuan_image3.py
+26 -9
View File
@@ -23,18 +23,31 @@ from tests.e2e.accuracy.helpers import (
download_images,
model_output_dir,
)
from tests.helpers.runtime import OmniRunner
from tests.helpers.mark import hardware_marks
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, *hardware_marks(res={"cuda": "H100"}, num_cards=8)]
# ============================================================================
# Configurable Parameters
# ============================================================================
AR_DEVICES = "0,1"
DIT_DEVICES = "2,3"
# Comma-separated logical CUDA device ids per stage: split visible GPUs (0..n-1), first half -> AR, second -> DiT.
def _default_ar_dit_devices() -> tuple[str, str]:
"""First floor(n/2) logical devices -> AR, rest -> DiT. ``device_count`` respects ``CUDA_VISIBLE_DEVICES``."""
n = torch.accelerator.device_count()
if n < 2:
return "0,1", "2,3"
split = n // 2
ar = ",".join(str(i) for i in range(split))
dit = ",".join(str(i) for i in range(split, n))
return ar, dit
AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices()
MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 2.5
@@ -75,7 +88,7 @@ _BASE_CONFIG = {
"scheduler_cls": "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler",
"gpu_memory_utilization": 0.95, "enforce_eager": True, "trust_remote_code": True,
"engine_output_type": "latent", "enable_prefix_caching": False,
"max_num_batched_tokens": 32768, "tensor_parallel_size": 2, "pipeline_parallel_size": 1,
"max_num_batched_tokens": 32768, "tensor_parallel_size": AR_TP_SIZE, "pipeline_parallel_size": 1,
"hf_overrides": {"rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"}},
},
"is_comprehension": False, "final_output": True, "final_output_type": "text",
@@ -91,7 +104,7 @@ _BASE_CONFIG = {
"engine_args": {
"model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM",
"enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp",
"parallel_config": {"tensor_parallel_size": 2, "enable_expert_parallel": True},
"parallel_config": {"tensor_parallel_size": DIT_TP_SIZE, "enable_expert_parallel": True},
},
"engine_input_source": [0],
"custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion",
@@ -134,12 +147,13 @@ def _make_config(enable_kv_reuse: bool, path: Path) -> None:
def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
from transformers import AutoTokenizer
from tests.helpers.runtime import OmniRunner
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
from vllm_omni.platforms import current_omni_platform
build_kwargs: dict = {"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified", "num_images": 2}
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)
result = build_prompt_tokens(
PROMPT,
tokenizer,
@@ -149,7 +163,7 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
system_prompt_type = result.system_prompt_type
ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer)
with OmniRunner(MODEL_NAME, stage_configs_path=stage_configs_path) as runner:
with OmniRunner(MODEL_PATH, stage_configs_path=stage_configs_path) as runner:
params_list = list(runner.omni.default_sampling_params_list)
for sp in params_list:
if isinstance(sp, OmniDiffusionSamplingParams):
@@ -206,7 +220,10 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
return image, cot_text, elapsed
@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Needs 4+ GPUs (2 AR + 2 DiT)")
@pytest.mark.skipif(
torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
)
def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
if importlib.util.find_spec("FlagEmbedding") is None:
raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
Loaded 100 of 224 files, more files were not shown because too many files have changed in this diff. Show more