Compare commits
33
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ac67e8874d | ||
|
|
dfe6e13cfc | ||
|
|
9d80efcbe5 | ||
|
|
ddbbfd29e5 | ||
|
|
9235b0ae65 | ||
|
|
1c2a81f6d8 | ||
|
|
b4581b293b | ||
|
|
9f978923f2 | ||
|
|
589afcef6c | ||
|
|
556c5f88a0 | ||
|
|
2eada79b99 | ||
|
|
7fb62db5c2 | ||
|
|
6510de344e | ||
|
|
740cb35a59 | ||
|
|
ae7bee0c43 | ||
|
|
6a604be610 | ||
|
|
f201b7171f | ||
|
|
0a7ac4c67e | ||
|
|
e788ef6e76 | ||
|
|
f24a6165ff | ||
|
|
bfaaa3c2c3 | ||
|
|
c9f2e5ad94 | ||
|
|
fbea788668 | ||
|
|
78c144f3a8 | ||
|
|
1c4aca7f6e | ||
|
|
6fdc4ca62d | ||
|
|
b6728db486 | ||
|
|
b9a51592f0 | ||
|
|
5215e03a91 | ||
|
|
d5d47e9be1 | ||
|
|
30cba84222 | ||
|
|
9f9f9f274a | ||
|
|
900a7f0813 |
No files matched your search
@@ -427,6 +427,12 @@ steps:
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · MiniMax H3 T2VA Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_t2va_similarity.py -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Other Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
@@ -477,12 +483,18 @@ steps:
|
||||
- pytest -sv tests/diffusion/offloader -m "full_model and cuda"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: Diffusion · Quantization Test with H100"
|
||||
- label: ":full_moon: Diffusion · Quantization Test with H100 · Single-GPU"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion · MiniMax H3 FP8 Quality Test with H100 · 2-GPU"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -sv tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py -m 'full_model and cuda and H100' --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Diffusion · Quantization Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
|
||||
+4
-4
@@ -16,15 +16,15 @@
|
||||
/.claude/ @hsliuustc0106
|
||||
|
||||
# Tests
|
||||
/tests/ @yenuo26
|
||||
/tests/ @yenuo26 @NickCao
|
||||
|
||||
# CI
|
||||
/.buildkite/ @yenuo26 @congw729
|
||||
/.buildkite/ @yenuo26 @congw729 @NickCao
|
||||
|
||||
# Engine and model_executor
|
||||
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
||||
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
|
||||
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks
|
||||
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
|
||||
/vllm_omni/model_executor/ @tzhouam @gcanlin
|
||||
/vllm_omni/worker/ @tzhouam @gcanlin
|
||||
|
||||
@@ -37,7 +37,7 @@
|
||||
|
||||
# Omni and TTS support
|
||||
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
|
||||
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin
|
||||
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
|
||||
|
||||
# Benchmarks
|
||||
/benchmarks/accuracy/ @david6666666
|
||||
|
||||
@@ -47,7 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes
|
||||
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
|
||||
|
||||
- **Generate Image** for text-to-image and image-to-image tasks
|
||||
- **Generate Video** for text-to-video and image-to-video tasks
|
||||
- **Generate Video** for text-to-video, first-frame/image-to-video, and reference-conditioned video
|
||||
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
|
||||
- **TTS** and **TTS Voice Clone** for TTS tasks
|
||||
|
||||
@@ -105,7 +105,7 @@ You can configure per-stage sampling parameters for multi-stage models.
|
||||
> [!TIP]
|
||||
> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
|
||||
|
||||
### Text-to-video and image-to-video generation (e.g., Wan)
|
||||
### Text-to-video and image-to-video generation (e.g., Wan, MiniMax-H3)
|
||||
|
||||
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
|
||||
|
||||
@@ -117,7 +117,11 @@ You can configure per-stage sampling parameters for multi-stage models.
|
||||
</p>
|
||||
|
||||
> [!TIP]
|
||||
> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
|
||||
> Connect a **frame** image for first-frame / image-to-video (e.g., Wan I2V, MiniMax-H3 FL2VA).
|
||||
>
|
||||
> For reference-conditioned generation (MiniMax-H3 Ref2VA), connect a **Video References** node instead.
|
||||
>
|
||||
> Do not use `frame` and `references` together. Task routing is automatic from which inputs you connect.
|
||||
|
||||
### TTS (e.g., Qwen TTS series)
|
||||
|
||||
|
||||
@@ -15,11 +15,13 @@ from .comfyui_vllm_omni.nodes import (
|
||||
VLLMOmniDiffusionSampling,
|
||||
VLLMOmniGenerateImage,
|
||||
VLLMOmniGenerateVideo,
|
||||
VLLMOmniMiniMaxH3Params,
|
||||
VLLMOmniQwenTTSParams,
|
||||
VLLMOmniRemoteLoRA,
|
||||
VLLMOmniSamplingParamsList,
|
||||
VLLMOmniTTS,
|
||||
VLLMOmniUnderstanding,
|
||||
VLLMOmniVideoReferences,
|
||||
VLLMOmniVoiceClone,
|
||||
VLLMOmniWanParams,
|
||||
)
|
||||
@@ -32,6 +34,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
|
||||
"VLLMOmniTTS": VLLMOmniTTS,
|
||||
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
|
||||
"VLLMOmniVideoReferences": VLLMOmniVideoReferences,
|
||||
# === Params ===
|
||||
"VLLMOmniARSampling": VLLMOmniARSampling,
|
||||
"VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
|
||||
@@ -39,6 +42,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
|
||||
"VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
|
||||
"VLLMOmniWanParams": VLLMOmniWanParams,
|
||||
"VLLMOmniMiniMaxH3Params": VLLMOmniMiniMaxH3Params,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
@@ -49,6 +53,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VLLMOmniUnderstanding": "Multimodality Understanding",
|
||||
"VLLMOmniTTS": "TTS (Text to Speech)",
|
||||
"VLLMOmniVoiceClone": "TTS Voice Cloning",
|
||||
"VLLMOmniVideoReferences": "Video References",
|
||||
# === Params ===
|
||||
"VLLMOmniARSampling": "AR Sampling Params",
|
||||
"VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
|
||||
@@ -56,6 +61,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VLLMOmniRemoteLoRA": "LoRA",
|
||||
"VLLMOmniQwenTTSParams": "Qwen TTS Params",
|
||||
"VLLMOmniWanParams": "Wan Video Params",
|
||||
"VLLMOmniMiniMaxH3Params": "MiniMax-H3 Video Params",
|
||||
}
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
@@ -10,7 +10,9 @@ from .utils.types import (
|
||||
AudioFormat,
|
||||
AutoregressionSamplingParams,
|
||||
DiffusionSamplingParams,
|
||||
MiniMaxH3ModelSpecificParams,
|
||||
QwenTTSModelSpecificParams,
|
||||
VideoReferences,
|
||||
WanModelSpecificParams,
|
||||
)
|
||||
from .utils.validators import (
|
||||
@@ -163,7 +165,8 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
|
||||
"num_frames": ("INT", {"default": 41, "min": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
"frame": ("IMAGE",),
|
||||
"references": ("VIDEO_REFERENCES",),
|
||||
"sampling_params": ("SAMPLING_PARAMS",),
|
||||
"lora": ("REMOTE_LORA",),
|
||||
"model_params": ("VIDEO_PARAMS",),
|
||||
@@ -174,6 +177,15 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
|
||||
RETURN_NAMES = ("video",)
|
||||
FUNCTION = "generate"
|
||||
|
||||
@classmethod
|
||||
def VALIDATE_INPUTS(cls, url, model, frame=None, references=None, **_kwargs) -> str | Literal[True]:
|
||||
base = super().VALIDATE_INPUTS(url, model)
|
||||
if base is not True:
|
||||
return base
|
||||
if frame is not None and references is not None:
|
||||
return "Provide only one of frame or references, not both."
|
||||
return True
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
url: str,
|
||||
@@ -184,7 +196,8 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
|
||||
fps: int,
|
||||
num_frames: int,
|
||||
negative_prompt: str | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
frame: torch.Tensor | None = None,
|
||||
references: dict | None = None,
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
model_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
@@ -207,14 +220,16 @@ class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
|
||||
|
||||
if sampling_params is not None:
|
||||
sampling_params.pop("type", None) # internal fields
|
||||
if model_params is not None:
|
||||
model_params.pop("type", None) # internal fields
|
||||
# model_params["type"] is also an internal field, but is used in api_client
|
||||
# if model_params is not None:
|
||||
# model_params.pop("type", None) # internal fields
|
||||
|
||||
client = VLLMOmniClient(url)
|
||||
output = await client.generate_video(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
image=image, # image present => i2v, absent => t2v
|
||||
frame=frame, # frame present => fl2va / Wan I2V
|
||||
references=references,
|
||||
width=width,
|
||||
height=height,
|
||||
num_frames=num_frames,
|
||||
@@ -734,3 +749,78 @@ class VLLMOmniWanParams:
|
||||
|
||||
def get_params(self, **kwargs):
|
||||
return (WanModelSpecificParams(kwargs),)
|
||||
|
||||
|
||||
class VLLMOmniMiniMaxH3Params:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"audio_flow_shift": (
|
||||
"FLOAT",
|
||||
{"default": 3.0, "min": 0.0, "max": 100.0, "step": 0.1},
|
||||
),
|
||||
"flow_shift": (
|
||||
"FLOAT",
|
||||
{"default": 12.0, "min": 0.0, "max": 100.0, "step": 0.1},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO_PARAMS",)
|
||||
RETURN_NAMES = ("MiniMax-H3 video params",)
|
||||
FUNCTION = "get_params"
|
||||
CATEGORY = "vLLM-Omni/Video Params"
|
||||
|
||||
def get_params(self, **kwargs):
|
||||
params = MiniMaxH3ModelSpecificParams(kwargs)
|
||||
params["type"] = "minimax_h3"
|
||||
return (params,)
|
||||
|
||||
|
||||
class VLLMOmniVideoReferences:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {},
|
||||
"optional": {
|
||||
"image_1": ("IMAGE",),
|
||||
"image_2": ("IMAGE",),
|
||||
"audio_1": ("AUDIO",),
|
||||
"audio_2": ("AUDIO",),
|
||||
"video_1": ("VIDEO",),
|
||||
"video_2": ("VIDEO",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO_REFERENCES",)
|
||||
RETURN_NAMES = ("references",)
|
||||
FUNCTION = "get_references"
|
||||
CATEGORY = "vLLM-Omni"
|
||||
|
||||
def get_references(
|
||||
self,
|
||||
image_1: torch.Tensor | None = None,
|
||||
image_2: torch.Tensor | None = None,
|
||||
audio_1: AudioInput | None = None,
|
||||
audio_2: AudioInput | None = None,
|
||||
video_1: VideoInput | None = None,
|
||||
video_2: VideoInput | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
refs = VideoReferences()
|
||||
if image_1 is not None:
|
||||
refs["image_1"] = image_1
|
||||
if image_2 is not None:
|
||||
refs["image_2"] = image_2
|
||||
if audio_1 is not None:
|
||||
refs["audio_1"] = audio_1
|
||||
if audio_2 is not None:
|
||||
refs["audio_2"] = audio_2
|
||||
if video_1 is not None:
|
||||
refs["video_1"] = video_1
|
||||
if video_2 is not None:
|
||||
refs["video_2"] = video_2
|
||||
return (refs,)
|
||||
@@ -24,6 +24,7 @@ from .format import (
|
||||
image_tensor_to_base64,
|
||||
image_tensor_to_png_bytes,
|
||||
video_to_base64,
|
||||
video_to_bytes,
|
||||
)
|
||||
from .logger import get_logger, pretty_printer
|
||||
from .models import lookup_model_spec
|
||||
@@ -63,7 +64,11 @@ async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get",
|
||||
|
||||
class VLLMOmniClient:
|
||||
def __init__(
|
||||
self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
|
||||
self,
|
||||
base_url: str,
|
||||
timeout: float | None = None,
|
||||
poll_interval: float = 5.0,
|
||||
max_poll_duration: float = 60 * 30,
|
||||
):
|
||||
self.base_url = base_url
|
||||
self.timeout = aiohttp.ClientTimeout(total=timeout)
|
||||
@@ -261,12 +266,17 @@ class VLLMOmniClient:
|
||||
num_frames: int,
|
||||
fps: int,
|
||||
negative_prompt: str | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
frame: torch.Tensor | None = None,
|
||||
references: dict | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
model_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
**extra_body,
|
||||
**extra_params,
|
||||
) -> VideoInput:
|
||||
if frame is not None and references is not None:
|
||||
raise ValueError("Provide only one of frame or references, not both.")
|
||||
|
||||
# === regular payload fields ===
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model", model)
|
||||
form.add_field("prompt", prompt)
|
||||
@@ -279,23 +289,91 @@ class VLLMOmniClient:
|
||||
if sampling_params is not None:
|
||||
for k, v in sampling_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if model_params is not None:
|
||||
for k, v in model_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if lora is not None:
|
||||
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
|
||||
if extra_body:
|
||||
form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
|
||||
|
||||
if image is not None:
|
||||
# === multimodal inputs (first-last-frames, references, etc.) ===
|
||||
input_reference_image: torch.Tensor | None = None
|
||||
audio_reference: AudioInput | None = None
|
||||
reference_videos: list[VideoInput] = []
|
||||
video_task: str | None = None
|
||||
|
||||
if frame is not None:
|
||||
input_reference_image = frame
|
||||
video_task = "fl2va"
|
||||
elif references is not None:
|
||||
images = [references[k] for k in ("image_1", "image_2") if k in references and references[k] is not None]
|
||||
audios = [references[k] for k in ("audio_1", "audio_2") if k in references and references[k] is not None]
|
||||
videos = [references[k] for k in ("video_1", "video_2") if k in references and references[k] is not None]
|
||||
|
||||
if not images and not audios and not videos:
|
||||
raise ValueError("references is empty; connect at least one image/audio/video.")
|
||||
|
||||
if videos:
|
||||
if images or audios:
|
||||
raise ValueError(
|
||||
"Video references cannot be combined with image or audio references. "
|
||||
"Connect only video_1/video_2 for multi-video Ref2VA."
|
||||
)
|
||||
reference_videos = videos
|
||||
video_task = "ref2va"
|
||||
elif len(images) == 1 and len(audios) == 1:
|
||||
input_reference_image = images[0]
|
||||
audio_reference = audios[0]
|
||||
video_task = "ref2va"
|
||||
else:
|
||||
raise ValueError(
|
||||
"Invalid references combination. Supported modes: "
|
||||
"(1) one or more videos only, or (2) exactly one image and one audio."
|
||||
)
|
||||
else:
|
||||
video_task = "t2va"
|
||||
|
||||
if input_reference_image is not None:
|
||||
image_filename = "image.png" # Required for multipart form
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_tensor_to_png_bytes(image, image_filename),
|
||||
image_tensor_to_png_bytes(input_reference_image, image_filename),
|
||||
filename=image_filename,
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
if audio_reference is not None:
|
||||
form.add_field(
|
||||
"audio_reference",
|
||||
json.dumps({"audio_url": audio_to_base64(audio_reference)}, ensure_ascii=False),
|
||||
)
|
||||
|
||||
for idx, video in enumerate(reference_videos, start=1):
|
||||
video_filename = f"reference_{idx}.mp4"
|
||||
form.add_field(
|
||||
"input_references",
|
||||
video_to_bytes(video, video_filename),
|
||||
filename=video_filename,
|
||||
content_type="video/mp4",
|
||||
)
|
||||
|
||||
# === model specific params. Either use a specialized builder, or add flattened fields as-is ===
|
||||
if model_params is not None:
|
||||
model_params = dict(model_params)
|
||||
model_params.pop("type", None)
|
||||
|
||||
spec, _ = lookup_model_spec(model)
|
||||
params_builder = spec.get("params_builder") if spec else None
|
||||
if params_builder is not None:
|
||||
form_fields = params_builder(
|
||||
model_params or {},
|
||||
extra_params={**extra_params, "task": video_task},
|
||||
)
|
||||
for k, v in form_fields.items():
|
||||
form.add_field(k, v if isinstance(v, str) else str(v))
|
||||
else:
|
||||
if model_params is not None:
|
||||
for k, v in model_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if extra_params:
|
||||
form.add_field("extra_params", json.dumps(extra_params, ensure_ascii=False))
|
||||
|
||||
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
||||
# Start the video generation job
|
||||
url = f"{self.base_url}/videos"
|
||||
|
||||
@@ -1,7 +1,12 @@
|
||||
import json
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from .logger import get_logger
|
||||
from .types import Modality, ModelMode, Spec
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _bagel_payload_preprocessor(payload: dict) -> dict:
|
||||
try:
|
||||
@@ -29,6 +34,30 @@ def _qwen25_payload_preprocessor(payload: dict) -> dict:
|
||||
return payload
|
||||
|
||||
|
||||
def _minimaxh3_params_builder(
|
||||
model_params: dict[str, Any],
|
||||
*,
|
||||
extra_params: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Build multipart form fields for MiniMax-H3 from model_params + routed task."""
|
||||
params = dict(model_params)
|
||||
params.pop("type", None)
|
||||
form_fields: dict[str, Any] = {}
|
||||
merged_extra_params: dict[str, Any] = dict(extra_params or {})
|
||||
|
||||
if "flow_shift" in params:
|
||||
form_fields["flow_shift"] = params.pop("flow_shift")
|
||||
for key in ("audio_flow_shift",):
|
||||
if key in params:
|
||||
merged_extra_params[key] = params.pop(key)
|
||||
if params:
|
||||
logger.warning("Unused MiniMax-H3 model params ignored: %s", sorted(params))
|
||||
|
||||
if merged_extra_params:
|
||||
form_fields["extra_params"] = json.dumps(merged_extra_params, ensure_ascii=False)
|
||||
return form_fields
|
||||
|
||||
|
||||
_MODEL_PIPELINE_SPECS: dict[str, Spec] = {
|
||||
r"BAGEL-7B-MoT": {
|
||||
"stages": [
|
||||
@@ -71,6 +100,21 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
|
||||
}
|
||||
],
|
||||
},
|
||||
r"MiniMax-H3(/FL2VA|/Ref2VA)?": {
|
||||
"stages": ["diffusion"],
|
||||
"modes": [
|
||||
{
|
||||
"mode": ModelMode.VIDEO_GENERATION,
|
||||
"input_modalities": [
|
||||
Modality.TEXT,
|
||||
Modality.IMAGE,
|
||||
Modality.VIDEO,
|
||||
Modality.AUDIO,
|
||||
],
|
||||
}
|
||||
],
|
||||
"params_builder": _minimaxh3_params_builder,
|
||||
},
|
||||
}
|
||||
# Convert dict keys to regex patterns
|
||||
MODEL_PIPELINE_SPECS: dict[re.Pattern, Spec] = {}
|
||||
@@ -80,12 +124,10 @@ del _MODEL_PIPELINE_SPECS
|
||||
|
||||
|
||||
def lookup_model_spec(model: str) -> tuple[Spec | None, str | None]:
|
||||
try:
|
||||
last_component = model.rstrip("/").rsplit("/", 1)[-1]
|
||||
except IndexError:
|
||||
last_component = model
|
||||
normalized = model.rstrip("/").lstrip("/")
|
||||
lookup_key = normalized.split("/", 1)[-1]
|
||||
for pattern, spec in MODEL_PIPELINE_SPECS.items():
|
||||
if pattern.search(last_component):
|
||||
if pattern.search(lookup_key):
|
||||
return spec, pattern.pattern
|
||||
return None, None
|
||||
|
||||
@@ -101,6 +143,9 @@ if __name__ == "__main__":
|
||||
"Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"Custom/Path/UnknownModel-Instruct",
|
||||
"Not/Matching/Anything",
|
||||
"/path/to/MiniMax-H3/FL2VA",
|
||||
"/path/to/MiniMax-H3/Ref2VA",
|
||||
"MiniMaxAI/MiniMax-H3",
|
||||
]
|
||||
|
||||
test_payload = {"messages": [{"role": "user", "content": "prompt"}]}
|
||||
@@ -112,7 +157,13 @@ if __name__ == "__main__":
|
||||
if preprocessor := spec.get("payload_preprocessor"):
|
||||
result = preprocessor(test_payload)
|
||||
print(f"✓ {path:<40} → {result}")
|
||||
elif params_builder := spec.get("params_builder"):
|
||||
result = params_builder(
|
||||
{"audio_flow_shift": 3.0, "flow_shift": 12.0, "type": "minimax_h3"},
|
||||
extra_params={"task": "t2va"},
|
||||
)
|
||||
print(f"✓ {path:<40} → {result}")
|
||||
else:
|
||||
print(f"✓ {path:<40} → No preprocessor")
|
||||
print(f"✓ {path:<40} → No preprocessor/params_builder")
|
||||
else:
|
||||
print(f"✗ {path:<40} → No match")
|
||||
@@ -23,6 +23,14 @@ class WanModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
class MiniMaxH3ModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
class VideoReferences(dict):
|
||||
pass
|
||||
|
||||
|
||||
class ModelMode(Enum):
|
||||
IMAGE_GENERATION = auto()
|
||||
VIDEO_GENERATION = auto()
|
||||
@@ -43,9 +51,11 @@ class ModelModeSpec(TypedDict):
|
||||
|
||||
|
||||
PayloadPreprocessor: TypeAlias = Callable[[dict[str, Any]], dict[str, Any]]
|
||||
ParamsBuilder: TypeAlias = Callable[..., dict[str, Any]]
|
||||
|
||||
|
||||
class Spec(TypedDict):
|
||||
stages: list[Literal["diffusion", "autoregression"]]
|
||||
modes: list[ModelModeSpec]
|
||||
payload_preprocessor: NotRequired[PayloadPreprocessor]
|
||||
params_builder: NotRequired[ParamsBuilder]
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 129 KiB After Width: | Height: | Size: 161 KiB |
@@ -1,513 +1 @@
|
||||
{
|
||||
"id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
|
||||
"revision": 0,
|
||||
"last_node_id": 14,
|
||||
"last_link_id": 28,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 10,
|
||||
"type": "MarkdownNote",
|
||||
"pos": [
|
||||
227.99306819575278,
|
||||
-231.24143306069843
|
||||
],
|
||||
"size": [
|
||||
240.3326478922474,
|
||||
136.3505820791881
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"title": "Note: Task and Input",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#000"
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
496.31859627609606,
|
||||
-229.71277089860084
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "image"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "choose file to upload",
|
||||
"name": "upload",
|
||||
"type": "IMAGEUPLOAD",
|
||||
"widget": {
|
||||
"name": "upload"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "IMAGE",
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
25
|
||||
]
|
||||
},
|
||||
{
|
||||
"localized_name": "MASK",
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"cute-cat.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "VLLMOmniDiffusionSampling",
|
||||
"pos": [
|
||||
465.5140218220927,
|
||||
148.80072646828987
|
||||
],
|
||||
"size": [
|
||||
284.205078125,
|
||||
226
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "n",
|
||||
"name": "n",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "n"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_inference_steps",
|
||||
"name": "num_inference_steps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_inference_steps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "guidance_scale",
|
||||
"name": "guidance_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "true_cfg_scale",
|
||||
"name": "true_cfg_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "true_cfg_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_slicing",
|
||||
"name": "vae_use_slicing",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_slicing"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_tiling",
|
||||
"name": "vae_use_tiling",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_tiling"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "seed",
|
||||
"name": "seed",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "seed"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "diffusion sampling params",
|
||||
"name": "diffusion sampling params",
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"links": [
|
||||
27
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniDiffusionSampling"
|
||||
},
|
||||
"widgets_values": [
|
||||
1,
|
||||
50,
|
||||
1,
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
1992,
|
||||
"randomize"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 14,
|
||||
"type": "VLLMOmniWanParams",
|
||||
"pos": [
|
||||
518.6005521845747,
|
||||
428.4455663719664
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
106
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "guidance_scale_2",
|
||||
"name": "guidance_scale_2",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale_2"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "boundary_ratio",
|
||||
"name": "boundary_ratio",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "boundary_ratio"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "flow_shift",
|
||||
"name": "flow_shift",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "flow_shift"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "Wan video params",
|
||||
"name": "Wan video params",
|
||||
"type": "VIDEO_PARAMS",
|
||||
"links": [
|
||||
28
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniWanParams"
|
||||
},
|
||||
"widgets_values": [
|
||||
4,
|
||||
0.875,
|
||||
5
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 12,
|
||||
"type": "VLLMOmniGenerateVideo",
|
||||
"pos": [
|
||||
827.2566336087867,
|
||||
-73.77449831827545
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
346
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"shape": 7,
|
||||
"type": "IMAGE",
|
||||
"link": 25
|
||||
},
|
||||
{
|
||||
"localized_name": "sampling_params",
|
||||
"name": "sampling_params",
|
||||
"shape": 7,
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"link": 27
|
||||
},
|
||||
{
|
||||
"localized_name": "lora",
|
||||
"name": "lora",
|
||||
"shape": 7,
|
||||
"type": "REMOTE_LORA",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model_params",
|
||||
"name": "model_params",
|
||||
"shape": 7,
|
||||
"type": "VIDEO_PARAMS",
|
||||
"link": 28
|
||||
},
|
||||
{
|
||||
"localized_name": "url",
|
||||
"name": "url",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "url"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model",
|
||||
"name": "model",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "model"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "prompt",
|
||||
"name": "prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "negative_prompt",
|
||||
"name": "negative_prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "negative_prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "width",
|
||||
"name": "width",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "width"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "height",
|
||||
"name": "height",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "height"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "fps",
|
||||
"name": "fps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "fps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_frames",
|
||||
"name": "num_frames",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_frames"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": [
|
||||
26
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniGenerateVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"http://localhost:8000/v1",
|
||||
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
|
||||
"Cartoonish",
|
||||
900,
|
||||
1200,
|
||||
12,
|
||||
41
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 13,
|
||||
"type": "SaveVideo",
|
||||
"pos": [
|
||||
1257.5441369501657,
|
||||
-44.12956394419863
|
||||
],
|
||||
"size": [
|
||||
374,
|
||||
372.34896826553796
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"link": 26
|
||||
},
|
||||
{
|
||||
"localized_name": "filename_prefix",
|
||||
"name": "filename_prefix",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "filename_prefix"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "format",
|
||||
"name": "format",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "format"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "codec",
|
||||
"name": "codec",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "codec"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"video/ComfyUI",
|
||||
"auto",
|
||||
"auto"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
25,
|
||||
4,
|
||||
0,
|
||||
12,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
26,
|
||||
12,
|
||||
0,
|
||||
13,
|
||||
0,
|
||||
"VIDEO"
|
||||
],
|
||||
[
|
||||
27,
|
||||
8,
|
||||
0,
|
||||
12,
|
||||
1,
|
||||
"SAMPLING_PARAMS"
|
||||
],
|
||||
[
|
||||
28,
|
||||
14,
|
||||
0,
|
||||
12,
|
||||
3,
|
||||
"VIDEO_PARAMS"
|
||||
]
|
||||
],
|
||||
"groups": [
|
||||
{
|
||||
"id": 1,
|
||||
"title": "Input",
|
||||
"bounding": [
|
||||
213.1706010087147,
|
||||
-313.20095750667554,
|
||||
560.8246471437027,
|
||||
407.4246532472182
|
||||
],
|
||||
"color": "#3f789e",
|
||||
"font_size": 24,
|
||||
"flags": {}
|
||||
}
|
||||
],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"workflowRendererVersion": "LG",
|
||||
"ds": {
|
||||
"scale": 1.1469075819486894,
|
||||
"offset": [
|
||||
101.7632950847089,
|
||||
492.29121889347806
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
{"id":"3aea1d08-7de5-4997-ba71-b2f66ad72a21","revision":0,"last_node_id":24,"last_link_id":35,"nodes":[{"id":13,"type":"SaveVideo","pos":[1315.729153459605,-137.92038948394355],"size":[374,313.7692307692308],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":32},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMFY_DYNAMICCOMBO_V3","widget":{"name":"codec"},"link":null}],"outputs":[{"localized_name":"video","name":"video","type":"VIDEO","links":null}],"properties":{},"widgets_values":["video/ComfyUI","auto","auto"]},{"id":10,"type":"MarkdownNote","pos":[183.74636797766794,-571.2529158046766],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":0,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both *text-to-video*, *frame-to-video*, and *reference-to-video* generation. The node will create corresponding payload depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[452.07189605801125,-569.7242536425791],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[29]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["rap.png","image"]},{"id":23,"type":"LoadImage","pos":[9.206110246057657,-7.6955291509170785],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[34]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["white-cat.jpg","image"]},{"id":22,"type":"MarkdownNote","pos":[-9.783113411416666,-148.8790300038605],"size":[297.3012007349086,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Reference to video (r2v)","properties":{},"widgets_values":["For r2v tasks, connect this Video Reference block to the main Generate Video block. Then disconnect the \"frame\" input above."],"color":"#432","bgcolor":"#000"},{"id":24,"type":"LoadAudio","pos":[317.9929742385805,-141.44821500065095],"size":[282.798828125,136],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[35]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["nothing-beats-a-jet2-holiday.mp3",null,""]},{"id":21,"type":"VLLMOmniVideoReferences","pos":[376.6129172897595,105.41557277556292],"size":[156.9380859375,126],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"image_1","name":"image_1","shape":7,"type":"IMAGE","link":34},{"localized_name":"image_2","name":"image_2","shape":7,"type":"IMAGE","link":null},{"localized_name":"audio_1","name":"audio_1","shape":7,"type":"AUDIO","link":35},{"localized_name":"audio_2","name":"audio_2","shape":7,"type":"AUDIO","link":null},{"localized_name":"video_1","name":"video_1","shape":7,"type":"VIDEO","link":null},{"localized_name":"video_2","name":"video_2","shape":7,"type":"VIDEO","link":null}],"outputs":[{"localized_name":"references","name":"references","type":"VIDEO_REFERENCES","links":null}],"properties":{"Node name for S&R":"VLLMOmniVideoReferences"},"widgets_values":[]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[83.24620208113915,422.08543477542605],"size":[284.205078125,226],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scLine truncated
|
||||
+42
-25
@@ -71,8 +71,7 @@ nav:
|
||||
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
|
||||
- LoRA: user_guide/diffusion/lora.md
|
||||
- Custom Pipeline: features/custom_pipeline.md
|
||||
- Request-Level Batching: user_guide/diffusion/request_batching.md
|
||||
- Step Execution: user_guide/diffusion/step_execution.md
|
||||
- Execution Modes: user_guide/diffusion/execution_modes.md
|
||||
- Quantization:
|
||||
- Overview: user_guide/quantization/overview.md
|
||||
- Online Quantization: user_guide/quantization/online.md
|
||||
@@ -102,31 +101,49 @@ nav:
|
||||
- CI: contributing/ci
|
||||
- Design Documents:
|
||||
- design/index.md
|
||||
- design/architecture_overview.md
|
||||
- Architecture Documents:
|
||||
- Architecture Overview: design/architecture_overview.md
|
||||
- Feature Design:
|
||||
- design/feature/disaggregated_inference.md
|
||||
- design/feature/ray_based_execution.md
|
||||
- design/feature/omni_connectors/
|
||||
- design/feature/prefix_caching.md
|
||||
- design/feature/cfg_parallel.md
|
||||
- design/feature/expert_parallel.md
|
||||
- design/feature/sequence_parallel.md
|
||||
- design/feature/tensor_parallel.md
|
||||
- design/feature/vae_parallel.md
|
||||
- design/feature/hsdp.md
|
||||
- design/feature/cache_dit.md
|
||||
- design/feature/teacache.md
|
||||
- design/feature/skip_softmax.md
|
||||
- design/feature/async_chunk.md
|
||||
- design/feature/vae_parallel.md
|
||||
- design/feature/diffusion_step_execution.md
|
||||
- design/feature/diffusion_request_level_batching.md
|
||||
- design/feature/diffusion_continuous_batching.md
|
||||
- Runtime and Stage Execution:
|
||||
- Disaggregated Inference: design/feature/disaggregated_inference.md
|
||||
- Ray-based Execution: design/feature/ray_based_execution.md
|
||||
- Async Chunk: design/feature/async_chunk.md
|
||||
- Async Diffusion Output: design/feature/async_diffusion_output.md
|
||||
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
|
||||
- Automatic Prefix Caching: design/feature/prefix_caching.md
|
||||
- Communication:
|
||||
- OmniConnector Implementations:
|
||||
- Mooncake Store Connector: design/feature/omni_connectors/mooncake_store_connector.md
|
||||
- Mooncake Transfer Engine Connector: design/feature/omni_connectors/mooncake_transfer_engine_connector.md
|
||||
- Mori Transfer Engine Connector: design/feature/omni_connectors/mori_transfer_engine_connector.md
|
||||
- Shared Memory Connector: design/feature/omni_connectors/shared_memory_connector.md
|
||||
- Yuanrong Store Connector: design/feature/omni_connectors/yuanrong_connector.md
|
||||
- Yuanrong Transfer Engine Connector: design/feature/omni_connectors/yuanrong_transfer_engine_connector.md
|
||||
- Diffusion Acceleration:
|
||||
- Parallelism:
|
||||
- CFG-Parallel: design/feature/cfg_parallel.md
|
||||
- Expert Parallel: design/feature/expert_parallel.md
|
||||
- Hybrid Sharded Data Parallel (HSDP): design/feature/hsdp.md
|
||||
- Pipeline Parallel: design/feature/pipeline_parallel.md
|
||||
- Sequence Parallel: design/feature/sequence_parallel.md
|
||||
- Tensor Parallel: design/feature/tensor_parallel.md
|
||||
- VAE Patch Parallelism: design/feature/vae_parallel.md
|
||||
- Attention Backends:
|
||||
- Skip-Softmax: design/feature/skip_softmax.md
|
||||
- Quantization: user_guide/quantization/overview.md
|
||||
- Cache-DiT: design/feature/cache_dit.md
|
||||
- TeaCache: design/feature/teacache.md
|
||||
- Diffusion Continuous Batching: design/feature/diffusion_continuous_batching.md
|
||||
- Infrastructure and Performance:
|
||||
- Prometheus Metrics: design/metrics.md
|
||||
- Speech Generation Performance Optimizations: design/qwen3_omni_tts_performance_optimization.md
|
||||
- Module Design:
|
||||
- design/module/ar_module.md
|
||||
- design/module/dit_module.md
|
||||
- design/module/entrypoint_module.md
|
||||
- design/module/async_omni_architecture.md
|
||||
- Runtime Modules:
|
||||
- AR Module: design/module/ar_module.md
|
||||
- DIT Module: design/module/dit_module.md
|
||||
- Orchestration:
|
||||
- Entrypoint Module: design/module/entrypoint_module.md
|
||||
- AsyncOmni Architecture: design/module/async_omni_architecture.md
|
||||
- Docs Guide: contributing/DOCS_GUIDE.md
|
||||
- API Reference:
|
||||
- api/README.md
|
||||
|
||||
@@ -40,6 +40,7 @@ Committers have write access and merge rights. They typically have deep expertis
|
||||
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
|
||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
|
||||
- [@NickCao](https://github.com/NickCao): CI, Tests, Entrypoint, and TTS Model Support
|
||||
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
||||
- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
|
||||
|
||||
@@ -387,8 +387,8 @@ that support request-level batching with `step_execution` disabled.
|
||||
|
||||
Use this together with `max_num_seqs > 1` for bursty serving traffic. `0`
|
||||
disables admission waiting and preserves the lowest first-request latency.
|
||||
For diffusion request-level batching tuning, see
|
||||
[Request-Level Batching](../user_guide/diffusion/request_batching.md).
|
||||
For diffusion execution and batching tuning, see
|
||||
[Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md).
|
||||
|
||||
Default: `0.0`
|
||||
|
||||
|
||||
@@ -44,9 +44,11 @@ Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bul
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [500, 800],
|
||||
"mean_audio_ttfp_ms": [2000, 3500],
|
||||
"mean_audio_rtf": [0.25, 0.35]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [500, 800],
|
||||
"mean_audio_ttfp_ms": [2000, 3500],
|
||||
"mean_audio_rtf": [0.25, 0.35]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -142,10 +144,6 @@ pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
|
||||
--test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
|
||||
--test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
|
||||
|
||||
# Optional baseline assertion (default off)
|
||||
pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --assert-baseline \
|
||||
--test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
```
|
||||
|
||||
See also [Markers for Tests](#markers-for-tests) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
|
||||
@@ -202,4 +200,4 @@ You can add any benchmark running parameters you need here. For all optional par
|
||||
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
|
||||
| request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
||||
| max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). |
|
||||
| baseline | object | No | see above | Optional per-metric thresholds; keys must match benchmark output fields. Scalar, list (per sweep step), or object (keyed by concurrency or QPS string).
|
||||
| baseline | object | No | see below | Optional thresholds. Prefer hardware-nested form: `{"H100": {"throughput_qps": …}, "A3": {…}}` (selected via runtime hardware label). Legacy flat `{metric: threshold}` still works. Per metric: scalar, list (sweep step), or object keyed by concurrency / QPS.
|
||||
@@ -755,7 +755,8 @@ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
|
||||
|
||||
### Step Execution
|
||||
|
||||
See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
|
||||
See the detailed
|
||||
[Diffusion Continuous Batching design guide](../../design/feature/diffusion_continuous_batching.md).
|
||||
|
||||
Use this only when your pipeline can be split into stable request-scoped and
|
||||
step-scoped phases. The reference implementation is
|
||||
@@ -769,9 +770,8 @@ step-scoped phases. The reference implementation is
|
||||
Do not enable `step_execution=True` until those four methods are implemented
|
||||
and validated against the request-level path.
|
||||
|
||||
If you want the pipeline to work with the experimental batched step-wise path
|
||||
(`max_num_seqs > 1`), also see:
|
||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
||||
The same design guide covers the experimental batched step-wise path used when
|
||||
`max_num_seqs > 1`.
|
||||
|
||||
If you expose this in example scripts or recipes, keep it opt-in. Surface
|
||||
runtime features like `step_execution` as optional flags instead of silently
|
||||
|
||||
@@ -1,165 +1,232 @@
|
||||
# Continuous Batching for Step-Wise Diffusion
|
||||
# Diffusion Continuous Batching
|
||||
|
||||
!!! warning "Experimental Feature"
|
||||
This feature is experimental. It currently applies only to native
|
||||
diffusion pipelines running with `step_execution=True`.
|
||||
|
||||
This document describes the batching extension built on top of
|
||||
[Diffusion Step Execution](diffusion_step_execution.md). The base
|
||||
step-execution contract is unchanged. The batching work is mainly in the
|
||||
scheduler and runner layers.
|
||||
|
||||
## Why It Helps
|
||||
|
||||
Step-wise execution breaks a long denoise loop into scheduler-visible units.
|
||||
That gives the runtime a place to admit other compatible requests between
|
||||
steps instead of waiting for an entire request to finish.
|
||||
|
||||
This matters most in low-MFU or bursty serving scenarios:
|
||||
|
||||
- one request's denoise step may not fully saturate the GPU
|
||||
- several compatible requests can share the same denoise forward pass
|
||||
- throughput and device utilization can improve without changing request-local
|
||||
scheduler state
|
||||
|
||||
This is **not** a guaranteed single-request latency win. The main benefit is
|
||||
usually higher utilization and better throughput when the workload contains
|
||||
multiple in-flight compatible requests.
|
||||
This document describes the unified diffusion execution architecture, including
|
||||
request execution, request-level batching, step execution, continuous batching,
|
||||
and output delivery. For user-facing configuration and CLI examples, see
|
||||
[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md).
|
||||
|
||||
## Overview
|
||||
|
||||
With continuous batching enabled:
|
||||
`DiffusionEngine` selects an execution mode from `step_execution` and uses
|
||||
`max_num_seqs` as that mode's scheduler capacity:
|
||||
|
||||
- the scheduler may keep multiple compatible requests active at the same time
|
||||
- the runner packs request-local step state into one `InputBatch`
|
||||
- `denoise_step()` runs on that batch
|
||||
- `step_scheduler()` and `post_decode()` still run per request
|
||||
| Configuration | Engine mode | Scheduler | Execution |
|
||||
|---|---|---|---|
|
||||
| `step_execution=False`, `max_num_seqs=1` | `REQUEST_BATCH` | `RequestScheduler` | One complete request-level `forward()` |
|
||||
| `step_execution=False`, `max_num_seqs>1` | `REQUEST_BATCH` | `RequestScheduler` | One fused `forward()` over compatible requests |
|
||||
| `step_execution=True`, `max_num_seqs=1` | `STEP_BATCH` | `StepScheduler` | One request advanced one denoise step per scheduler tick |
|
||||
| `step_execution=True`, `max_num_seqs>1` | `STEP_BATCH` | `StepScheduler` | Compatible requests advanced together in step waves |
|
||||
|
||||
The current implementation is conservative:
|
||||
Serial request execution is handled inside `REQUEST_BATCH`; it is not a
|
||||
separate engine mode. Similarly, continuous batching is the multi-request
|
||||
configuration of `STEP_BATCH`, not an independent execution mode.
|
||||
|
||||
- only compatible requests are batched together
|
||||
- per-request progress and completion remain independent
|
||||
The engine performs the following initialization:
|
||||
|
||||
Here, "continuous batching" means the step-wise path enabled by
|
||||
`step_execution=True`. Request-mode `DiffusionRequestBatch` is static
|
||||
request-level batching for one full pipeline `forward()` call; it does not
|
||||
admit or remove requests between denoise steps.
|
||||
1. Resolve model-specific pre- and post-processing hooks.
|
||||
2. Select `REQUEST_BATCH` or `STEP_BATCH`.
|
||||
3. Construct the configured `DiffusionExecutor`.
|
||||
4. Construct `RequestScheduler` or `StepScheduler`, unless a scheduler was
|
||||
explicitly injected.
|
||||
5. Bind `execute_batch` or `execute_step`.
|
||||
|
||||
## Enablement
|
||||
`DiffusionEngine.make_engine()` owns startup warmup after engine construction.
|
||||
This keeps construction, backend selection, and warmup failure cleanup in one
|
||||
lifecycle.
|
||||
|
||||
Use `--step-execution` as the feature gate, then increase `--max-num-seqs`
|
||||
above `1` if you want batching:
|
||||
## Unified Output Stream
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
All asynchronous diffusion requests use the same per-request output stream,
|
||||
regardless of whether the caller exposes streaming results to the user.
|
||||
|
||||
```python
|
||||
async for output in engine.step_streaming(request):
|
||||
...
|
||||
```
|
||||
|
||||
`--max-num-seqs 1` keeps the step-wise path without enabling batching.
|
||||
When a request is admitted, the engine creates an
|
||||
`asyncio.Queue[DiffusionOutput]` keyed by its request ID. The busy loop sends
|
||||
both intermediate chunks and final results through that queue:
|
||||
|
||||
For a reproducible replay flow using the bundled serving benchmark, see the
|
||||
Qwen-Image replay commands in
|
||||
[`benchmarks/diffusion/README.md`](gh-file:benchmarks/diffusion/README.md)
|
||||
and
|
||||
[`benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md`](gh-file:benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md).
|
||||
```text
|
||||
request
|
||||
-> scheduler.add_request()
|
||||
-> scheduler.schedule()
|
||||
-> execute_batch() or execute_step()
|
||||
-> scheduler.update_from_output()
|
||||
-> per-request output queue
|
||||
-> step_streaming()
|
||||
```
|
||||
|
||||
## Scheduler
|
||||
Streaming callers forward every yielded output. Non-streaming callers consume
|
||||
the same generator to completion and return only its final output. This avoids
|
||||
maintaining separate future-based and queue-based result paths.
|
||||
|
||||
The scheduler derives its batch capacity from `max_num_seqs` through
|
||||
`max_num_running_reqs`.
|
||||
`DiffusionEngine.step()` and
|
||||
`async_add_req_and_wait_for_response()` remain deprecated compatibility
|
||||
wrappers. New integrations should consume `step_streaming()` or
|
||||
`async_add_req_and_stream_response()`.
|
||||
|
||||
Batch admission is gated by
|
||||
[`StepBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
|
||||
which is built from shape-sensitive and CFG-sensitive sampling fields. This is
|
||||
the core correctness rule for batching: requests are only co-batched when they
|
||||
share the same denoise tensor contract.
|
||||
### Completion and Cancellation
|
||||
|
||||
There are three important details:
|
||||
Output delivery and scheduler lifecycle are related but separate:
|
||||
|
||||
- `num_inference_steps` is not part of the key, so requests with different
|
||||
total step counts can still share a batch
|
||||
- requests also do not need to be at the same current denoise progress; active
|
||||
requests can continue batching even when their current step indices diverge
|
||||
- admission is still FIFO, so an incompatible request at the head of the
|
||||
waiting queue blocks later compatible requests
|
||||
- A terminal scheduler result always finalizes request state.
|
||||
- A disconnected or cancelled consumer removes only its delivery queue.
|
||||
- The scheduler may still finish and clean up the request after its queue has
|
||||
been removed.
|
||||
- Engine shutdown sends an error output to all remaining streams.
|
||||
|
||||
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. 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.
|
||||
This separation prevents cancelled consumers from leaving active request IDs in
|
||||
the scheduler.
|
||||
|
||||
The scheduler batching unit is one logical `OmniDiffusionRequest`. In the
|
||||
step-wise path, runtime tensor batching is represented as `StepInputBatch`. For
|
||||
request-mode prompt semantics, see
|
||||
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
|
||||
## Request-Batch Execution
|
||||
|
||||
## Runner
|
||||
Request-batch execution runs the complete diffusion pipeline for each scheduler
|
||||
wave. Each `OmniDiffusionRequest` remains an independent logical request with
|
||||
its own prompt, sampling parameters, seed, request ID, output, error, and abort
|
||||
state.
|
||||
|
||||
The runner keeps persistent per-request execution state in
|
||||
[`StepRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
|
||||
while the scheduler owns a separate lightweight
|
||||
[`SchedulerRequestState`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
for queueing and lifecycle tracking.
|
||||
With `max_num_seqs=1`, `execute_batch()` handles the serial request path. With a
|
||||
larger value, a pipeline can opt into fused execution by declaring:
|
||||
|
||||
For each step, the runner builds an
|
||||
[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the
|
||||
active request states:
|
||||
```python
|
||||
supports_request_batch = True
|
||||
```
|
||||
|
||||
- prompt embeddings and masks are normalized and padded
|
||||
- dynamic tensors such as `latents` and `timesteps` are gathered each step
|
||||
- buffers are reused when batch composition stays the same
|
||||
Its `forward()` method must accept `DiffusionRequestBatch` and return one
|
||||
`DiffusionOutput` per request. The runner validates the result count and maps
|
||||
outputs back to their original request IDs with `BatchRunnerOutput`.
|
||||
|
||||
The step-wise batched path is:
|
||||
Current in-tree request-batch implementations include the Flux, LTX-2, SD3, and
|
||||
Qwen-Image pipelines. Treat the pipeline capability flag in the source as
|
||||
authoritative because support continues to expand.
|
||||
|
||||
1. Run `prepare_encode()` for newly admitted requests.
|
||||
### Request-Batch Data Flow
|
||||
|
||||
- `DiffusionSchedulerOutput` carries scheduled request IDs and payloads.
|
||||
- `DiffusionRequestBatch` provides the pipeline-facing static batch.
|
||||
- `BatchRunnerOutput` carries per-request results back to the engine.
|
||||
|
||||
Request-local setup, such as seed handling and output/error mapping, remains
|
||||
per request. Cache refresh, homogeneous LoRA activation, and
|
||||
`pipeline.forward(req_batch)` can be shared by the batch.
|
||||
|
||||
### Compatibility and Admission
|
||||
|
||||
`RequestBatchSamplingParamsKey` contains fields that affect the tensor contract,
|
||||
including shape, guidance, output count, and LoRA identity. Only compatible
|
||||
requests can share a fused forward call.
|
||||
|
||||
Admission is FIFO and conservative. An incompatible request at the front of the
|
||||
waiting queue can prevent later compatible requests from joining the current
|
||||
wave.
|
||||
|
||||
`request_batch_max_wait_ms` optionally delays the first schedule operation of a
|
||||
new wave so bursty arrivals can accumulate. It applies only when:
|
||||
|
||||
- the pipeline supports request batching;
|
||||
- step execution is disabled;
|
||||
- no requests are currently running; and
|
||||
- the configured wait is greater than zero.
|
||||
|
||||
The wait ends when the queue reaches capacity, stabilizes briefly, reaches its
|
||||
deadline, or the engine stops.
|
||||
|
||||
## Step-Batch Execution
|
||||
|
||||
Step execution exposes denoising progress to the scheduler. A supporting
|
||||
pipeline implements four stateful operations:
|
||||
|
||||
| Operation | Responsibility |
|
||||
|---|---|
|
||||
| `prepare_encode(state)` | Validate input, encode prompts, initialize latents and timesteps, and create request-local scheduler state |
|
||||
| `denoise_step(input_batch, *, states=...)` | Run one denoise forward for the scheduler-provided request states |
|
||||
| `step_scheduler(state, noise_pred)` | Update latents and advance request progress |
|
||||
| `post_decode(state)` | Decode and postprocess a completed request or output boundary |
|
||||
|
||||
Persistent request state lives in `StepRequestState`. Pipeline-specific fields
|
||||
that do not belong in the shared contract should be stored in `state.extra`.
|
||||
Queueing and lifecycle metadata remain in the scheduler's request state.
|
||||
|
||||
Current native pipelines that explicitly enable step execution include
|
||||
Qwen-Image, HunyuanImage3, and Helios. Step execution alone does not imply
|
||||
continuous-batching support: Qwen-Image accepts batched step states.
|
||||
HunyuanImage3 accepts batched step states only when its resolved self-attention
|
||||
backend is `TORCH_SDPA`; otherwise it rejects groups larger than one request.
|
||||
Configure `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or
|
||||
`diffusion_attention_config.default.backend=TORCH_SDPA` when
|
||||
`max_num_seqs>1`. See the
|
||||
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
|
||||
for its validated configuration. Helios supports only a single active step
|
||||
request and must use `max_num_seqs=1`.
|
||||
|
||||
### Continuous Batching
|
||||
|
||||
When `max_num_seqs>1`, `StepScheduler` can keep multiple compatible requests
|
||||
active. The runner gathers their state into `InputBatch`, performs one batched
|
||||
denoise forward, then applies scheduler updates per request:
|
||||
|
||||
1. Receive transferred KV payloads, then run `prepare_encode()` for newly
|
||||
admitted requests.
|
||||
2. Build or refresh `InputBatch`.
|
||||
3. Run one batched `denoise_step(input_batch)`.
|
||||
4. Slice the batched `noise_pred` back per request.
|
||||
5. Run per-request `step_scheduler()`.
|
||||
6. Run `post_decode()` only for requests that finished denoising.
|
||||
7. Scatter updated latents back into persistent request state with
|
||||
[`scatter_latents()`](gh-file:vllm_omni/diffusion/worker/input_batch.py).
|
||||
3. Run one batched `denoise_step(input_batch, states=states)`.
|
||||
4. Slice noise predictions back to each request.
|
||||
5. Run each request's `step_scheduler()`.
|
||||
6. Run `post_decode()` at a chunk boundary or request completion.
|
||||
7. Scatter updated latents back into persistent request state.
|
||||
|
||||
This keeps the shared work limited to the denoise forward pass while preserving
|
||||
request-local scheduler state and outputs.
|
||||
`StepBatchSamplingParamsKey` protects the batched tensor contract. Requests can
|
||||
have different total step counts and current step indices, but shape-sensitive,
|
||||
CFG-sensitive, and LoRA fields must be compatible. FIFO head-of-line blocking
|
||||
also applies here.
|
||||
|
||||
## Engine
|
||||
Chunk-capable pipelines may emit intermediate results. Final-only step
|
||||
pipelines emit when the request finishes even if `streaming_output=True`.
|
||||
|
||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) provides
|
||||
the background loop and async add-request path needed for multiple requests to
|
||||
accumulate in the scheduler.
|
||||
## Model Author Guidelines
|
||||
|
||||
When `step_execution=True`, the engine routes work through the step-wise
|
||||
executor path. The continuous batching behavior is defined by scheduler-side
|
||||
compatibility gating and runner-side `StepInputBatch` packing.
|
||||
When converting a request-level pipeline:
|
||||
|
||||
## Current Limitations
|
||||
- Reuse the same helpers as `forward()` to avoid behavior drift.
|
||||
- Copy request-scoped scheduler state into `state.scheduler`.
|
||||
- Advance `state.step_index` only in `step_scheduler()`.
|
||||
- Keep model forward work in `denoise_step()`.
|
||||
- Keep latent mutation in `step_scheduler()`.
|
||||
- Keep final decoding equivalent to the tail of `forward()`.
|
||||
- Store masks, condition latents, and other request-local tensors in the
|
||||
request state rather than on the shared pipeline.
|
||||
|
||||
- Experimental feature; use `max_num_seqs=1` for the conservative single-request
|
||||
step path.
|
||||
- Only native pipelines that already support `step_execution=True`.
|
||||
- Only homogeneous batches keyed by `StepBatchSamplingParamsKey` are supported.
|
||||
- `cache_backend`, KV transfer, and other request-mode extras are not wired
|
||||
into the batched step-wise path yet.
|
||||
- Future work can relax the current same-shape restriction with richer
|
||||
heterogeneous batching policies such as bucketing or padded execution for
|
||||
different resolutions.
|
||||
Before setting `supports_step_execution = True`, validate output parity for the
|
||||
same seed and sampling parameters, scheduler isolation across concurrent
|
||||
requests, abort cleanup, progress reporting, and all supported CFG paths.
|
||||
|
||||
## Related Files
|
||||
Before setting `supports_request_batch = True`, validate the batch input
|
||||
contract, one-output-per-request result shape, request identity and error
|
||||
mapping, seeded concurrency, LoRA compatibility, and tensor IPC.
|
||||
|
||||
- Scheduler base:
|
||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
||||
- Scheduler interface:
|
||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Step scheduler:
|
||||
[`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
||||
- Runner:
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Input batch:
|
||||
[`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
||||
- Tests:
|
||||
[`tests/diffusion/test_diffusion_scheduler.py`](gh-file:tests/diffusion/test_diffusion_scheduler.py)
|
||||
## Limitations
|
||||
|
||||
- Request batching requires an explicit pipeline capability when
|
||||
`max_num_seqs>1`.
|
||||
- Step execution requires the four-method stateful pipeline contract.
|
||||
- Both batching paths currently require homogeneous compatibility keys.
|
||||
- FIFO scheduling can reduce batching opportunities.
|
||||
- `request_batch_max_wait_ms` trades first-request latency for burst
|
||||
coalescing.
|
||||
- All diffusion cache backends are currently unsupported in step mode.
|
||||
- KV transfer is supported for newly admitted step requests; some other
|
||||
request-mode extras remain unsupported.
|
||||
- Step continuous batching remains experimental; use `max_num_seqs=1` for the
|
||||
conservative step path.
|
||||
|
||||
## Related Implementation
|
||||
|
||||
- Engine: [`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py)
|
||||
- Model contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
|
||||
- Scheduler interface: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Request scheduler: [`vllm_omni/diffusion/sched/request_scheduler.py`](gh-file:vllm_omni/diffusion/sched/request_scheduler.py)
|
||||
- Step scheduler: [`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
||||
- Runner: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Request batch: [`vllm_omni/diffusion/worker/request_batch.py`](gh-file:vllm_omni/diffusion/worker/request_batch.py)
|
||||
- Step input batch: [`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
||||
@@ -1,175 +0,0 @@
|
||||
# Request-Level Batching for Diffusion
|
||||
|
||||
This document describes the request-mode batching path for diffusion pipelines.
|
||||
For end-user enablement and tuning, see
|
||||
[Request-Level Batching](../../user_guide/diffusion/request_batching.md).
|
||||
|
||||
This is separate from
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
Request-level batching runs one full pipeline `forward()` over a static batch of
|
||||
compatible requests. Step-wise continuous batching admits work between denoise
|
||||
steps when `step_execution=True`.
|
||||
|
||||
## Why It Helps
|
||||
|
||||
The request-level design avoids coupling several logical requests to one request
|
||||
object. This keeps request identity, abort/error handling, and per-request
|
||||
metadata unambiguous while still allowing one fused pipeline forward pass for
|
||||
bursty or concurrent traffic.
|
||||
|
||||
## Overview
|
||||
|
||||
With request-level batching enabled:
|
||||
|
||||
- each `OmniDiffusionRequest` contains one `prompt` and one `request_id`
|
||||
- the scheduler groups compatible waiting requests into one scheduler wave
|
||||
- `DiffusionRequestBatch` wraps the scheduled requests for pipeline `forward()`
|
||||
- batch-capable pipelines return `list[DiffusionOutput]`, one output per
|
||||
request
|
||||
- `BatchRunnerOutput` maps each result back to its original `request_id`
|
||||
|
||||
Pipelines opt in with `supports_request_batch = True` and a `forward()` method
|
||||
that accepts `DiffusionRequestBatch` and returns `list[DiffusionOutput]`.
|
||||
Pipelines that do not opt in keep the existing per-request execution path.
|
||||
|
||||
## Enablement
|
||||
|
||||
Request-level batching is the request-mode path, so `step_execution` must remain
|
||||
disabled. Increase `max_num_seqs` above `1` to let the scheduler keep multiple
|
||||
compatible requests active:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 4
|
||||
```
|
||||
|
||||
For bursty online ingress, `request_batch_max_wait_ms` can add a bounded
|
||||
admission wait before the first `schedule()` of a scheduler wave:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 4 \
|
||||
--request-batch-max-wait-ms 20
|
||||
```
|
||||
|
||||
`request_batch_max_wait_ms=0` disables this wait and is the default.
|
||||
|
||||
## Request Contract
|
||||
|
||||
`OmniDiffusionRequest` represents one logical request. It owns one prompt,
|
||||
sampling parameters, request id, and request-local metadata. Runtime batches are
|
||||
formed by the scheduler and represented separately from the request payload.
|
||||
|
||||
Runtime batching is represented by:
|
||||
|
||||
- [`DiffusionSchedulerOutput`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
for scheduled request ids and request payloads
|
||||
- [`DiffusionRequestBatch`](gh-file:vllm_omni/diffusion/worker/request_batch.py)
|
||||
for the pipeline-facing request batch
|
||||
- [`BatchRunnerOutput`](gh-file:vllm_omni/diffusion/worker/utils.py) for
|
||||
per-request results
|
||||
|
||||
`DiffusionRequestBatch` intentionally exposes compatibility properties such as
|
||||
`prompts`, `sampling_params`, `request_id`, and `kv_sender_info` so migrated
|
||||
pipelines can stay close to upstream code while using a batch-aware contract.
|
||||
|
||||
## Scheduler
|
||||
|
||||
The scheduler derives its capacity from `max_num_seqs` through
|
||||
`max_num_running_reqs`. It exposes waiting/running queue counters so the engine
|
||||
can decide whether admission wait is useful before scheduling a new wave.
|
||||
|
||||
Batch compatibility is controlled by
|
||||
[`RequestBatchSamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py). The key
|
||||
contains shape-sensitive and guidance-sensitive fields, including output count
|
||||
and LoRA identity. Requests with incompatible shapes, CFG settings, output
|
||||
counts, LoRA adapters, or LoRA scales are kept in separate batches.
|
||||
|
||||
Admission is conservative:
|
||||
|
||||
- the scheduler only batches compatible requests
|
||||
- FIFO ordering is preserved
|
||||
- an incompatible request at the head of the waiting queue blocks later
|
||||
compatible requests
|
||||
|
||||
## Engine
|
||||
|
||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) resolves
|
||||
request-batch capability during initialization from the configured pipeline
|
||||
class, including custom pipeline classes.
|
||||
|
||||
The capability check uses the pipeline class attribute
|
||||
`supports_request_batch = True`. Pipelines that set this attribute must implement
|
||||
a request-batch-compatible `forward()` contract and return one
|
||||
`DiffusionOutput` per request; the runner validates that return shape at runtime.
|
||||
|
||||
When the selected pipeline is batch-capable and `step_execution=False`, request
|
||||
mode routes scheduler waves through the batch executor path. Otherwise it keeps
|
||||
the per-request executor path.
|
||||
|
||||
The optional admission wait runs only when:
|
||||
|
||||
- request batching is supported
|
||||
- `step_execution=False`
|
||||
- `request_batch_max_wait_ms > 0`
|
||||
- no requests are currently running
|
||||
|
||||
The wait exits early when the waiting queue reaches capacity, when the queue is
|
||||
stable for a short window, when the deadline expires, or when the engine stops.
|
||||
|
||||
## Executor And Runner
|
||||
|
||||
The executor exposes two request-mode entries:
|
||||
|
||||
- `execute_request`: one worker call per scheduled request
|
||||
- `execute_batch`: one worker call for the whole `DiffusionSchedulerOutput`
|
||||
|
||||
On the batch path, the worker builds a `DiffusionRequestBatch` and runs the
|
||||
pipeline once. Request-local setup remains per request:
|
||||
|
||||
- KV transfer metadata
|
||||
- random generator and seed handling
|
||||
- request output/error/abort mapping
|
||||
|
||||
Shared batch setup happens once per batch when possible:
|
||||
|
||||
- cache refresh
|
||||
- LoRA activation for the homogeneous adapter key
|
||||
- pipeline `forward(req_batch)`
|
||||
|
||||
Large tensor IPC still uses the shared-memory packing path. The packer traverses
|
||||
both normal `RunnerOutput.result` wrappers and nested batch results so batched
|
||||
outputs do not fall back to pickle IPC for tensor payloads.
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- Only pipelines that declare the request-batch contract use fused batch
|
||||
execution.
|
||||
- Batches are homogeneous under `RequestBatchSamplingParamsKey`; heterogeneous resolution or
|
||||
incompatible guidance settings do not co-batch yet.
|
||||
- FIFO scheduling can reduce batching opportunities when an incompatible
|
||||
request is at the front of the queue.
|
||||
- `request_batch_max_wait_ms` improves burst coalescing but can add latency to
|
||||
the first request in a scheduler wave. Keep it small for latency-sensitive
|
||||
serving.
|
||||
- Step-wise continuous batching is documented separately and only applies when
|
||||
`step_execution=True`.
|
||||
|
||||
## Related Files
|
||||
|
||||
- Request object and request batch:
|
||||
[`vllm_omni/diffusion/request.py`](gh-file:vllm_omni/diffusion/request.py)
|
||||
- Scheduler interface:
|
||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Scheduler base:
|
||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
||||
- Engine:
|
||||
[`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py)
|
||||
- Worker runner:
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Executor interface:
|
||||
[`vllm_omni/diffusion/executor/abstract.py`](gh-file:vllm_omni/diffusion/executor/abstract.py)
|
||||
- Tests:
|
||||
[`tests/diffusion/test_diffusion_engine.py`](gh-file:tests/diffusion/test_diffusion_engine.py)
|
||||
@@ -1,131 +0,0 @@
|
||||
# Adding Step Execution Support for Diffusion Pipelines
|
||||
|
||||
This guide documents vLLM-Omni's stepwise diffusion contract for model authors
|
||||
and contributors implementing `step_execution=True` support for a diffusion
|
||||
pipeline.
|
||||
|
||||
For end-user enablement, supported models, and current limitations, see
|
||||
[Step Execution](../../user_guide/diffusion/step_execution.md).
|
||||
|
||||
This document describes the base step-execution contract only. For the
|
||||
experimental batching policy layered on top of the step-wise path, see
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
|
||||
## Current Support Scope
|
||||
|
||||
`step_execution` is **not** a generic diffusion toggle. It only works for
|
||||
pipelines that implement the segmented stateful contract in
|
||||
[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
|
||||
|
||||
This page is intentionally author-facing. Treat runtime enablement
|
||||
(`step_execution=True` in Python or `--step-execution` in serving) as an
|
||||
opt-in user knob layered on top of the implementation contract below.
|
||||
|
||||
Current in-tree support:
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
|----------|----------------|----------------|
|
||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
||||
|
||||
Current engine/runtime limitations:
|
||||
|
||||
- Continuous batching with `max_num_seqs > 1` is experimental and documented in
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
Keep `max_num_seqs=1` for the conservative single-request step path.
|
||||
- `cache_backend` is not supported in step mode.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
- Unsupported pipelines now fail early during model loading instead of failing on the first request.
|
||||
|
||||
## Execution Contract
|
||||
|
||||
Step mode is driven by four pipeline methods plus the shared mutable request
|
||||
state object:
|
||||
|
||||
- `prepare_encode(state)`: one-time request preparation.
|
||||
- `denoise_step(state)`: compute the noise prediction for the current step.
|
||||
- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
|
||||
- `post_decode(state)`: decode the final output after denoising is complete.
|
||||
|
||||
The state lives in
|
||||
[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
|
||||
as `StepRequestState`. Store request-scoped tensors there, or use
|
||||
`state.extra` for model-specific fields that do not justify extending the core
|
||||
dataclass.
|
||||
|
||||
The worker-side step loop lives in
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
|
||||
|
||||
1. `prepare_encode()` runs once for a new request.
|
||||
2. `denoise_step()` runs every scheduler tick.
|
||||
3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
|
||||
4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
|
||||
|
||||
## Recommended Split
|
||||
|
||||
When converting an existing request-level `forward()` pipeline, keep the split
|
||||
strict and mechanical:
|
||||
|
||||
| Request-level phase | Stepwise method | What belongs there |
|
||||
|---------------------|-----------------|--------------------|
|
||||
| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
|
||||
| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
|
||||
| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
|
||||
| VAE decode / postprocess | `post_decode()` | Final decode only |
|
||||
|
||||
Keep the stepwise path reusing the same helpers as the request-level path
|
||||
whenever possible. Reimplementing the denoise loop from scratch is the easiest
|
||||
way to introduce behavioral drift.
|
||||
|
||||
## Qwen-Image Reference
|
||||
|
||||
[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
|
||||
is the reference implementation and is split correctly for the current
|
||||
contract:
|
||||
|
||||
- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
|
||||
latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
|
||||
with `forward()`.
|
||||
- `prepare_encode()` deep-copies `self.scheduler` **after**
|
||||
`prepare_timesteps()` so request-specific scheduler state is isolated.
|
||||
- `denoise_step()` reuses `_build_denoise_kwargs()` plus
|
||||
`predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
|
||||
non-CFG behavior stay identical to the request-level path.
|
||||
- `step_scheduler()` only calls
|
||||
`scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
|
||||
and increments `state.step_index`.
|
||||
- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
|
||||
the normal `forward()` path.
|
||||
|
||||
That decomposition is the target pattern for future models.
|
||||
|
||||
## Rules For New Pipelines
|
||||
|
||||
- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
|
||||
`state.scheduler` during `prepare_encode()`.
|
||||
- Do not mutate `state.step_index` inside `denoise_step()`. Only
|
||||
`step_scheduler()` should advance the step.
|
||||
- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
|
||||
- If the request-level pipeline has condition latents, masks, or edit-specific
|
||||
tensors, store them in `state` or `state.extra`, not in global pipeline
|
||||
attributes.
|
||||
- Preserve CFG behavior by sharing the same helper path used by `forward()`.
|
||||
- Keep `post_decode()` equivalent to the tail of `forward()`.
|
||||
|
||||
## Validation Checklist
|
||||
|
||||
Before marking a pipeline as `supports_step_execution = True`, verify:
|
||||
|
||||
- Stepwise output matches request-level output for the same seed and sampling params.
|
||||
- Per-request scheduler state is isolated across concurrent requests.
|
||||
- Abort during denoise does not leak cached state.
|
||||
- `step_index` reported by `RunnerOutput` matches the scheduler progress.
|
||||
- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
|
||||
|
||||
## Related Files
|
||||
|
||||
- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
|
||||
- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
|
||||
- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
|
||||
@@ -0,0 +1,386 @@
|
||||
# Async Omni Output Materialization
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Overview](#overview)
|
||||
2. [Performance](#performance)
|
||||
3. [Architecture](#architecture)
|
||||
4. [Enabling the Feature](#enabling-the-feature)
|
||||
5. [Compatibility and Fallbacks](#compatibility-and-fallbacks)
|
||||
6. [Related Files](#related-files)
|
||||
|
||||
## Overview
|
||||
|
||||
Async Omni output materialization moves CPU-side construction of
|
||||
`OmniModelRunnerOutput` out of the autoregressive (AR) decode critical path.
|
||||
It complements [Async Chunk](async_chunk.md): async chunk pipelines partial
|
||||
outputs between stages, while async output materialization prevents the
|
||||
construction of those partial outputs from blocking the next decode step.
|
||||
|
||||
An AR stage must return sampled token IDs quickly because the scheduler needs
|
||||
them for the next decode step. However, an Omni stage can also produce hidden
|
||||
states, multimodal tensors, streaming inter-stage payloads, and connector
|
||||
metadata. Building these payloads can require device-to-host (D2H) copies,
|
||||
tensor slicing, flattening, and Python object construction.
|
||||
|
||||
Without this feature, all of that work runs inline in `sample_tokens()`:
|
||||
|
||||
```text
|
||||
sample tokens
|
||||
-> update decode state
|
||||
-> copy hidden and multimodal outputs to CPU
|
||||
-> build per-request payloads
|
||||
-> build the streaming wire payload
|
||||
-> collect connector signals
|
||||
-> return sampled tokens
|
||||
-> launch the next decode step
|
||||
```
|
||||
|
||||
With async output materialization, only work required by the next decode step
|
||||
stays on the critical path:
|
||||
|
||||
```text
|
||||
sample tokens
|
||||
-> update decode state
|
||||
-> snapshot output state and start asynchronous D2H copies
|
||||
-> register sampled tokens
|
||||
-> launch the next decode step
|
||||
|
||||
background output path
|
||||
-> wait for the payload snapshot
|
||||
-> build per-request payloads
|
||||
-> build the streaming wire payload
|
||||
-> collect connector signals
|
||||
-> construct OmniModelRunnerOutput
|
||||
```
|
||||
|
||||
This allows payload construction for step `N` to overlap with GPU decode work
|
||||
for step `N + 1`. It does not change model computation or generated output.
|
||||
|
||||
!!! note "Platform and validation scope"
|
||||
The performance and correctness validation in this document covers CUDA
|
||||
and ROCm. This is a validation scope, not a runtime platform guard: XPU
|
||||
inherits `GPUARModelRunner`, and MUSA selects `GPUARWorker`, so those
|
||||
backends may enter the asynchronous path when its other guards pass. XPU
|
||||
and MUSA have not been validated for this optimization. Ascend NPU uses a
|
||||
separate `NPUARModelRunner` path that still constructs
|
||||
`OmniModelRunnerOutput` and drains connector output synchronously before
|
||||
wrapping the result as an asynchronous runner output.
|
||||
|
||||
<p align="center">
|
||||
<img
|
||||
alt="Decode step gap before and after async output materialization"
|
||||
src="../../../source/architecture/qwen3-omni-async-output-step-gap.svg"
|
||||
width="100%"
|
||||
>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<em>
|
||||
Figure 1: Moving payload construction off the decode path reduces the
|
||||
observed gap between consecutive Talker steps from about 2.8 ms to 41 µs.
|
||||
Adapted from the
|
||||
<a href="https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization">
|
||||
Qwen3-Omni optimization blog
|
||||
</a>.
|
||||
</em>
|
||||
</p>
|
||||
|
||||
For Qwen3-Omni, the feature is enabled for both AR stages:
|
||||
|
||||
- **Thinker**: Defers hidden-state and multimodal payload materialization.
|
||||
- **Talker**: Runs its decode-state postprocess eagerly, then defers codec
|
||||
payload materialization. Because Code2Wav consumes codec codes rather than
|
||||
Talker hidden states, the Talker also skips unnecessary hidden-state D2H.
|
||||
- **Code2Wav**: Does not use this path because it is a generation stage rather
|
||||
than an AR stage.
|
||||
|
||||
Qwen3-TTS uses the same mechanism in its AR Talker stage. The Talker updates
|
||||
the decode state it needs for the next step before returning, then constructs
|
||||
the codec payload for Code2Wav in the background.
|
||||
|
||||
## Performance
|
||||
|
||||
In the controlled Qwen3-Omni optimization sweep, enabling async output
|
||||
materialization on top of CUDA Graph and async chunk produced the following
|
||||
results at concurrency 64:
|
||||
|
||||
| Configuration | Request throughput | Mean audio TTFP | Mean audio RTF |
|
||||
|---|---:|---:|---:|
|
||||
| CUDA Graph + async chunk | 9.3 req/s | 655 ms | 0.63 |
|
||||
| + async output materialization | 11.3 req/s | 631 ms | 0.47 |
|
||||
| Change | **+22%** | **-4%** | **-25%** |
|
||||
|
||||
The sweep used `Qwen/Qwen3-Omni-30B-A3B-Instruct`, 640 Seed-TTS English
|
||||
prompts, and one replica each for Thinker, Talker, and Code2Wav. Async output
|
||||
materialization recovers throughput by removing CPU payload construction from
|
||||
the decode critical path without regressing time to first audio packet (TTFP).
|
||||
See the
|
||||
[Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization)
|
||||
for the complete benchmark methodology and results.
|
||||
|
||||
The implementation was also validated against the synchronous safe-copy path.
|
||||
The text and audio output hashes matched, confirming that overlap does not
|
||||
change the generated result. See
|
||||
[PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476) for the
|
||||
correctness and profiling details.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Output Lifecycle
|
||||
|
||||
`OmniAsyncGPUModelRunnerOutput` extends vLLM's
|
||||
`AsyncGPUModelRunnerOutput`. It preserves the existing asynchronous
|
||||
sampled-token feedback while adding a background builder for the complete Omni
|
||||
output.
|
||||
|
||||
The output lifecycle is:
|
||||
|
||||
1. `GPUARModelRunner.sample_tokens()` samples tokens and performs the
|
||||
bookkeeping required by the next decode step.
|
||||
2. The runner snapshots step-local metadata such as request IDs, token spans,
|
||||
scheduler output, and query start locations.
|
||||
3. CUDA payload tensors are cloned before reusable model or CUDA Graph buffers
|
||||
can be overwritten.
|
||||
4. The cloned payload is copied to pinned CPU memory on a dedicated CUDA
|
||||
stream. A CUDA event records when the copy is ready.
|
||||
5. `OmniAsyncGPUModelRunnerOutput` starts a background thread that waits for
|
||||
the payload event and builds `OmniModelRunnerOutput`.
|
||||
6. The AR runner immediately registers the sampled-token CPU copy with the
|
||||
input batch and returns, allowing async scheduling to advance.
|
||||
7. When the engine calls `get_output()`, it joins the background builder,
|
||||
propagates any builder exception, and finalizes sampled tokens and logprobs
|
||||
through the upstream output implementation.
|
||||
|
||||
```text
|
||||
GPU / decode path Background output path
|
||||
|
||||
forward + sample
|
||||
|
|
||||
+-- clone output tensors
|
||||
+-- enqueue D2H copy --------------> wait for D2H event
|
||||
+-- register sampled tokens |
|
||||
+-- return async output +-- slice per request
|
||||
| +-- build multimodal payloads
|
||||
next decode step +-- partition streaming payloads
|
||||
+-- build the wire payload
|
||||
+-- drain connector signals
|
||||
+-- build OmniModelRunnerOutput
|
||||
```
|
||||
|
||||
### Safe State Snapshots
|
||||
|
||||
The next scheduler step can mutate runner state while the background builder is
|
||||
still active. Most step-local builder inputs are therefore detached snapshots.
|
||||
The snapshotted state includes:
|
||||
|
||||
- scheduler token counts and speculative-token metadata
|
||||
- request IDs and request-to-batch-index mappings
|
||||
- sampled token IDs, logprobs, and prompt logprobs
|
||||
- query start locations and scheduled-token spans
|
||||
- hidden states and multimodal output tensors
|
||||
- KV-connector and encoder-cache outputs captured for the step
|
||||
|
||||
CUDA tensor cloning is required because CUDA Graph and model output buffers can
|
||||
be reused by the next forward pass. The dedicated copy stream and pinned host
|
||||
buffers make the D2H transfer asynchronous; the snapshot retains its cloned
|
||||
CUDA sources until the transfer event completes.
|
||||
|
||||
### Live Runner State and Ownership
|
||||
|
||||
The complete builder input is not a detached snapshot. Connector output state
|
||||
remains runner-owned: `get_omni_connector_output()` drains live, per-cycle
|
||||
connector signals after the streaming payload is built. Correctness therefore
|
||||
depends on that builder being the only consumer that drains connector signals
|
||||
for its output cycle. Receive-side connector state written by the background
|
||||
receiver is coordinated by the connector mixin's `_lock`.
|
||||
`OmniAsyncGPUModelRunnerOutput.get_output()` joins the builder and is the
|
||||
completion and exception boundary before the materialized output is consumed.
|
||||
|
||||
The shared output builder also contains a full-payload accumulation branch
|
||||
that reads `self.requests`. That branch is not reachable when async Omni output
|
||||
materialization is enabled: this feature requires `async_chunk`, while
|
||||
`should_accumulate_full_payload_output()` returns `False` whenever
|
||||
`async_chunk` is enabled. The feature path partitions each step's payload into
|
||||
streaming inter-stage and client handoffs instead.
|
||||
|
||||
### Output Builder
|
||||
|
||||
The background builder performs the work that previously ran inline:
|
||||
|
||||
- resolves which requests need downstream payloads
|
||||
- converts hidden and multimodal tensors into per-request CPU payloads
|
||||
- applies model-specific payload processing
|
||||
- partitions per-step payloads into inter-stage and client streams
|
||||
- creates the tensor-only `multimodal_outputs` wire payload
|
||||
- drains live connector readiness signals after building the streaming payload
|
||||
- constructs the final `OmniModelRunnerOutput`
|
||||
|
||||
Input-side connector operations remain synchronous with model execution. In
|
||||
particular, receiving inputs and flushing previously completed connector
|
||||
outputs are not deferred because they affect scheduler-visible request state.
|
||||
|
||||
### Qwen3-Omni Stage Behavior
|
||||
|
||||
| Stage | Async output behavior | Reason |
|
||||
|---|---|---|
|
||||
| Thinker | Snapshots hidden states and multimodal outputs; builds the downstream payload in the background | The Talker needs the payload, but the next Thinker decode step only needs sampled-token feedback |
|
||||
| Talker | Runs lightweight postprocess eagerly; snapshots codec outputs; omits hidden states from the downstream payload | `hidden_states.last` is needed by the next Talker step, while Code2Wav only needs codec codes |
|
||||
| Code2Wav | Uses the normal generation-stage output path | Code2Wav is not executed by `GPUARModelRunner` |
|
||||
|
||||
The Qwen3-TTS Talker follows the same pattern as the Qwen3-Omni Talker:
|
||||
postprocess runs eagerly because the next decode step needs its updated state,
|
||||
while codec payload construction runs asynchronously. Its Code2Wav stage also
|
||||
uses the normal generation-stage path.
|
||||
|
||||
## Enabling the Feature
|
||||
|
||||
There is no separate `async_omni_output` command-line or YAML option. The
|
||||
feature is selected automatically for model stages that opt in when the runtime
|
||||
conditions are safe. On the validated CUDA and ROCm configurations, using the
|
||||
bundled deployment profile is enough to turn it on for the supported
|
||||
Qwen3-Omni and Qwen3-TTS pipelines. The selection logic has no CUDA/ROCm-only
|
||||
guard; see [Compatibility and Fallbacks](#compatibility-and-fallbacks) for the
|
||||
other platform paths. The same profile does not enable background Omni
|
||||
materialization on Ascend NPU.
|
||||
|
||||
### Example 1: Qwen3-Omni
|
||||
|
||||
Start the full Qwen3-Omni pipeline:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
The model registry automatically loads
|
||||
`vllm_omni/deploy/qwen3_omni_moe.yaml`. Its relevant settings are:
|
||||
|
||||
```yaml
|
||||
async_chunk: true
|
||||
|
||||
stages:
|
||||
- stage_id: 0 # Thinker
|
||||
enable_prefix_caching: false
|
||||
# async_scheduling defaults to true for AR stages
|
||||
|
||||
- stage_id: 1 # Talker
|
||||
enable_prefix_caching: false
|
||||
# async_scheduling defaults to true for AR stages
|
||||
|
||||
- stage_id: 2 # Code2Wav
|
||||
enable_prefix_caching: false
|
||||
async_scheduling: false
|
||||
```
|
||||
|
||||
With this profile, async output materialization activates automatically for:
|
||||
|
||||
- **Stage 0, Thinker**: Defers hidden-state and multimodal payload
|
||||
construction.
|
||||
- **Stage 1, Talker**: Defers codec payload construction after eagerly
|
||||
updating its decode state.
|
||||
|
||||
Stage 2 does not use this feature because Code2Wav is not an AR stage.
|
||||
|
||||
### Example 2: Qwen3-TTS
|
||||
|
||||
Start Qwen3-TTS with any checkpoint that uses the Qwen3-TTS Talker pipeline.
|
||||
For example:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
The model registry automatically loads `vllm_omni/deploy/qwen3_tts.yaml`. Its
|
||||
relevant settings are:
|
||||
|
||||
```yaml
|
||||
async_chunk: true
|
||||
|
||||
stages:
|
||||
- stage_id: 0 # Talker
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
|
||||
- stage_id: 1 # Code2Wav
|
||||
enable_prefix_caching: false
|
||||
```
|
||||
|
||||
Async output materialization activates only for Stage 0, the AR Talker. The
|
||||
Talker keeps the state required by its next decode step on the GPU, skips an
|
||||
unnecessary hidden-state D2H copy, and builds the codec payload for Code2Wav in
|
||||
the background. Stage 1 uses the generation-stage output path.
|
||||
|
||||
The same behavior applies to the Base and VoiceDesign Qwen3-TTS checkpoints
|
||||
because they use the same Talker implementation.
|
||||
|
||||
!!! warning
|
||||
Do not pass `--no-async-chunk` or enable prefix caching when you want this
|
||||
optimization. Either change causes the runner to fall back to synchronous
|
||||
output construction. Enabling async chunk alone cannot activate the feature
|
||||
for a model that has not opted into async Omni output.
|
||||
|
||||
!!! note "Prefix cache compatibility"
|
||||
Async Omni output materialization and Omni prefix caching cannot currently
|
||||
run together. `_should_use_async_omni_output()` returns `False` whenever
|
||||
`self.omni_prefix_cache` is present, so prefix caching selects synchronous
|
||||
output materialization. Supporting both features requires snapshotting or
|
||||
otherwise synchronizing prefix-cache merge and update state before the
|
||||
background output path can consume it safely.
|
||||
|
||||
## Compatibility and Fallbacks
|
||||
|
||||
`GPUARModelRunner` uses the async output path only when all of the following
|
||||
runtime conditions hold:
|
||||
|
||||
| Requirement | Reason |
|
||||
|---|---|
|
||||
| AR async scheduling is enabled | The optimization relies on the scheduler advancing while the prior output is materialized |
|
||||
| `async_chunk` is enabled | The feature targets incremental downstream Omni payloads |
|
||||
| The model stage opts in with `use_async_omni_output` | Models must declare that their output lifecycle is safe to defer |
|
||||
| Omni prefix cache is disabled | Prefix-cache merge and update ordering currently requires synchronous materialization |
|
||||
| Speculative decoding is disabled | Speculative output state is not included in this deferred path |
|
||||
| Routed-expert output is disabled | Routed-expert extraction currently requires the synchronous path |
|
||||
| Postprocess is absent or explicitly runs eagerly | State needed by the next decode step must be updated before the runner returns |
|
||||
|
||||
These checks are evaluated per stage on runners based on `GPUARModelRunner`.
|
||||
An unsupported combination does not prevent serving; it only falls back to
|
||||
synchronous output construction for that stage. CUDA and ROCm are the validated
|
||||
platforms. XPU inherits `GPUARModelRunner`, and MUSA selects `GPUARWorker`, so
|
||||
they may select the async path when these checks pass, but this optimization is
|
||||
not yet validated on either backend.
|
||||
|
||||
Ascend NPU does not use this selection path. `NPUARModelRunner.sample_tokens()`
|
||||
fully constructs `OmniModelRunnerOutput`, calls
|
||||
`get_omni_connector_output()`, and only then creates the asynchronous wrapper,
|
||||
so Omni payload materialization remains synchronous.
|
||||
|
||||
!!! note
|
||||
`use_async_omni_output`,
|
||||
`eager_omni_postprocess_before_async_output`, and
|
||||
`omni_pooler_payload_include_hidden` are model implementation contracts,
|
||||
not user-facing configuration fields.
|
||||
|
||||
## Related Files
|
||||
|
||||
- `vllm_omni/worker/gpu_ar_model_runner.py`: Async output object, safe tensor
|
||||
snapshots, compatibility guards, and deferred Omni output builder.
|
||||
- `vllm_omni/platforms/npu/worker/npu_ar_model_runner.py`: Separate Ascend NPU
|
||||
runner, where Omni output materialization remains synchronous.
|
||||
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Thinker and
|
||||
Talker opt-in behavior.
|
||||
- `vllm_omni/deploy/qwen3_omni_moe.yaml`: Default Qwen3-Omni deployment
|
||||
settings.
|
||||
- `vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py`: Qwen3-TTS
|
||||
Talker opt-in and eager postprocess behavior.
|
||||
- `vllm_omni/deploy/qwen3_tts.yaml`: Default Qwen3-TTS deployment settings.
|
||||
- `tests/worker/test_gpu_ar_model_runner.py`: Snapshot, guard, connector
|
||||
ordering, and background error-propagation tests.
|
||||
- [Async Chunk](async_chunk.md): Inter-stage chunking and scheduling design.
|
||||
- [Qwen3-Omni optimization blog](https://vllm.ai/blog/2026-07-01-qwen3-omni-optimization):
|
||||
Optimization context and controlled performance results.
|
||||
- [PR #4476](https://github.com/vllm-project/vllm-omni/pull/4476): Feature
|
||||
implementation, validation, and review history.
|
||||
+52
-5
@@ -1,6 +1,8 @@
|
||||
# Design Documents
|
||||
|
||||
This section contains design documents and architecture specifications for vLLM-Omni.
|
||||
This section contains design documents and architecture specifications for
|
||||
vLLM-Omni. The sidebar groups documents by the system concern they describe,
|
||||
while the source paths remain stable so existing links continue to work.
|
||||
|
||||
## Architecture Documents
|
||||
|
||||
@@ -8,18 +10,63 @@ This section contains design documents and architecture specifications for vLLM-
|
||||
|
||||
## Feature Design Documents
|
||||
|
||||
### Runtime and stage execution
|
||||
|
||||
- [Disaggregated Inference](feature/disaggregated_inference.md)
|
||||
- [Ray-based Execution](feature/ray_based_execution.md)
|
||||
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
||||
- [Request-Level Batching for Diffusion](feature/diffusion_request_level_batching.md)
|
||||
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
||||
- [Async Chunk](feature/async_chunk.md)
|
||||
- [Async Diffusion Output](feature/async_diffusion_output.md)
|
||||
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
|
||||
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
|
||||
|
||||
## Infrastructure Design Documents
|
||||
### Communication
|
||||
|
||||
#### OmniConnector implementations
|
||||
|
||||
- [Mooncake Store Connector](feature/omni_connectors/mooncake_store_connector.md)
|
||||
- [Mooncake Transfer Engine Connector](feature/omni_connectors/mooncake_transfer_engine_connector.md)
|
||||
- [Mori Transfer Engine Connector](feature/omni_connectors/mori_transfer_engine_connector.md)
|
||||
- [Shared Memory Connector](feature/omni_connectors/shared_memory_connector.md)
|
||||
- [Yuanrong Store Connector](feature/omni_connectors/yuanrong_connector.md)
|
||||
- [Yuanrong Transfer Engine Connector](feature/omni_connectors/yuanrong_transfer_engine_connector.md)
|
||||
|
||||
### Diffusion acceleration
|
||||
|
||||
#### Parallelism
|
||||
|
||||
- [CFG-Parallel](feature/cfg_parallel.md)
|
||||
- [Expert Parallel](feature/expert_parallel.md)
|
||||
- [Hybrid Sharded Data Parallel (HSDP)](feature/hsdp.md)
|
||||
- [Pipeline Parallel](feature/pipeline_parallel.md)
|
||||
- [Sequence Parallel](feature/sequence_parallel.md)
|
||||
- [Tensor Parallel](feature/tensor_parallel.md)
|
||||
- [VAE Patch Parallelism](feature/vae_parallel.md)
|
||||
|
||||
#### Attention Backends
|
||||
|
||||
- [Skip-Softmax](feature/skip_softmax.md)
|
||||
|
||||
#### Quantization
|
||||
|
||||
- [Quantization Overview](../user_guide/quantization/overview.md)
|
||||
|
||||
- [Cache-DiT](feature/cache_dit.md)
|
||||
- [TeaCache](feature/teacache.md)
|
||||
- [Diffusion Continuous Batching](feature/diffusion_continuous_batching.md)
|
||||
|
||||
## Infrastructure and Performance
|
||||
|
||||
- [Prometheus Metrics](metrics.md)
|
||||
- [Speech Generation Performance Optimizations](qwen3_omni_tts_performance_optimization.md)
|
||||
|
||||
## Module Design Documents
|
||||
|
||||
### Runtime Modules
|
||||
|
||||
- [AR Module](module/ar_module.md)
|
||||
- [DIT Module](module/dit_module.md)
|
||||
|
||||
### Orchestration
|
||||
|
||||
- [Entrypoint Module](module/entrypoint_module.md)
|
||||
- [AsyncOmni Architecture (Qwen3-Omni Example)](module/async_omni_architecture.md)
|
||||
@@ -252,8 +252,7 @@ while True:
|
||||
batch of compatible independent requests for one full pipeline forward when
|
||||
the pipeline supports request-level batching; the step-wise path can also
|
||||
admit or remove compatible requests between denoise steps. See
|
||||
[Request-Level Batching](../feature/diffusion_request_level_batching.md) and
|
||||
[Continuous Batching for Step-Wise Diffusion](../feature/diffusion_continuous_batching.md).
|
||||
[Diffusion Continuous Batching](../feature/diffusion_continuous_batching.md).
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -87,9 +87,8 @@ if __name__ == "__main__":
|
||||
|
||||
!!! info
|
||||
|
||||
For diffusion request-level batching controls such as `max_num_seqs` and
|
||||
`request_batch_max_wait_ms`, see
|
||||
[Request-Level Batching](../user_guide/diffusion/request_batching.md).
|
||||
For diffusion request batching, step execution, and streaming controls, see
|
||||
[Diffusion Execution Modes](../user_guide/diffusion/execution_modes.md).
|
||||
|
||||
For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
|
||||
|
||||
|
||||
@@ -44,7 +44,7 @@ th {
|
||||
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
|
||||
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, FL2VA, and Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
|
||||
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" width="1024" height="600" viewBox="0 0 1024 600" role="img" aria-labelledby="title desc">
|
||||
<title id="title">Qwen3-Omni Talker step gap before and after asynchronous output materialization</title>
|
||||
<desc id="desc">Before the optimization, synchronous D2H copy, per-request payload construction, streaming payload handoff, and connector-signal drain separate consecutive Talker forward and sample steps by about 2.8 milliseconds. After the optimization, that Omni output work runs on a background builder while the next Talker step executes, leaving a boundary of about 41 microseconds for snapshot, D2H enqueue, state update, and sampled-token feedback.</desc>
|
||||
<defs>
|
||||
<marker id="time-arrow" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="7" markerHeight="7" orient="auto">
|
||||
<path d="M 0 0 L 10 5 L 0 10 z" fill="#64748b"/>
|
||||
</marker>
|
||||
<marker id="red-end" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="5" markerHeight="5" orient="auto">
|
||||
<path d="M 0 0 L 10 5 L 0 10 z" fill="#b91c1c"/>
|
||||
</marker>
|
||||
<marker id="red-start" viewBox="0 0 10 10" refX="2" refY="5" markerWidth="5" markerHeight="5" orient="auto">
|
||||
<path d="M 10 0 L 0 5 L 10 10 z" fill="#b91c1c"/>
|
||||
</marker>
|
||||
<marker id="blue-end" viewBox="0 0 10 10" refX="8" refY="5" markerWidth="6" markerHeight="6" orient="auto">
|
||||
<path d="M 0 0 L 10 5 L 0 10 z" fill="#2563eb"/>
|
||||
</marker>
|
||||
<style>
|
||||
.canvas { fill: #ffffff; stroke: #e5e7eb; }
|
||||
.title { font: 700 22px system-ui, -apple-system, Segoe UI, sans-serif; fill: #111827; }
|
||||
.subtitle { font: 500 13px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
|
||||
.phase { font: 700 16px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1f2937; }
|
||||
.phase-sub { font: 500 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
|
||||
.lane { font: 600 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #374151; }
|
||||
.block-title { font: 700 10px system-ui, -apple-system, Segoe UI, sans-serif; letter-spacing: .25px; fill: #ffffff; }
|
||||
.block-op { font: 500 9.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
|
||||
.axis { font: 500 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #94a3b8; }
|
||||
.measure { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #b91c1c; }
|
||||
.inset-title { font: 700 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
|
||||
.inset-step { font: 600 10px system-ui, -apple-system, Segoe UI, sans-serif; fill: #ffffff; }
|
||||
.note-title { font: 700 12px system-ui, -apple-system, Segoe UI, sans-serif; fill: #1e3a8a; }
|
||||
.note-text { font: 500 10.5px system-ui, -apple-system, Segoe UI, sans-serif; fill: #475569; }
|
||||
.stat-cap { font: 600 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
|
||||
.stat-big { font: 800 16px system-ui, -apple-system, Segoe UI, sans-serif; fill: #111827; }
|
||||
.stat-sub { font: 500 11px system-ui, -apple-system, Segoe UI, sans-serif; fill: #64748b; }
|
||||
.before-card { fill: #eef4ff; stroke: #c7dbf6; }
|
||||
.after-card { fill: #eafaf0; stroke: #bfe7cf; }
|
||||
.talker { fill: #16a34a; }
|
||||
.sync-output { fill: #ea580c; }
|
||||
.async-output { fill: #2563eb; }
|
||||
.zoom-card { fill: #fff7f7; stroke: #f3caca; }
|
||||
.note-card { fill: #eff6ff; stroke: #bfdbfe; }
|
||||
.stat-blue { fill: #eef4ff; stroke: #c7dbf6; }
|
||||
.stat-green { fill: #eafaf0; stroke: #bfe7cf; }
|
||||
.stat-gray { fill: #f4f5f7; stroke: #e2e5ea; }
|
||||
.guide { stroke: #e2e8f0; stroke-width: 1; }
|
||||
.time-line { stroke: #64748b; stroke-width: 1.3; fill: none; marker-end: url(#time-arrow); }
|
||||
.measure-line { stroke: #b91c1c; stroke-width: 1.4; marker-start: url(#red-start); marker-end: url(#red-end); }
|
||||
.zoom-lead { stroke: #dc2626; stroke-width: 1.15; stroke-dasharray: 4 3; fill: none; }
|
||||
.handoff { stroke: #2563eb; stroke-width: 1.3; stroke-dasharray: 3 2; fill: none; marker-end: url(#blue-end); }
|
||||
</style>
|
||||
</defs>
|
||||
|
||||
<rect x="14" y="14" width="996" height="572" rx="14" class="canvas"/>
|
||||
|
||||
<text x="44" y="52" class="title">Async output: what moves off the Talker decode path</text>
|
||||
<text x="44" y="74" class="subtitle">Qwen3-Omni Talker stage · consecutive decode steps in a fixed profiler window</text>
|
||||
|
||||
<!-- ===================== BEFORE ===================== -->
|
||||
<rect x="44" y="94" width="144" height="110" rx="10" class="before-card"/>
|
||||
<text x="60" y="120" class="phase">Before</text>
|
||||
<text x="60" y="139" class="phase-sub">Omni output built</text>
|
||||
<text x="60" y="156" class="phase-sub">synchronously on the</text>
|
||||
<text x="60" y="173" class="phase-sub">AR critical path</text>
|
||||
|
||||
<text x="206" y="122" class="lane">AR path</text>
|
||||
<line x1="300" y1="88" x2="300" y2="210" class="guide"/>
|
||||
|
||||
<rect x="300" y="104" width="190" height="54" rx="5" class="talker"/>
|
||||
<text x="395" y="124" text-anchor="middle" class="block-title">TALKER STAGE · STEP N</text>
|
||||
<text x="395" y="145" text-anchor="middle" class="block-op">forward → sample</text>
|
||||
|
||||
<rect x="490" y="104" width="250" height="54" rx="5" class="sync-output"/>
|
||||
<text x="615" y="120" text-anchor="middle" class="block-title">SYNC OMNI OUTPUT · STEP N</text>
|
||||
<text x="615" y="138" text-anchor="middle" class="block-op">D2H → per-request slice / payload build</text>
|
||||
<text x="615" y="152" text-anchor="middle" class="block-op">streaming handoff → connector-signal drain</text>
|
||||
|
||||
<rect x="740" y="104" width="220" height="54" rx="5" class="talker"/>
|
||||
<text x="850" y="124" text-anchor="middle" class="block-title">TALKER STAGE · STEP N+1</text>
|
||||
<text x="850" y="145" text-anchor="middle" class="block-op">forward → sample</text>
|
||||
|
||||
<line x1="490" y1="174" x2="740" y2="174" class="measure-line"/>
|
||||
<text x="615" y="194" text-anchor="middle" class="measure">~2.8 ms GPU gap: output work is inline</text>
|
||||
|
||||
<!-- ===================== AFTER ===================== -->
|
||||
<rect x="44" y="218" width="144" height="140" rx="10" class="after-card"/>
|
||||
<text x="60" y="244" class="phase">After</text>
|
||||
<text x="60" y="263" class="phase-sub">Only step feedback</text>
|
||||
<text x="60" y="280" class="phase-sub">stays inline; payload</text>
|
||||
<text x="60" y="297" class="phase-sub">work runs in the</text>
|
||||
<text x="60" y="314" class="phase-sub">background</text>
|
||||
|
||||
<text x="206" y="254" class="lane">Talker GPU</text>
|
||||
<text x="206" y="326" class="lane">Output builder</text>
|
||||
<line x1="300" y1="212" x2="300" y2="370" class="guide"/>
|
||||
|
||||
<rect x="300" y="226" width="235" height="54" rx="5" class="talker"/>
|
||||
<text x="418" y="246" text-anchor="middle" class="block-title">TALKER STAGE · STEP N</text>
|
||||
<text x="418" y="267" text-anchor="middle" class="block-op">forward → sample</text>
|
||||
|
||||
<rect x="540" y="226" width="235" height="54" rx="5" class="talker"/>
|
||||
<text x="658" y="246" text-anchor="middle" class="block-title">TALKER STAGE · STEP N+1</text>
|
||||
<text x="658" y="267" text-anchor="middle" class="block-op">forward → sample</text>
|
||||
|
||||
<path d="M 535 280 L 548 298" class="handoff"/>
|
||||
<rect x="540" y="298" width="360" height="58" rx="5" class="async-output"/>
|
||||
<text x="720" y="315" text-anchor="middle" class="block-title">BACKGROUND OMNI OUTPUT · STEP N</text>
|
||||
<text x="720" y="334" text-anchor="middle" class="block-op">wait for D2H → per-request slice / payload build</text>
|
||||
<text x="720" y="349" text-anchor="middle" class="block-op">streaming handoff → connector-signal drain</text>
|
||||
|
||||
<line x1="300" y1="366" x2="930" y2="366" class="time-line"/>
|
||||
<text x="922" y="382" text-anchor="end" class="axis">time</text>
|
||||
|
||||
<!-- leads from the real 41 us boundary to the magnified inset -->
|
||||
<line x1="535" y1="280" x2="485" y2="424" class="zoom-lead"/>
|
||||
<line x1="540" y1="280" x2="525" y2="424" class="zoom-lead"/>
|
||||
|
||||
<!-- ===================== 41 µs BOUNDARY ===================== -->
|
||||
<rect x="300" y="386" width="400" height="92" rx="10" class="zoom-card"/>
|
||||
<text x="316" y="404" class="inset-title">Zoom · Talker step boundary (magnified)</text>
|
||||
|
||||
<rect x="320" y="416" width="165" height="26" rx="4" class="talker"/>
|
||||
<rect x="525" y="416" width="155" height="26" rx="4" class="talker"/>
|
||||
<text x="403" y="433" text-anchor="middle" class="inset-step">Talker step N</text>
|
||||
<text x="603" y="433" text-anchor="middle" class="inset-step">Talker step N+1</text>
|
||||
|
||||
<line x1="485" y1="452" x2="525" y2="452" class="measure-line"/>
|
||||
<text x="505" y="469" text-anchor="middle" class="measure">~41 µs</text>
|
||||
|
||||
<rect x="720" y="386" width="240" height="92" rx="10" class="note-card"/>
|
||||
<text x="736" y="405" class="note-title">Inline at the ~41 µs boundary</text>
|
||||
<text x="736" y="424" class="note-text">• eager Talker state update</text>
|
||||
<text x="736" y="441" class="note-text">• output snapshot + D2H enqueue</text>
|
||||
<text x="736" y="458" class="note-text">• sampled-token feedback</text>
|
||||
|
||||
<!-- ===================== SUMMARY ===================== -->
|
||||
<rect x="44" y="494" width="296" height="70" rx="9" class="stat-blue"/>
|
||||
<text x="60" y="514" class="stat-cap">Step-to-step gap</text>
|
||||
<text x="60" y="538" class="stat-big">2.8 ms → 41 µs</text>
|
||||
<text x="60" y="555" class="stat-sub">~68× tighter</text>
|
||||
|
||||
<rect x="356" y="494" width="296" height="70" rx="9" class="stat-green"/>
|
||||
<text x="372" y="514" class="stat-cap">Talker GPU</text>
|
||||
<text x="372" y="538" class="stat-big">next step starts immediately</text>
|
||||
<text x="372" y="555" class="stat-sub">decode workers stay busy</text>
|
||||
|
||||
<rect x="668" y="494" width="312" height="70" rx="9" class="stat-gray"/>
|
||||
<text x="684" y="514" class="stat-cap">Background builder</text>
|
||||
<text x="684" y="538" class="stat-big">materializes Omni output</text>
|
||||
<text x="684" y="555" class="stat-sub">payload work overlaps the next Talker step</text>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 9.8 KiB |
@@ -0,0 +1,252 @@
|
||||
# Diffusion Execution Modes
|
||||
|
||||
vLLM-Omni supports complete-request execution and step-wise diffusion
|
||||
execution. Both modes use the same asynchronous engine output stream; the
|
||||
configuration controls how work is scheduled and whether intermediate outputs
|
||||
are exposed.
|
||||
|
||||
## Choose a Mode
|
||||
|
||||
| Goal | CLI configuration |
|
||||
|---|---|
|
||||
| Serial request execution | `--max-num-seqs 1` |
|
||||
| Fused request-level batching | `--max-num-seqs N` |
|
||||
| Single-request step execution | `--step-execution --max-num-seqs 1` |
|
||||
| Step-wise continuous batching | `--step-execution --max-num-seqs N` |
|
||||
| Chunked diffusion output | `--diffusion-streaming-output` |
|
||||
|
||||
`N` must be greater than `1` to allow batching. The selected pipeline must
|
||||
support request-level batching or batched-step execution as appropriate;
|
||||
single-request step support alone is not sufficient for step-wise continuous
|
||||
batching.
|
||||
|
||||
## Request Execution
|
||||
|
||||
Request execution is the default when `--step-execution` is omitted. The
|
||||
pipeline performs a complete `forward()` for each scheduler wave.
|
||||
|
||||
### Serial Requests
|
||||
|
||||
Use the conservative serial path with any request-mode pipeline:
|
||||
|
||||
```bash
|
||||
vllm serve MODEL --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 1
|
||||
```
|
||||
|
||||
### Request-Level Batching
|
||||
|
||||
Set `max_num_seqs` above one to combine compatible independent requests into a
|
||||
single pipeline forward:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 4
|
||||
```
|
||||
|
||||
For bursty traffic, add a small admission window:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 4 \
|
||||
--request-batch-max-wait-ms 20
|
||||
```
|
||||
|
||||
`--request-batch-max-wait-ms 0` is the default. A nonzero value can improve
|
||||
batch formation but adds up to that much latency before a new scheduler wave.
|
||||
It has no effect in step mode.
|
||||
|
||||
Each prompt remains a separate logical request. Do not submit a top-level list
|
||||
as one packed prompt. Concurrent serving requests are batched internally when
|
||||
their shapes, guidance settings, output counts, and LoRA settings are
|
||||
compatible.
|
||||
|
||||
Only pipelines that declare request-batch support accept
|
||||
`max_num_seqs>1` in request mode. Unsupported pipelines fail during engine
|
||||
initialization; use `max_num_seqs=1` instead.
|
||||
|
||||
## Step Execution
|
||||
|
||||
Step execution lets the scheduler advance and abort requests between denoise
|
||||
steps:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 1
|
||||
```
|
||||
|
||||
Set a larger capacity to allow compatible requests to share denoise waves:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
```
|
||||
|
||||
Step continuous batching is experimental. Start with `max_num_seqs=1` when
|
||||
validating a model or debugging correctness, then increase it for
|
||||
multi-request throughput.
|
||||
|
||||
Step execution is capability-based, not a generic switch for every diffusion
|
||||
model. Qwen-Image supports step-wise continuous batching. HunyuanImage3 also
|
||||
supports it, but only when its resolved self-attention backend is `TORCH_SDPA`;
|
||||
set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure
|
||||
`diffusion_attention_config.default.backend=TORCH_SDPA` before using
|
||||
`--max-num-seqs >1`. See the
|
||||
[HunyuanImage-3.0 recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/HunyuanImage-3.0-Instruct.md)
|
||||
for its validated configuration. Helios supports single-request step execution only: use
|
||||
`--step-execution --max-num-seqs 1` for Helios. Consult the selected pipeline's
|
||||
documentation and source for the latest support status.
|
||||
|
||||
## Streaming Output
|
||||
|
||||
Use `--diffusion-streaming-output` for a pipeline that can produce intermediate
|
||||
diffusion outputs:
|
||||
|
||||
```bash
|
||||
vllm serve BestWishYsh/Helios-Distilled --omni \
|
||||
--port 8000 \
|
||||
--diffusion-streaming-output
|
||||
```
|
||||
|
||||
Streaming output requires step execution. If
|
||||
`--diffusion-streaming-output` is set without `--step-execution`, the engine
|
||||
enables step execution automatically. Model initialization fails if the
|
||||
pipeline does not implement step execution.
|
||||
|
||||
Chunk-capable pipelines emit intermediate and final outputs through the same
|
||||
request stream. Final-only step pipelines emit only the final result even when
|
||||
streaming output is enabled.
|
||||
|
||||
The normal non-streaming serving path uses the same internal stream but drains
|
||||
it and returns only the final output. Users do not need to enable
|
||||
`--diffusion-streaming-output` to benefit from the unified engine lifecycle.
|
||||
|
||||
## Send Requests
|
||||
|
||||
Execution-mode flags configure the server. They do not change the
|
||||
OpenAI-compatible client request format. For example, after starting any of
|
||||
the non-streaming configurations above:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "a cat sitting on a windowsill",
|
||||
"size": "1024x1024",
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42
|
||||
}' | jq -r '.data[0].b64_json' | base64 -d > cat.png
|
||||
```
|
||||
|
||||
To create a request-level batch, send multiple independent requests
|
||||
concurrently. Do not put several prompts into one `prompt` field:
|
||||
|
||||
```bash
|
||||
for prompt in "a red fox" "a blue bird" "a green frog"; do
|
||||
curl -s -X POST http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "{\"prompt\":\"${prompt}\",\"size\":\"1024x1024\",\"seed\":42}" \
|
||||
> /dev/null &
|
||||
done
|
||||
wait
|
||||
```
|
||||
|
||||
The scheduler may batch these requests when their sampling parameters are
|
||||
compatible. See the
|
||||
[Image Generation API](../../serving/image_generation_api.md) for response
|
||||
formats and additional client examples. Streaming models can use
|
||||
model-specific streaming endpoints documented by their serving guide.
|
||||
|
||||
## Python API
|
||||
|
||||
The Python arguments mirror the CLI flags:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
max_num_seqs=4,
|
||||
request_batch_max_wait_ms=20.0,
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
[
|
||||
"a cup of coffee on a table",
|
||||
"a toy dinosaur on a sandy beach",
|
||||
"a fox waking up in bed and yawning",
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
Each list item passed to `Omni.generate()` becomes an independent logical
|
||||
request that the scheduler may batch with compatible items.
|
||||
|
||||
For step execution:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
step_execution=True,
|
||||
max_num_seqs=1,
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
## Stage Configuration
|
||||
|
||||
Use the equivalent engine arguments in a deployment YAML:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: diffusion
|
||||
engine_args:
|
||||
step_execution: false
|
||||
max_num_seqs: 4
|
||||
request_batch_max_wait_ms: 20
|
||||
```
|
||||
|
||||
For step execution, set `step_execution: true` and remove
|
||||
`request_batch_max_wait_ms`.
|
||||
|
||||
## CLI Reference
|
||||
|
||||
| Flag | Default | Effect |
|
||||
|---|---:|---|
|
||||
| `--step-execution` | disabled | Select step-wise scheduling |
|
||||
| `--max-num-seqs` | `1` for diffusion stages | Set request- or step-scheduler capacity |
|
||||
| `--request-batch-max-wait-ms` | `0` | Wait for burst coalescing in request mode |
|
||||
| `--diffusion-streaming-output` | disabled | Expose supported intermediate diffusion outputs and require step execution |
|
||||
|
||||
## Limitations and Troubleshooting
|
||||
|
||||
- Requests batch only when their compatibility-sensitive parameters match.
|
||||
- Different LoRA adapters or scales run in separate batches.
|
||||
- FIFO scheduling can cause an incompatible request to block later compatible
|
||||
requests.
|
||||
- All diffusion cache backends are unsupported in step mode. KV transfer is
|
||||
supported for newly admitted step requests; some other request-mode extras
|
||||
remain unsupported.
|
||||
- If request-mode startup reports that the pipeline does not support batching,
|
||||
use `--max-num-seqs 1`.
|
||||
- If step-mode startup mentions `prepare_encode()`, `denoise_step()`,
|
||||
`step_scheduler()`, or `post_decode()`, the pipeline does not implement the
|
||||
required step contract.
|
||||
|
||||
For implementation details and model-author guidance, see
|
||||
[Diffusion Continuous Batching](../../design/feature/diffusion_continuous_batching.md).
|
||||
@@ -1,119 +0,0 @@
|
||||
# Request-Level Batching
|
||||
|
||||
Request-level batching lets diffusion serving combine multiple compatible
|
||||
logical requests into one pipeline `forward()` call. Each prompt remains a
|
||||
separate request with its own `request_id`, sampling parameters, seed, output,
|
||||
error, and abort state. The scheduler decides which requests can run together.
|
||||
|
||||
!!! warning "Prompt List Semantics"
|
||||
|
||||
Diffusion request-level batching does not support a top-level packed
|
||||
list-prompt request. Submit multiple prompts as independent requests and let
|
||||
the scheduler batch compatible in-flight requests. Multimodal payloads stay
|
||||
inside a single prompt dict, for example
|
||||
`{"prompt": "...", "multi_modal_data": {"image": image}}`.
|
||||
|
||||
## Enablement
|
||||
|
||||
Request-level execution is the default diffusion request mode when
|
||||
`step_execution=False`. `max_num_seqs` controls the maximum number of compatible
|
||||
requests in one scheduler wave:
|
||||
|
||||
- `max_num_seqs=1` is the serial request path.
|
||||
- `max_num_seqs>1` allows fused request-level batching for pipelines that
|
||||
declare request-batch support.
|
||||
|
||||
Increase `max_num_seqs` above `1` to allow the request scheduler to keep more
|
||||
than one compatible request active:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 4
|
||||
```
|
||||
|
||||
For bursty online traffic, you can also set a small admission wait window. This
|
||||
lets the engine wait briefly before the first `schedule()` of a new wave so
|
||||
nearby compatible requests can arrive and share the same fused forward pass:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--max-num-seqs 4 \
|
||||
--request-batch-max-wait-ms 20
|
||||
```
|
||||
|
||||
`--request-batch-max-wait-ms 0` is the default and disables admission waiting,
|
||||
so there is no added wait latency.
|
||||
|
||||
For deploy YAMLs, configure the diffusion stage engine args:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: diffusion
|
||||
engine_args:
|
||||
max_num_seqs: 4
|
||||
request_batch_max_wait_ms: 20
|
||||
```
|
||||
|
||||
## Compatibility
|
||||
|
||||
Only pipelines that declare request-batch support can use the fused
|
||||
multi-request `pipeline.forward(batch)` path. Pipelines that do not support this
|
||||
contract can still use request-level execution with `max_num_seqs=1`, which is
|
||||
the serial request path. Setting `max_num_seqs>1` for a pipeline without
|
||||
request-batch support fails early during engine initialization.
|
||||
|
||||
The scheduler batches only compatible requests. Compatibility is based on
|
||||
shape-sensitive and guidance-sensitive sampling fields, including resolution,
|
||||
frame count, CFG settings, output count, and LoRA identity. Requests with
|
||||
different LoRA adapters or scales are kept in separate batches so the worker
|
||||
activates one adapter per batch.
|
||||
|
||||
Request-level batching applies when `step_execution=False`. For the separate
|
||||
step-wise runtime, see [Step Execution](step_execution.md).
|
||||
|
||||
## Tuning
|
||||
|
||||
- `max_num_seqs` caps the number of active compatible requests in one scheduler
|
||||
wave. `1` means serial request execution; values above `1` enable fused
|
||||
request-level batching when the pipeline supports it.
|
||||
- `request_batch_max_wait_ms` is an upper bound on extra admission wait before a
|
||||
new wave starts. Keep it small for latency-sensitive serving; values such as
|
||||
`10` to `50` ms are a practical starting range for bursty HTTP ingress.
|
||||
- `0` disables admission waiting and preserves the lowest first-request latency.
|
||||
- FIFO ordering is conservative: an incompatible request at the front of the
|
||||
waiting queue can block later compatible requests from joining the current
|
||||
batch.
|
||||
|
||||
## Python API
|
||||
|
||||
When constructing `Omni`, pass the same engine arguments:
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
max_num_seqs=4,
|
||||
request_batch_max_wait_ms=20.0,
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
[
|
||||
"a cup of coffee on a table",
|
||||
"a toy dinosaur on a sandy beach",
|
||||
"a fox waking up in bed and yawning",
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
`Omni.generate([...])` submits each list item as its own logical diffusion
|
||||
request. The runtime may batch those requests internally when their sampling
|
||||
parameters are compatible.
|
||||
|
||||
## For Contributors
|
||||
|
||||
For implementation details and model-author guidance, see
|
||||
[Request-Level Batching for Diffusion](../../design/feature/diffusion_request_level_batching.md).
|
||||
@@ -1,100 +0,0 @@
|
||||
# Step Execution
|
||||
|
||||
Step execution is an opt-in diffusion execution mode enabled with
|
||||
`step_execution=True` when constructing `Omni`.
|
||||
|
||||
It is not a generic diffusion toggle for every pipeline. Only pipelines that
|
||||
implement the stepwise contract support it today.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
step_execution=True,
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=50,
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
### Serving
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
```
|
||||
|
||||
For serving, `--step-execution` enables the step-wise runtime. Continuous
|
||||
batching only becomes relevant when `--max-num-seqs > 1`.
|
||||
|
||||
`max_num_seqs` controls the step scheduler capacity:
|
||||
|
||||
- `max_num_seqs=1` runs one request through the step-wise state machine.
|
||||
- `max_num_seqs>1` allows compatible active requests to share step waves.
|
||||
|
||||
## Supported Pipelines
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
|----------|----------------|----------------|
|
||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
||||
|
||||
!!! warning "Experimental continuous batching"
|
||||
When `--step-execution` is enabled and `max_num_seqs > 1` is configured,
|
||||
the step-wise path can batch
|
||||
compatible requests together. This is experimental. Requests with
|
||||
incompatible sampling parameters are intentionally kept in separate batches,
|
||||
and `max_num_seqs=1` remains the conservative single-request step path.
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- Continuous batching under `step_execution` is experimental and only batches
|
||||
compatible requests.
|
||||
- `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
|
||||
|
||||
Use step execution only when you specifically need the pipeline to run through
|
||||
its stepwise request state machine. For normal diffusion inference, leave it
|
||||
disabled unless your workflow depends on this mode.
|
||||
|
||||
For Qwen-Image online serving, the usual progression is:
|
||||
|
||||
- start with `--step-execution --max-num-seqs 1` if you only need the step-wise path
|
||||
- increase `--max-num-seqs` after that if you want the experimental compatible-request batching behavior
|
||||
|
||||
If you are looking for general diffusion speedups, see
|
||||
[Diffusion Features Overview](../diffusion_features.md).
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
If model loading fails with a message mentioning `prepare_encode()`,
|
||||
`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected
|
||||
pipeline does not support step execution.
|
||||
|
||||
## For Model Authors
|
||||
|
||||
If you want to add step execution support to a new diffusion pipeline, see the
|
||||
implementation guide:
|
||||
[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).
|
||||
|
||||
If you also want that pipeline to participate in the experimental batched
|
||||
step-wise path, see:
|
||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
||||
@@ -80,12 +80,11 @@ denoise steps.
|
||||
|
||||
| Method | Description | Best For |
|
||||
|--------|-------------|----------|
|
||||
| **[Request-Level Batching](diffusion/request_batching.md)** | Scheduler batches compatible independent diffusion requests into one pipeline forward pass | Bursty online serving and multi-request throughput |
|
||||
| **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control |
|
||||
| **[Diffusion Execution Modes](diffusion/execution_modes.md)** | Configures serial requests, request batching, step execution, continuous batching, and streaming output | Matching latency, throughput, cancellation, and output-delivery requirements |
|
||||
|
||||
**Note:** Request-level batching is available for pipelines that declare the
|
||||
request-batch forward contract. Step execution is currently supported by
|
||||
QwenImagePipeline only. See [Supported Models](#supported-models) for details.
|
||||
**Note:** Request-level batching and step execution are capability-based.
|
||||
Consult the execution guide and the selected pipeline's documentation for
|
||||
current support.
|
||||
|
||||
### Quantization Methods
|
||||
|
||||
@@ -160,7 +159,7 @@ The following tables show which models support each feature:
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
|
||||
| **MiniMax-H3** | ❌ | ✅ | ✅ | ❌ | ✅ (DiT/TE) | ❌ | ✅ | ✅ | ✅ (tile) | ✅ (DiT) | ❌ |
|
||||
|
||||
**Frame Interpolation Support**
|
||||
|
||||
@@ -206,7 +205,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). 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).
|
||||
6. Step Execution is not compatible with any diffusion cache backend. 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
|
||||
@@ -286,7 +285,7 @@ Measured on NVIDIA H800:
|
||||
|
||||
**Execution Modes:**
|
||||
|
||||
- **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
|
||||
- **[Diffusion Execution Modes](diffusion/execution_modes.md)** - Configure request batching, step execution, continuous batching, and streaming output
|
||||
|
||||
**Startup Optimization:**
|
||||
|
||||
|
||||
@@ -184,7 +184,7 @@ if __name__ == "__main__":
|
||||
!!! info
|
||||
|
||||
For diffusion request-level batching controls such as `max_num_seqs`, see
|
||||
[Request-Level Batching](../../diffusion/request_batching.md).
|
||||
[Diffusion Execution Modes](../../diffusion/execution_modes.md).
|
||||
|
||||
### Negative Prompts
|
||||
|
||||
|
||||
@@ -93,6 +93,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
|
||||
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
|
||||
| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
|
||||
| MiniMax H3 | `MiniMaxAI/MiniMax-H3` | Yes | Not validated | DiT linears except patch, timestep, and final projections; not compatible with layerwise offload | None | |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline PersonaPlex inference through the vllm-omni fullduplex backend.
|
||||
|
||||
This is the runnable end-to-end proof for the integration: it drives a user WAV
|
||||
through :class:`PersonaPlexEngine` + :class:`PersonaPlexSession` (our wrappers
|
||||
around the Moshi/Mimi lockstep loop) and writes the agent's audio + inner-monologue
|
||||
text. It mirrors ``python -m moshi.offline`` but goes through our ``FrameStepper``
|
||||
seam, so a green run validates that the session/adapter plumbing is faithful.
|
||||
|
||||
Prereqs: HF access to ``nvidia/personaplex-7b-v1`` (``export HF_TOKEN=...``). The
|
||||
stack is moshi-free (native temporal + depformer + streaming Mimi via
|
||||
``transformers``); decoding is greedy.
|
||||
|
||||
Example:
|
||||
python examples/offline_inference/personaplex/personaplex_offline.py \\
|
||||
--input-wav input_assistant.wav --output-wav out.wav \\
|
||||
--voice-prompt NATF2.pt --persona "You are a wise and friendly teacher."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vllm_omni.experimental.fullduplex.personaplex import (
|
||||
PersonaPlexConfig,
|
||||
PersonaPlexEngine,
|
||||
PersonaPlexSession,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.personaplex.config import DEFAULT_PERSONA
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--input-wav", required=True, help="User-side input WAV")
|
||||
parser.add_argument("--output-wav", required=True, help="Agent audio output WAV")
|
||||
parser.add_argument("--output-text", default=None, help="Optional JSON for agent text")
|
||||
parser.add_argument("--voice-prompt", default="NATF2.pt", help="Voice clone (bundled .pt or .wav)")
|
||||
parser.add_argument("--persona", default=DEFAULT_PERSONA, help="System role text")
|
||||
parser.add_argument("--hf-repo", default="nvidia/personaplex-7b-v1")
|
||||
parser.add_argument("--device", default="cuda")
|
||||
parser.add_argument("--cpu-offload", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
import sphn # local import: only the runnable driver needs the audio I/O dep
|
||||
|
||||
config = PersonaPlexConfig(
|
||||
hf_repo=args.hf_repo,
|
||||
voice_prompt=args.voice_prompt,
|
||||
persona=args.persona,
|
||||
device=args.device,
|
||||
cpu_offload=args.cpu_offload,
|
||||
)
|
||||
|
||||
engine = PersonaPlexEngine(config).load()
|
||||
session = PersonaPlexSession(engine, config)
|
||||
session.open()
|
||||
|
||||
# Load exactly like the reference (moshi.lm.load_audio): native read then
|
||||
# explicit resample, so greedy runs stay bit-faithful to `python -m moshi.offline`.
|
||||
user_pcm, src_sr = sphn.read(args.input_wav)
|
||||
user_pcm = sphn.resample(user_pcm, src_sample_rate=src_sr, dst_sample_rate=config.sample_rate)
|
||||
audio, text = session.run(np.asarray(user_pcm[0], dtype=np.float32))
|
||||
|
||||
sphn.write_wav(args.output_wav, audio, config.sample_rate)
|
||||
print(f"wrote {audio.shape[0] / config.sample_rate:.2f}s of agent audio to {args.output_wav}")
|
||||
print(f"agent text: {text!r}")
|
||||
if args.output_text:
|
||||
with open(args.output_text, "w") as f:
|
||||
json.dump({"text": text}, f, ensure_ascii=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -260,7 +260,7 @@ if __name__ == "__main__":
|
||||
For diffusion pipelines, the input list is sliced into single-item requests
|
||||
before feeding into the diffusion pipeline. For request-level batching
|
||||
controls such as `max_num_seqs`, see
|
||||
[Request-Level Batching](../../../docs/user_guide/diffusion/request_batching.md).
|
||||
[Diffusion Execution Modes](../../../docs/user_guide/diffusion/execution_modes.md).
|
||||
|
||||
### Negative Prompts
|
||||
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
# PersonaPlex full-duplex online serving
|
||||
|
||||
Serve [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
|
||||
(a Moshi-based full-duplex speech-to-speech model) with the native vLLM-Omni engine.
|
||||
|
||||
The server hosts the **official PersonaPlex web client** at `/` (auto-downloaded from
|
||||
the model repo) and implements its WebSocket protocol at `/api/chat`, so you talk to the
|
||||
model live in the browser — mic in, agent speech + inner-monologue text out.
|
||||
|
||||
> Experimental. Requires a GPU and Hugging Face access to the gated repo
|
||||
> (`HF_TOKEN` with access to `nvidia/personaplex-7b-v1`).
|
||||
|
||||
## Start the server
|
||||
|
||||
```bash
|
||||
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 bash run_server.sh
|
||||
# or directly:
|
||||
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 python -m \
|
||||
vllm_omni.experimental.fullduplex.personaplex.serving.server --port 8124
|
||||
```
|
||||
|
||||
## Talk to it
|
||||
|
||||
- **Browser (recommended):** open `http://localhost:8124/` and allow the microphone.
|
||||
Use headphones so the agent does not hear itself.
|
||||
- **Headless:** stream a 24 kHz mono WAV and save the reply:
|
||||
|
||||
```bash
|
||||
python duplex_client.py --url ws://localhost:8124/api/chat --input user.wav --out reply.wav
|
||||
```
|
||||
|
||||
## Unified Realtime API
|
||||
|
||||
The default `vllm_omni/deploy/personaplex.yaml` also enables main's engine-owned
|
||||
full-duplex control plane:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
|
||||
/path/to/personaplex-7b-v1 \
|
||||
--omni \
|
||||
--stage-configs-path vllm_omni/deploy/personaplex.yaml
|
||||
```
|
||||
|
||||
Validate the actual `/v1/realtime?duplex=1` scheduler path with paced 24 kHz
|
||||
PCM, two concurrent sessions, overflow admission, per-session slot recycling,
|
||||
and non-silent output:
|
||||
|
||||
```bash
|
||||
python tests/e2e/online_serving/personaplex_realtime_duplex.py \
|
||||
--model /path/to/personaplex-7b-v1 \
|
||||
--input-wav /path/to/speech.wav \
|
||||
--output-dir /tmp/personaplex-realtime-duplex
|
||||
```
|
||||
|
||||
This is distinct from the standalone `/api/chat` compatibility server above.
|
||||
The unified endpoint advertises `supports_barge_in=false`: overlapping speech
|
||||
is native model behavior, but destructive output interruption and model-state
|
||||
rewind have not been validated for PersonaPlex.
|
||||
|
||||
## Protocol (`/api/chat`)
|
||||
|
||||
Binary WebSocket messages, first byte is a tag (identical to `moshi.server`):
|
||||
|
||||
| Direction | Tag | Payload |
|
||||
| --- | --- | --- |
|
||||
| server → client | `\x00` | ready handshake (after the system-prompt prefill) |
|
||||
| client → server | `\x01` | Opus-encoded mic audio (24 kHz) |
|
||||
| server → client | `\x01` | Opus-encoded agent audio |
|
||||
| server → client | `\x02` | UTF-8 inner-monologue text |
|
||||
|
||||
Connect with query params `text_prompt` (persona) and `voice_prompt` (voice file).
|
||||
A raw-PCM endpoint (`/v1/audio/duplex`, JSON `open` + float32 frames) is also available
|
||||
for clients that do not want an Opus dependency.
|
||||
|
||||
## Notes
|
||||
|
||||
- **Run the client near the server.** Real-time 80 ms audio is sensitive to network
|
||||
latency/jitter; over a high-latency remote link playback can stutter regardless of
|
||||
engine speed. On localhost it is smooth.
|
||||
- Greedy decoding by default. One conversation per server by default (a second
|
||||
concurrent connection is rejected with WS close 1013 until the first ends);
|
||||
pass `--batch-size N` to host N concurrent conversations on one engine
|
||||
(elastic batching; connections beyond N are rejected until a slot frees).
|
||||
- Full runbook: [`recipes/NVIDIA/PersonaPlex.md`](../../../recipes/NVIDIA/PersonaPlex.md).
|
||||
@@ -0,0 +1,114 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Headless client for the PersonaPlex duplex server (official Moshi protocol).
|
||||
|
||||
Streams a 24 kHz mono WAV to ``/api/chat`` exactly as the browser client does
|
||||
(Opus over WebSocket, tags ``\\x00`` ready / ``\\x01`` audio / ``\\x02`` text),
|
||||
collects the agent's reply audio + inner-monologue text, and writes a WAV.
|
||||
|
||||
The browser UI at ``http://<host>:<port>/`` is the primary demo; this script is
|
||||
for headless testing / scripting.
|
||||
|
||||
python duplex_client.py --url ws://localhost:8124/api/chat \
|
||||
--input user.wav --out reply.wav
|
||||
|
||||
Deps: ``websockets``, ``sphn``, ``soundfile``, ``numpy``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import sphn
|
||||
|
||||
SR = 24000 # PersonaPlex / Mimi codec rate
|
||||
|
||||
|
||||
async def converse(url: str, pcm: np.ndarray) -> tuple[np.ndarray, str]:
|
||||
import websockets
|
||||
|
||||
reader = sphn.OpusStreamReader(SR) # decode agent audio
|
||||
writer = sphn.OpusStreamWriter(SR) # encode our mic audio
|
||||
text: list[str] = []
|
||||
ready = asyncio.Event()
|
||||
last = {"t": 0.0}
|
||||
|
||||
async with websockets.connect(url, max_size=None) as ws:
|
||||
|
||||
async def receive() -> None:
|
||||
loop = asyncio.get_event_loop()
|
||||
while True:
|
||||
try:
|
||||
msg = await ws.recv()
|
||||
except Exception:
|
||||
return
|
||||
if not isinstance(msg, (bytes, bytearray)) or not msg:
|
||||
continue
|
||||
tag, payload = msg[0], msg[1:]
|
||||
if tag == 0:
|
||||
ready.set()
|
||||
elif tag == 1:
|
||||
reader.append_bytes(bytes(payload))
|
||||
last["t"] = loop.time()
|
||||
elif tag == 2:
|
||||
text.append(payload.decode("utf8", errors="replace"))
|
||||
|
||||
recv_task = asyncio.create_task(receive())
|
||||
await asyncio.wait_for(ready.wait(), timeout=120) # waits for system-prompt prefill
|
||||
|
||||
frame = 1920 # 80 ms
|
||||
for i in range(0, len(pcm), frame):
|
||||
chunk = pcm[i : i + frame]
|
||||
if len(chunk) < frame:
|
||||
chunk = np.concatenate([chunk, np.zeros(frame - len(chunk), dtype=np.float32)])
|
||||
writer.append_pcm(chunk)
|
||||
data = writer.read_bytes()
|
||||
if data:
|
||||
await ws.send(b"\x01" + data)
|
||||
await asyncio.sleep(frame / SR) # pace at real time
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
last["t"] = loop.time()
|
||||
while loop.time() - last["t"] < 2.5: # drain until the agent stops talking
|
||||
await asyncio.sleep(0.1)
|
||||
recv_task.cancel()
|
||||
|
||||
chunks: list[np.ndarray] = []
|
||||
while True:
|
||||
out = reader.read_pcm()
|
||||
if out is None or out.shape[-1] == 0:
|
||||
break
|
||||
chunks.append(np.asarray(out, dtype=np.float32))
|
||||
audio = np.concatenate(chunks) if chunks else np.zeros(0, dtype=np.float32)
|
||||
return audio, "".join(text)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--url", default="ws://localhost:8124/api/chat")
|
||||
ap.add_argument("--input", required=True, help="24 kHz mono WAV of the user turn")
|
||||
ap.add_argument("--out", default="personaplex_reply.wav")
|
||||
ap.add_argument("--tail", type=float, default=10.0, help="trailing silence (s) for the agent reply")
|
||||
ap.add_argument("--persona", default="")
|
||||
ap.add_argument("--voice", default="NATF2.pt")
|
||||
args = ap.parse_args()
|
||||
|
||||
wav, sr = sf.read(args.input, dtype="float32")
|
||||
if wav.ndim > 1:
|
||||
wav = wav.mean(axis=1)
|
||||
if sr != SR:
|
||||
raise SystemExit(f"input must be {SR} Hz mono; got {sr} Hz. Resample it first.")
|
||||
wav = np.concatenate([wav, np.zeros(int(args.tail * SR), dtype=np.float32)])
|
||||
|
||||
url = f"{args.url}?text_prompt={args.persona}&voice_prompt={args.voice}"
|
||||
audio, text = asyncio.run(converse(url, wav))
|
||||
sf.write(args.out, audio if audio.size else np.zeros(SR, dtype=np.float32), SR)
|
||||
print(f"agent reply: {audio.size / SR:.2f}s -> {args.out}")
|
||||
print(f"inner-monologue: {text.strip()!r}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
# PersonaPlex full-duplex online serving.
|
||||
#
|
||||
# Serves the OFFICIAL PersonaPlex web client at / (auto-downloaded from
|
||||
# nvidia/personaplex-7b-v1) and the Moshi WebSocket protocol at /api/chat,
|
||||
# backed by the native vLLM-Omni engine. Requires HF access to the gated repo.
|
||||
|
||||
HOST="${HOST:-0.0.0.0}"
|
||||
PORT="${PORT:-8124}"
|
||||
|
||||
echo "Starting PersonaPlex duplex server on ${HOST}:${PORT}..."
|
||||
echo "Open http://localhost:${PORT}/ for the official web client (run it near the server)."
|
||||
|
||||
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
|
||||
--host "$HOST" --port "$PORT"
|
||||
@@ -71,6 +71,12 @@ plugins:
|
||||
- search
|
||||
- autorefs
|
||||
- awesome-nav
|
||||
- redirects:
|
||||
redirect_maps:
|
||||
user_guide/diffusion/request_batching.md: user_guide/diffusion/execution_modes.md
|
||||
user_guide/diffusion/step_execution.md: user_guide/diffusion/execution_modes.md
|
||||
design/feature/diffusion_request_level_batching.md: design/feature/diffusion_continuous_batching.md
|
||||
design/feature/diffusion_step_execution.md: design/feature/diffusion_continuous_batching.md
|
||||
- glightbox
|
||||
- git-revision-date-localized:
|
||||
# exclude files
|
||||
|
||||
+2
-2
@@ -70,8 +70,7 @@ dev = [
|
||||
"FlagEmbedding==1.4.0", # for hunyuan_image3 tests
|
||||
"kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
||||
"s3tokenizer==0.3.0", # for tts test
|
||||
"step-audio2==1.0.0", # Step-Audio2 Token2Wav (flashcosyvoice / hyperpyyaml)
|
||||
"stepaudio2-minicpmo", # for MiniCPM-o 4.5 talker vocoder (stepaudio2-minicpmo Token2wav)
|
||||
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
|
||||
]
|
||||
|
||||
demo = [
|
||||
@@ -109,6 +108,7 @@ docs = [
|
||||
"mkdocstrings-python",
|
||||
"mkdocs-gen-files",
|
||||
"mkdocs-awesome-nav",
|
||||
"mkdocs-redirects",
|
||||
"mkdocs-glightbox",
|
||||
"mkdocs-git-revision-date-localized-plugin",
|
||||
"mkdocs-minify-plugin",
|
||||
|
||||
@@ -332,6 +332,6 @@ bundled offline CLI do not currently expose `sigmas`.
|
||||
- <https://huggingface.co/Lightricks/LTX-2.3>
|
||||
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
|
||||
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
|
||||
- [Request batching](../../docs/user_guide/diffusion/request_batching.md)
|
||||
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
|
||||
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
|
||||
- [I2V offline example](../../examples/offline_inference/image_to_video/README.md)
|
||||
@@ -0,0 +1,61 @@
|
||||
# MiniMax-H3 on RTX 5090
|
||||
|
||||
This recipe uses BF16 weights, tiled VAE decode, tensor parallelism where a
|
||||
second GPU is available, and distributed layerwise offload (DLO). It is a
|
||||
memory-first serving configuration; lower resident counts reduce HBM use and
|
||||
increase CPU-to-GPU transfer time.
|
||||
|
||||
## Capacity requirements
|
||||
|
||||
| Resource | One RTX 5090 | Two RTX 5090s |
|
||||
| --- | ---: | ---: |
|
||||
| GPU HBM | 32 GiB | 32 GiB per GPU |
|
||||
| Checkpoint storage | 135 GiB per partition | 135 GiB per partition |
|
||||
| Available system RAM | 200 GiB minimum | 200 GiB minimum |
|
||||
| Recommended system RAM | 384 GiB | 384 GiB |
|
||||
|
||||
`FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one
|
||||
server at a time. DLO keeps rank-local weights in pinned host memory; increasing
|
||||
`--dlo-resident-layers` improves latency but does **not** reduce host RAM in the
|
||||
current implementation because resident layers retain pinned CPU master copies.
|
||||
|
||||
> **Modular H3:** after #5720 lands, preserve this recipe's one-partition
|
||||
> behavior with `--task-type fl2va` or `--task-type ref2va`.
|
||||
|
||||
## One RTX 5090: 1344x768, 5 seconds
|
||||
|
||||
Use 12 resident DiT layers. A 50-step B300 allocation test with this exact
|
||||
single-rank topology peaked at 26.50 GiB; re-measure peak HBM on the target
|
||||
card before increasing the resident count.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve /path/to/MiniMax-H3/FL2VA \
|
||||
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
|
||||
--num-gpus 1 --tensor-parallel-size 1 --text-encoder-tp-size 1 \
|
||||
--usp 1 --ring 1 --vae-patch-parallel-size 1 \
|
||||
--vae-parallel-mode tile --vae-use-tiling \
|
||||
--enable-distributed-layerwise-offload --dlo-no-use-allgather \
|
||||
--dlo-resident-layers 12 --enforce-eager \
|
||||
--diffusion-attention-backend CUDNN_ATTN
|
||||
```
|
||||
|
||||
## Two RTX 5090s: 1344x768, 5 seconds
|
||||
|
||||
Use TP2 and 20 resident DiT layers. The two-rank B300 capacity run peaked at
|
||||
27,726 MiB per rank for this shape and 50 steps. This is a memory/correctness
|
||||
proxy, not a consumer-GPU latency claim.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/MiniMax-H3/FL2VA \
|
||||
--omni --trust-remote-code --host 0.0.0.0 --port 8000 \
|
||||
--num-gpus 2 --tensor-parallel-size 2 --text-encoder-tp-size 2 \
|
||||
--usp 1 --ring 1 --vae-patch-parallel-size 2 \
|
||||
--vae-parallel-mode tile --vae-use-tiling \
|
||||
--enable-distributed-layerwise-offload --dlo-no-use-allgather \
|
||||
--dlo-resident-layers 20 --enforce-eager \
|
||||
--diffusion-attention-backend CUDNN_ATTN
|
||||
```
|
||||
|
||||
For Ref2VA, stop the FL2VA server and restart the same command with
|
||||
`/path/to/MiniMax-H3/Ref2VA`. Ref2VA reference video count and prompt length
|
||||
can increase activation memory; begin with one request at a time.
|
||||
@@ -0,0 +1,147 @@
|
||||
# MiniMax H3 on Moore Threads MUSA
|
||||
|
||||
> Joint video and audio generation on Moore Threads GPUs
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: MiniMaxAI
|
||||
- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
|
||||
- Tasks: T2VA, FL2VA, and Ref2VA
|
||||
- Mode: OpenAI-compatible `/v1/videos` HTTP serving
|
||||
- Hardware: 4x Moore Threads MTT S5000 for the validated Ref2VA profile
|
||||
- Maintainer: Community
|
||||
|
||||
This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for MUSA environments.
|
||||
MiniMax H3 has two independently served checkpoint partitions:
|
||||
|
||||
- `FL2VA` serves text-to-video+audio (`t2va`) and
|
||||
first-frame-to-video+audio (`fl2va`).
|
||||
- `Ref2VA` serves image/audio or video-reference generation (`ref2va`).
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### Checkpoint
|
||||
|
||||
Only download the partition needed by the server. Downloading the whole model
|
||||
repository also retrieves duplicate shared components and the original
|
||||
top-level pipeline layout.
|
||||
|
||||
For T2VA and FL2VA:
|
||||
|
||||
```bash
|
||||
python -m pip install modelscope
|
||||
export MODEL_ROOT=/path/to/MiniMax-H3
|
||||
modelscope download MiniMax/MiniMax-H3 \
|
||||
--local_dir "${MODEL_ROOT}" \
|
||||
--max-workers 16 \
|
||||
--include 'FL2VA/**'
|
||||
```
|
||||
|
||||
For Ref2VA, download `Ref2VA/**` instead.
|
||||
|
||||
### Environment
|
||||
|
||||
Install a compatible PyTorch/torch-musa, torchada, MATE, and vLLM-MUSA
|
||||
stack before installing vLLM-Omni. Importing vLLM-Omni should report that the
|
||||
`musa` vLLM platform plugin is active.
|
||||
|
||||
Install vLLM-Omni from a checkout containing MiniMax H3 support:
|
||||
|
||||
```bash
|
||||
python -m pip install -e .
|
||||
```
|
||||
|
||||
MiniMax H3 reference inputs and MP4 output require `soundfile`, `ffmpeg`, and
|
||||
`ffprobe`. The Python dependency is installed by vLLM-Omni; install the two
|
||||
executables with the operating system package manager and verify them:
|
||||
|
||||
```bash
|
||||
python -c 'import soundfile; print(soundfile.__version__)'
|
||||
ffmpeg -version
|
||||
ffprobe -version
|
||||
```
|
||||
|
||||
When torchaudio cannot load TorchCodec, vLLM-Omni automatically falls back to
|
||||
soundfile. Formats that libsndfile cannot read are demuxed through ffmpeg.
|
||||
|
||||
## Start a server
|
||||
|
||||
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
|
||||
and FL2VA, or to `Ref2VA` for Ref2VA.
|
||||
|
||||
The validated Ref2VA configuration uses tensor parallelism across four MTT
|
||||
S5000 GPUs and offloads inactive model components to CPU:
|
||||
|
||||
```bash
|
||||
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||
export PORT=8091
|
||||
|
||||
MUSA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
|
||||
vllm serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "${PORT}" \
|
||||
--trust-remote-code \
|
||||
--num-gpus 4 \
|
||||
--tensor-parallel-size 4 \
|
||||
--vae-patch-parallel-size 4 \
|
||||
--enable-cpu-offload \
|
||||
--vae-use-tiling \
|
||||
--diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
On MUSA, `FLASH_ATTN` selects the MATE/FlashAttention-3 implementation. Do
|
||||
not install the CUDA-only `fa4` optional dependency.
|
||||
|
||||
For request examples and parameter definitions, see
|
||||
[MiniMax-H3.md](MiniMax-H3.md#http-api-examples).
|
||||
|
||||
## MUSA validation
|
||||
|
||||
The following weight-independent gates have been validated on one MTT S5000:
|
||||
|
||||
- packed variable-length attention selects MATE FlashAttention-3 and returns
|
||||
finite BF16 output;
|
||||
- MiniMax H3 RoPE casts CPU FP64 position metadata to FP32 before device
|
||||
arithmetic and returns finite output;
|
||||
- the video-VAE seeded RNG context is deterministic on MUSA and restores both
|
||||
CPU and MUSA RNG states;
|
||||
- the soundfile fallback preserves stereo channel layout, FP32 samples, and
|
||||
the native sample rate when torchaudio/TorchCodec is unavailable.
|
||||
|
||||
Ref2VA image+audio generation was also validated end to end on one MTT S5000
|
||||
with CPU offload and VAE tiling. The smoke used a 448x256 reference image, a
|
||||
four-second 32 kHz stereo audio reference, two denoising steps, and seed 42.
|
||||
The synchronous HTTP request completed in 50.6 seconds and returned an MP4
|
||||
with the following decoded properties:
|
||||
|
||||
- 107 video frames at 448x256 and 24 FPS;
|
||||
- non-static video (`pixel_std=13.07`, `temporal_delta=0.49`);
|
||||
- finite stereo audio at 32 kHz (`audio_rms=0.0458`).
|
||||
|
||||
Official Ref2VA video-reference serving was validated with the expanded
|
||||
MiniMax prompt at 1344x768, 5 seconds, seed 0, and 50 steps. A true TP4 run on
|
||||
four MTT S5000 GPUs returned HTTP 200 in 638.6 seconds and produced 124 H.264
|
||||
frames at 24 FPS with AAC stereo audio at 32 kHz.
|
||||
|
||||
Keep `--vae-use-tiling` enabled for this serving profile.
|
||||
|
||||
## Known limitations
|
||||
|
||||
- Keep Ring Attention at degree 1. The current Ring path does not preserve
|
||||
MiniMax H3 packed padding boundaries.
|
||||
- H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
|
||||
- Each server process loads only one checkpoint partition.
|
||||
- H3 currently executes one generation request per diffusion batch.
|
||||
- FP8 quantization has not been enabled for MiniMax H3.
|
||||
- MP3, M4A, MP4, and reference-video audio fallback requires `ffmpeg` on
|
||||
`PATH`; WAV inputs can be read directly through soundfile.
|
||||
|
||||
## Additional resources
|
||||
|
||||
- [MiniMax-H3.md](MiniMax-H3.md) — full model and API guide
|
||||
- [MiniMax-H3-NPU.md](MiniMax-H3-NPU.md) — Ascend NPU deployment guide
|
||||
- [Supported models](../../docs/models/supported_models.md)
|
||||
- [Video API](../../docs/serving/videos_api.md)
|
||||
@@ -1,7 +1,7 @@
|
||||
# MiniMax H3 on Ascend NPU
|
||||
|
||||
> Joint video and audio generation with text, first-frame, image/audio, or
|
||||
> multi-video conditions — Ascend NPU deployment guide
|
||||
> Joint video and audio generation with text, first/last keyframes, and mixed
|
||||
> image/video/audio references — Ascend NPU deployment guide
|
||||
|
||||
## Summary
|
||||
|
||||
@@ -149,9 +149,9 @@ throughput guarantee.
|
||||
- H3 currently executes one generation request per diffusion batch.
|
||||
- The first regional-compile request is a warmup and should not be included
|
||||
in steady-state performance measurements.
|
||||
- Image+audio Ref2VA accepts exactly one image and one audio reference.
|
||||
- Video Ref2VA accepts one or more video files, but not an additional
|
||||
standalone audio reference.
|
||||
- The official H3 input matrix and media limits are documented in the [GPU
|
||||
recipe](MiniMax-H3.md#official-input-matrix-and-limits); this NPU path uses
|
||||
the same HTTP request contract.
|
||||
- VAE patch parallelism requires size 1 or the full DiT group size and
|
||||
supports the H3 native `tile` mode only.
|
||||
|
||||
|
||||
+257
-18
@@ -1,7 +1,7 @@
|
||||
# MiniMax H3
|
||||
|
||||
> Joint video and audio generation with text, first-frame, image/audio, or
|
||||
> multi-video conditions
|
||||
> Joint video and audio generation with text, first/last keyframes, and
|
||||
> mixed image/video/audio references
|
||||
|
||||
## Summary
|
||||
|
||||
@@ -16,8 +16,8 @@ checkpoint has two independently served partitions:
|
||||
|
||||
- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
|
||||
(`fl2va`)
|
||||
- `Ref2VA`: image+audio or one-or-more-video reference-to-video+audio
|
||||
(`ref2va`)
|
||||
- `Ref2VA`: up to 9 images, 3 videos, and 3 audio references in supported
|
||||
image/video/audio combinations (`ref2va`; audio-only is rejected)
|
||||
|
||||
The generated MP4 contains H.264 video and synchronized stereo audio.
|
||||
|
||||
@@ -33,8 +33,10 @@ export MODEL_ROOT=/path/to/MiniMax-H3
|
||||
hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
|
||||
```
|
||||
|
||||
Install vLLM-Omni from the checkout containing MiniMax H3 support. On
|
||||
Blackwell, install the optional FlashAttention-4 dependency:
|
||||
Install vLLM-Omni from the checkout containing MiniMax H3 support. The
|
||||
two-GPU RTX 5090/4090 profiles use cuDNN attention and do not need
|
||||
FlashAttention-4. Install the optional dependency only for the four-GPU
|
||||
B300/GB200 `FLASH_ATTN` profile:
|
||||
|
||||
```bash
|
||||
uv venv
|
||||
@@ -50,6 +52,26 @@ reference-video preparation and MP4 output.
|
||||
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
|
||||
and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
|
||||
|
||||
### Memory and storage requirements
|
||||
|
||||
Treat GPU HBM, host RAM, and checkpoint storage as separate requirements. Each
|
||||
H3 checkpoint partition (`FL2VA` or `Ref2VA`) contains about **134 GiB** of
|
||||
BF16 safetensors (about **135 GiB** on disk). Keeping both partitions locally
|
||||
therefore needs roughly **270 GiB** of model storage, although the examples
|
||||
serve only one partition at a time.
|
||||
|
||||
CPU offload and distributed layerwise offload reduce GPU residency; they do
|
||||
not make the model weights disappear. With `--dlo-no-use-allgather`, each
|
||||
worker retains its standard-loader rank-local weights in host memory, including
|
||||
pinned CPU buffers used for H2D streaming. Use at least **200 GiB available
|
||||
system RAM** before starting the two-GPU recipe; a **384 GiB host is
|
||||
recommended** to leave room for the OS, CUDA/PyTorch allocations, request
|
||||
inputs, and filesystem cache. Do not run the FL2VA and Ref2VA servers at the
|
||||
same time on a host sized for this minimum.
|
||||
|
||||
The consumer-GPU profiles below are HBM budgets only. They still require the
|
||||
host-RAM budget above.
|
||||
|
||||
### Single GPU: accuracy and memory first
|
||||
|
||||
The single-GPU configuration uses model-level CPU offload.
|
||||
@@ -77,6 +99,98 @@ Use a GPU with enough memory for the active H3 component and enough system RAM
|
||||
for the offloaded components. CPU offload reduces GPU memory pressure but adds
|
||||
PCIe/NVLink transfer latency.
|
||||
|
||||
### Two 24/32 GB GPUs: TP2 distributed layerwise offload
|
||||
|
||||
For two PCIe consumer GPUs, combine TP2 with distributed layerwise offload
|
||||
(DLO). The standard loader first creates the rank-local TP shard. DLO keeps
|
||||
that shard in pinned host memory and streams the 30 tail DiT blocks through a
|
||||
shared two-buffer window without a DP AllGather. The first 20 DiT blocks are
|
||||
copied to the GPUs once per denoise stage, reused by every sampling step, and
|
||||
released before VAE decode so the decoder can reuse their HBM.
|
||||
|
||||
```bash
|
||||
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||
export PORT=8091
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
vllm serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "${PORT}" \
|
||||
--trust-remote-code \
|
||||
--num-gpus 2 \
|
||||
--tensor-parallel-size 2 \
|
||||
--usp 1 \
|
||||
--ring 1 \
|
||||
--text-encoder-tp-size 2 \
|
||||
--vae-patch-parallel-size 2 \
|
||||
--vae-parallel-mode tile \
|
||||
--vae-use-tiling \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--dlo-no-use-allgather \
|
||||
--dlo-resident-layers 20 \
|
||||
--enforce-eager \
|
||||
--diffusion-attention-backend CUDNN_ATTN
|
||||
```
|
||||
|
||||
Use the profile that matches the per-GPU memory capacity:
|
||||
|
||||
| Profile | GPUs | Starting shape | Resident DiT blocks | Attention | Execution | Status |
|
||||
|---|---:|---:|---:|---|---|---|
|
||||
| `rtx5090` | 2 x 32 GB | 1344x768 | 20 | cuDNN attention | eager | Target-hardware validated |
|
||||
| `rtx4090` | 2 x 24 GB | 1024x576 | 12 | cuDNN attention | eager | Capacity-proxy starting point |
|
||||
|
||||
This topology uses all available parallel capacity: TP2 shards both the DiT
|
||||
and text encoder, `--dlo-no-use-allgather` streams each rank's local TP shard
|
||||
without reconstructing full blocks, and VAE patch parallelism splits tiled
|
||||
decode across both GPUs. cuDNN attention is selected explicitly for the RTX
|
||||
consumer path; the server stays eager to avoid an unqualified compile path.
|
||||
|
||||
The resident count changes placement and transfer frequency only; it does not
|
||||
quantize or change the BF16/FP32 denoise math. Re-measure peak memory before
|
||||
increasing it on a different request shape.
|
||||
|
||||
### RTX 5090 target-hardware validation
|
||||
|
||||
At vLLM-Omni commit `ae6577ea`, one full 50-step T2VA request completed on
|
||||
2 x RTX 5090 without OOM:
|
||||
|
||||
| Shape | Frames | Client E2E | Sampled peak/GPU | Output validation |
|
||||
|---:|---:|---:|---:|---|
|
||||
| 1344x768 | 124 at 24 FPS | 8 min 38 s | approximately 22.6 GiB | H.264 video + 32 kHz stereo AAC; full `ffmpeg` decode passed |
|
||||
|
||||
This is a single end-to-end validation run, not a warmed multi-run latency
|
||||
benchmark. The sampled `nvidia-smi` peak is also not a CUDA allocator
|
||||
high-water mark. The environment used vLLM 0.26.0, vLLM-Omni
|
||||
`0.26.1.dev14+gae6577ea`, and PyTorch 2.11.0+cu130. The
|
||||
[run record](https://github.com/lishunyang12/vllm-omni-rankings/blob/dcd06d7e83cb069842535918c0169ee9f3f29ba0/scripts/%E5%BE%AE%E4%BF%A1%E5%9B%BE%E7%89%87_20260805000034_86_237.png)
|
||||
captures the environment, output contract, elapsed time, and sampled peak.
|
||||
|
||||
Before the target run, both profiles were exercised on two B300 ranks as an
|
||||
allocation and correctness proxy. At 1344x768, 124 frames, and 50 steps, the
|
||||
20-layer profile peaked at 27,726 MiB per rank. At 1024x576, the 12-layer
|
||||
profile peaked at 18,888 MiB per rank in a 5-step capacity run. The resident
|
||||
and fully streamed placements produced identical decoded video-frame and audio
|
||||
hashes for the same shape, step count, prompt, and seed. The B300 result does
|
||||
not establish RTX 4090 PCIe latency; treat the 4090 profile as a conservative
|
||||
starting point until it is measured on that GPU.
|
||||
|
||||
To run T2VA, FL2VA, image+audio Ref2VA, and two-video Ref2VA in order, validate
|
||||
every MP4's H.264/AAC streams, and retain live server and GPU-memory logs:
|
||||
|
||||
```bash
|
||||
RUN_ROOT=/path/to/run-root \
|
||||
MODEL_ROOT=/path/to/MiniMax-H3 \
|
||||
GPU_IDS=0,1 \
|
||||
PROFILE=rtx5090 \
|
||||
bash examples/offline_inference/minimax_h3/run_h3_2gpu_all_tasks.sh
|
||||
```
|
||||
|
||||
The script selects 20 resident layers for `PROFILE=rtx5090` and 12 for
|
||||
`PROFILE=rtx4090`; `DLO_RESIDENT_LAYERS=N` overrides either default.
|
||||
|
||||
### Four GPUs: measured best-practice throughput
|
||||
|
||||
The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
|
||||
@@ -85,7 +199,8 @@ The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
|
||||
- Ulysses sequence parallelism degree 4;
|
||||
- native tiled VAE patch parallelism degree 4;
|
||||
- regional `torch.compile` for the repeated DiT blocks;
|
||||
- FlashAttention, with Ring and TP left at 1.
|
||||
- CuTe FlashAttention-4 through the `FLASH_ATTN` backend, with Ring and TP
|
||||
left at 1.
|
||||
|
||||
```bash
|
||||
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||
@@ -109,6 +224,9 @@ vllm serve "${MODEL}" \
|
||||
--diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
On Blackwell, `FLASH_ATTN` selects FA4. Confirm the server log contains
|
||||
`Using CuTe FlashAttention-4 on Blackwell` before recording measurements.
|
||||
|
||||
Do not add `--enforce-eager` to this performance configuration. The first
|
||||
request includes regional compilation; warm the server once before measuring
|
||||
steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
|
||||
@@ -164,6 +282,42 @@ command with:
|
||||
export MODEL="${MODEL_ROOT}/Ref2VA"
|
||||
```
|
||||
|
||||
### Online FP8 quantization
|
||||
|
||||
MiniMax H3 supports load-time FP8 quantization of the DiT. The checkpoint
|
||||
remains BF16 on disk; vLLM-Omni quantizes eligible weights while loading and
|
||||
uses dynamic activation scaling during inference. By default, attention and
|
||||
MLP linears in the token refiner and main DiT blocks, the condition
|
||||
projection, and all AdaLN projections use FP8. Patch, timestep, and final
|
||||
projections remain FP32; the text encoder and VAEs are unchanged.
|
||||
|
||||
Add this option to an existing H3 server command:
|
||||
|
||||
```bash
|
||||
--quantization fp8
|
||||
```
|
||||
|
||||
Use `ignored_layers` to keep any otherwise eligible linear in BF16. H3
|
||||
resolves the `transformer` component before constructing the DiT, so names do
|
||||
not start with `transformer.`. Entries are exact runtime linear prefixes; a
|
||||
parent name such as `blocks.0.attn` does not exclude its children.
|
||||
|
||||
Eligible names are the `attn.qkv_proj`, `attn.out_proj`, `mlp.fc1`, and
|
||||
`mlp.fc2` children under `token_refiner.blocks.<0-1>` or `blocks.<0-49>`.
|
||||
The other eligible names are `condition_proj`,
|
||||
`blocks.<0-49>.adaln_proj.linear`, and `final_layer.adaln_proj.linear`.
|
||||
For example, keep the first main block's attention projections in BF16 with:
|
||||
|
||||
```bash
|
||||
--diffusion-quantization-config \
|
||||
'{"transformer":{"method":"fp8","ignored_layers":["blocks.0.attn.qkv_proj","blocks.0.attn.out_proj"]}}'
|
||||
```
|
||||
|
||||
The structured option replaces `--quantization fp8`. Online FP8 is currently
|
||||
incompatible with H3 layerwise offload because the offload path produces a
|
||||
weight stride rejected by the Cutlass FP8 kernel. Use resident FP8 with tensor
|
||||
parallelism and VAE tiling instead.
|
||||
|
||||
## HTTP API examples
|
||||
|
||||
The following requests use the synchronous endpoint so the returned body can
|
||||
@@ -189,6 +343,7 @@ curl -sS -X POST "${API_URL}" \
|
||||
-F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \
|
||||
-F 'width=1344' \
|
||||
-F 'height=768' \
|
||||
-F 'aspect_ratio=16:9' \
|
||||
-F 'fps=24' \
|
||||
-F 'num_inference_steps=50' \
|
||||
-F 'flow_shift=12' \
|
||||
@@ -217,11 +372,31 @@ curl -sS -X POST "${API_URL}" \
|
||||
-o fl2va.mp4
|
||||
```
|
||||
|
||||
### 3. Ref2VA: image and audio reference
|
||||
Use the same `FL2VA` partition for the official tail-keyframe forms. A single
|
||||
image with `frame_indices=[-1]` conditions the last frame; two ordered images
|
||||
with `frame_indices=[0,-1]` condition the first and last frames:
|
||||
|
||||
Run this request against the `Ref2VA` partition. `audio_reference` accepts an
|
||||
HTTP(S) URL or a `data:` URL. In one terminal, expose the local reference
|
||||
assets to the serving host:
|
||||
```bash
|
||||
export LAST_FRAME=/path/to/fl2va_last_frame.png
|
||||
export FIRST_FRAME=/path/to/fl2va_first_frame.png
|
||||
|
||||
curl -sS -X POST "${API_URL}" \
|
||||
-F 'prompt=The subject moves naturally from the first image to the last image.' \
|
||||
-F 'num_inference_steps=50' \
|
||||
-F 'flow_shift=12' \
|
||||
-F 'seed=2102' \
|
||||
-F 'extra_params={"task":"fl2va","duration":8.7,"frame_indices":[0,-1],"audio_flow_shift":3.0}' \
|
||||
-F "input_references=@${FIRST_FRAME};type=image/png" \
|
||||
-F "input_references=@${LAST_FRAME};type=image/png" \
|
||||
-o fl2va_first_last.mp4
|
||||
```
|
||||
|
||||
### 3. Ref2VA: image-only, image/audio, or mixed references
|
||||
|
||||
Run these requests against the `Ref2VA` partition. Image-only Ref2VA omits
|
||||
`audio_reference`; adding one or more audio references is optional. The typed
|
||||
fields accept one object or an ordered JSON list. `audio_reference` accepts an
|
||||
HTTP(S) URL or a `data:` URL.
|
||||
|
||||
```bash
|
||||
python -m http.server 8092 \
|
||||
@@ -229,7 +404,24 @@ python -m http.server 8092 \
|
||||
--directory /path/to/reference_assets
|
||||
```
|
||||
|
||||
Then submit the image and audio request from another terminal:
|
||||
Then submit an image-only request from another terminal:
|
||||
|
||||
```bash
|
||||
export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
|
||||
|
||||
curl -sS -X POST "${API_URL}" \
|
||||
-F 'prompt=A white cat sits on a beige couch and slowly looks toward the camera.' \
|
||||
-F 'aspect_ratio=adaptive' \
|
||||
-F 'short_edge=768' \
|
||||
-F 'num_inference_steps=50' \
|
||||
-F 'flow_shift=12' \
|
||||
-F 'seed=3100' \
|
||||
-F 'extra_params={"task":"ref2va","duration":8.0,"audio_flow_shift":3.0}' \
|
||||
-F "input_reference=@${REF_IMAGE};type=image/png" \
|
||||
-o ref2va_image_only.mp4
|
||||
```
|
||||
|
||||
An image-plus-audio request is:
|
||||
|
||||
```bash
|
||||
export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
|
||||
@@ -252,7 +444,7 @@ curl -sS -X POST "${API_URL}" \
|
||||
The requested duration should cover the complete audio. If `duration` is
|
||||
shorter, the reference soundtrack is truncated to the generated clip.
|
||||
|
||||
### 4. Ref2VA: two video references
|
||||
### 4. Ref2VA: video, separate audio, and mixed references
|
||||
|
||||
Run this request against the `Ref2VA` partition. Repeat the
|
||||
`input_references` multipart field once per source video. H3 consumes the
|
||||
@@ -277,10 +469,40 @@ curl -sS -X POST "${API_URL}" \
|
||||
-o ref2va_video_video.mp4
|
||||
```
|
||||
|
||||
The server stores uploaded multi-video references only for the lifetime of the
|
||||
request and deletes the temporary files after generation. Video Ref2VA uses
|
||||
the source-video soundtracks and does not accept a separate
|
||||
`audio_reference`.
|
||||
The server stores uploaded references only for the lifetime of the request and
|
||||
deletes temporary files after generation. A video may use its embedded
|
||||
soundtrack, a separate `audio_reference`, or both. To send a mixed multipart
|
||||
request, repeat `input_references` for each image, video, or audio file; the
|
||||
server classifies them by MIME type and preserves the per-type order.
|
||||
|
||||
Reference videos must be MP4/MOV with H.264/H.265 video, optional AAC/MP3
|
||||
audio, 2–15 seconds each, and at most 15 seconds combined. They may still be
|
||||
longer than the generated clip. Use `start_time_seconds` to select a
|
||||
synchronized segment; for multiple typed video references, pass one value per
|
||||
video in `extra_params.start_time_seconds`.
|
||||
|
||||
Reference images accept JPG/JPEG, PNG, WEBP, HEIC, or HEIF up to 30 MiB. Standalone
|
||||
audio references accept WAV or MP3 up to 15 MiB, with 2–15 seconds per file and
|
||||
at most 15 seconds combined.
|
||||
|
||||
## Official input matrix and limits
|
||||
|
||||
| Task | Supported references | Limits |
|
||||
|------|----------------------|--------|
|
||||
| T2VA | text only | prompt must be non-empty |
|
||||
| FL2VA | first image, last image, or ordered first+last images | at most 2 images; `frame_indices` is `[0]`, `[-1]`, or `[0,-1]` |
|
||||
| Ref2VA | image-only, image+image, image+video, video+audio, and mixed image/video/audio | images ≤9, videos ≤3, audios ≤3, total references ≤12; audio requires a visual reference |
|
||||
|
||||
The H3 output contract is 4–15 seconds at 24 FPS, stereo 32 kHz audio, and a
|
||||
32-pixel canvas multiple. T2VA requires one named output ratio from `21:9`,
|
||||
`16:9`, `4:3`, `1:1`, `3:4`, or `9:16`. FL2VA always follows the first input
|
||||
image's ratio and ignores a generic `aspect_ratio` override. Ref2VA defaults to
|
||||
`16:9`; `adaptive` and SGLang's `auto` spelling are accepted aliases for that
|
||||
default. `short_edge` controls the 768-pixel canvas and must be `768`.
|
||||
`num_outputs_per_prompt` accepts 1–10 and derives each output seed as
|
||||
`seed + output_index`. The asynchronous endpoint returns all
|
||||
outputs; the synchronous raw-MP4 endpoint returns the first output when more
|
||||
than one is requested.
|
||||
|
||||
## Key parameters
|
||||
|
||||
@@ -293,7 +515,15 @@ the source-video soundtracks and does not accept a separate
|
||||
| `flow_shift` | `12` | Video sigma shift |
|
||||
| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
|
||||
| `seed` | Task-specific | Use a fixed value for reproducibility |
|
||||
| `width`, `height` | Multiples of 32 | Aspect ratio must be between 1:4 and 4:1 |
|
||||
| `aspect_ratio` | Task-specific | T2VA requires a named ratio; FL2VA follows the input image; Ref2VA defaults to `16:9` |
|
||||
| `short_edge` | `768` | H3 shape policy requires exactly 768 when `width`/`height` are omitted |
|
||||
| `num_outputs_per_prompt` | `1` | 1–10; async API returns every output |
|
||||
| `start_time_seconds` | `0` | Reference-video segment start; use a list in `extra_params` for multiple videos |
|
||||
| `width`, `height` | Multiples of 32 | Output aspect ratio must be between 1:4 and 4:1 |
|
||||
|
||||
## ComfyUI Frontend
|
||||
|
||||
Users can also use a ComfyUI frontend to interact with a hosted MiniMax-H3 service. The ComfyUI frontend can run in a separate environment or machine. Refer to [vLLM-Omni ComfyUI Integration](../../docs/features/comfyui.md) for details.
|
||||
|
||||
## Validated four-GPU evidence
|
||||
|
||||
@@ -310,12 +540,21 @@ throughput guarantee. Multi-video Ref2VA is much slower because the two
|
||||
reference videos expand both the Qwen3-VL vision sequence and the packed DiT
|
||||
attention sequence.
|
||||
|
||||
## Validated FP8 evidence
|
||||
|
||||
With eager DiT/text-encoder TP2 and VAE tiling, the 384x672, 107-frame,
|
||||
10-step quality case measured LPIPS 0.1156 (limit 0.20), PSNR 23.6316 dB,
|
||||
audio spectral cosine 0.9589 (minimum 0.80), and audio RMS ratio 0.9342. The
|
||||
resident per-GPU peak was 68.52 GiB for BF16 and 53.51 GiB for FP8, a 22%
|
||||
reduction.
|
||||
|
||||
## Known limitations
|
||||
|
||||
- Each server process loads only one checkpoint partition.
|
||||
- H3 currently executes one generation request per diffusion batch.
|
||||
- The first regional-compile request is a warmup and should not be included in
|
||||
steady-state performance measurements.
|
||||
- Online FP8 is not compatible with layerwise offload.
|
||||
- Image+audio Ref2VA accepts exactly one image and one audio reference.
|
||||
- Video Ref2VA accepts one or more video files, but not an additional standalone
|
||||
audio reference.
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
# PersonaPlex
|
||||
|
||||
> Live full-duplex speech-to-speech (talk to the model in your browser, it listens
|
||||
> while it speaks) with persona and voice control, on the native vLLM-Omni engine.
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: NVIDIA
|
||||
- Model: [`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
|
||||
(gated; a Moshi finetune)
|
||||
- Task: Full-duplex speech-to-speech. 24 kHz mic audio in, agent speech +
|
||||
inner-monologue text out, in 80 ms lockstep. Persona (role text) and voice
|
||||
(zero-shot clone) are set per session.
|
||||
- Mode: Live duplex WebSocket server with the official browser client, single or
|
||||
multi-session (`--batch-size`), plus an offline WAV-in / WAV-out example
|
||||
- Maintainer: [`@linyueqian`](https://github.com/linyueqian)
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to run `nvidia/personaplex-7b-v1` as a live voice agent on one GPU.
|
||||
The integration is **moshi-free**: the temporal transformer, depformer, and
|
||||
streaming Mimi codec are all native vLLM-Omni modules (Mimi runs on
|
||||
`transformers.MimiModel`), so no vendored fork is installed. Decoding is greedy,
|
||||
matching the reference implementation frame for frame on the golden replays used
|
||||
as the acceptance gate.
|
||||
|
||||
PersonaPlex is the first Moshi-class (pure-lockstep) model in the experimental
|
||||
full-duplex framework (`vllm_omni/experimental/fullduplex/`), and the first with
|
||||
elastic multi-session batching: one engine hosts `--batch-size` concurrent
|
||||
conversations, callers join or leave on any 80 ms tick without disturbing live
|
||||
slots (greedy isolation is bit-exact).
|
||||
|
||||
## References
|
||||
|
||||
- Serving package:
|
||||
[`vllm_omni/experimental/fullduplex/personaplex/`](../../vllm_omni/experimental/fullduplex/personaplex/)
|
||||
- Native model modules (temporal / depformer / streaming Mimi):
|
||||
[`vllm_omni/model_executor/models/personaplex/`](../../vllm_omni/model_executor/models/personaplex/)
|
||||
- Online example (server + headless client):
|
||||
[`examples/online_serving/personaplex/`](../../examples/online_serving/personaplex/)
|
||||
- Offline example:
|
||||
[`examples/offline_inference/personaplex/personaplex_offline.py`](../../examples/offline_inference/personaplex/personaplex_offline.py)
|
||||
- Staged offline pipeline (talker -> Mimi code2wav, async-chunk streaming):
|
||||
[`vllm_omni/deploy/personaplex.yaml`](../../vllm_omni/deploy/personaplex.yaml)
|
||||
- Integration PR: [vllm-project/vllm-omni#4771](https://github.com/vllm-project/vllm-omni/pull/4771)
|
||||
- Framework context: [RFC #3745](https://github.com/vllm-project/vllm-omni/issues/3745)
|
||||
(duplex adapter patterns), [#1335](https://github.com/vllm-project/vllm-omni/issues/1335)
|
||||
(full-duplex target), [PR #3907](https://github.com/vllm-project/vllm-omni/pull/3907)
|
||||
(fullduplex core contracts this adapter conforms to)
|
||||
- Upstream model card:
|
||||
[`nvidia/personaplex-7b-v1`](https://huggingface.co/nvidia/personaplex-7b-v1)
|
||||
|
||||
## Why it is different from MiniCPM-o / JoyVL duplex
|
||||
|
||||
| | MiniCPM-o 4.5 (#3907) | JoyVL | **PersonaPlex (this)** |
|
||||
|---|---|---|---|
|
||||
| Cadence | 1 s chunk groups | ~1 fps frames | **80 ms lockstep (12.5 Hz)** |
|
||||
| Turn control | learned `⟨listen⟩`/`⟨speak⟩` | `</silence>`/`</response>` | **none, pure lockstep** |
|
||||
| Per step | variable-length token group | text decision | **1 user frame in -> 1 agent frame + 1 text token** |
|
||||
| Barge-in | at chunk boundary | n/a | **native (model always hears the user)** |
|
||||
| Session state | chunk-group KV | per-tick HTTP | **persistent ring KV + streaming Mimi state** |
|
||||
|
||||
This is the lockstep ("parallel-frame joint") shape of the duplex adapter
|
||||
patterns: the adapter declares `DuplexCapability.continuous = True`, and
|
||||
`core.DuplexRuntime` runs ONE eternal response that consumes input frames as they
|
||||
arrive and drains on close. The flag is a small, default-off lifecycle mode, so
|
||||
turn-based adapters (JoyVL, MiniCPM-o) are unaffected.
|
||||
|
||||
## Architecture (Moshi RQ-Transformer, all native)
|
||||
|
||||
- **Mimi codec** 24 kHz, 12.5 Hz, 1920 samples/frame, 8 active codebooks
|
||||
(card 2048). Streaming encode + decode via `transformers.MimiModel` with the
|
||||
PersonaPlex reference checkpoint weights.
|
||||
- **Temporal transformer** (Helium backbone) 4096-d / 32 layers / 32 heads,
|
||||
sliding window 3000 frames (a 4-minute rolling context) held in a ring KV cache
|
||||
with per-slot recycle.
|
||||
- **Depformer** 1024-d / 6 layers / 16 heads, autoregressive over 16 codebooks
|
||||
per frame (8 vocoded), reset each frame.
|
||||
- Per-frame token column `[B, 17, 1]`: row 0 = inner-monologue text,
|
||||
rows 1-8 = agent audio, rows 9-16 = user audio.
|
||||
- Persona + voice are injected once at session open through the same lockstep
|
||||
step: voice clone forces the agent stream from reference-audio Mimi codes;
|
||||
persona forces the text stream from `<system> ... <system>` tokens.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
Verified on GPU (Hopper-class). The serving path is plain PyTorch eager plus the
|
||||
native modules, so other CUDA GPUs with enough memory are expected to work.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x 141 GB Hopper-class GPU (H20-class, verified)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.10+ (CI covers 3.11 / 3.12)
|
||||
- vLLM-Omni: PR #4771 branch or current `main` once merged
|
||||
- Extra dependency for the live server: `pip install sphn` (Opus framing;
|
||||
`aiohttp` already ships with vLLM)
|
||||
- `export HF_TOKEN=...` with access to the gated
|
||||
`nvidia/personaplex-7b-v1` repo (accept the license on the model page)
|
||||
|
||||
First run auto-downloads from the model repo: `model.safetensors`,
|
||||
`tokenizer_spm_32k_3.model`, `voices.tgz`, and the Mimi reference checkpoint;
|
||||
`kyutai/mimi` is fetched via `transformers` for the codec module graph.
|
||||
|
||||
#### Command
|
||||
|
||||
Live duplex server (official browser client served at `/`):
|
||||
|
||||
```bash
|
||||
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 \
|
||||
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
|
||||
--port 8124 --voice NATF2.pt
|
||||
```
|
||||
|
||||
Multi-session on the same GPU (elastic batching, callers join/leave any tick):
|
||||
|
||||
```bash
|
||||
HF_TOKEN=... CUDA_VISIBLE_DEVICES=0 \
|
||||
python -m vllm_omni.experimental.fullduplex.personaplex.serving.server \
|
||||
--port 8124 --batch-size 4
|
||||
```
|
||||
|
||||
Then open `http://localhost:8124/` and allow the microphone (use headphones so
|
||||
the agent does not hear itself). Headless alternative:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/personaplex/duplex_client.py \
|
||||
--url ws://localhost:8124/api/chat --input user.wav --out reply.wav
|
||||
```
|
||||
|
||||
Offline WAV-in / WAV-out (no server):
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/personaplex/personaplex_offline.py \
|
||||
--input-wav question.wav --output-wav reply.wav --output-text reply.json \
|
||||
--voice-prompt NATF2.pt --persona "You are a wise and friendly teacher."
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
# GPU-free contract tests (stubbed FrameStepper)
|
||||
pytest tests/e2e/features/fullduplex/test_personaplex_adapter.py \
|
||||
tests/e2e/features/fullduplex/test_personaplex_session.py \
|
||||
tests/e2e/features/fullduplex/test_personaplex_batched.py -q
|
||||
|
||||
# GPU e2e: batched isolation + slot recycle (coherence check needs a real
|
||||
# question WAV via PPLEX_QUESTION_WAV)
|
||||
pytest tests/e2e/features/fullduplex/test_personaplex_elastic.py -q
|
||||
```
|
||||
|
||||
A green offline run producing a coherent spoken reply (and inner-monologue text
|
||||
in `reply.json`) is the quickest end-to-end check.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Realtime budget is 80 ms/frame. Verified eager per-tick latency at
|
||||
`--batch-size 4` is ~70-74 ms on this hardware class, i.e. all four
|
||||
conversations stay realtime; single-session has comfortable headroom.
|
||||
- Decoding is greedy only (temperature/top-k knobs are intentionally not
|
||||
exposed): greedy is what the parity gates pin against the reference
|
||||
implementation, and sampling amplifies sub-bit numeric drift into divergence.
|
||||
- With `--batch-size N`, connections beyond N are rejected until a slot frees;
|
||||
a recycled slot is bit-exact with a fresh engine for the same inputs.
|
||||
- Voice prompts: bundled `voices.tgz` names (`NATF*`/`NATM*` natural,
|
||||
`VARF*`/`VARM*` varied) or a path to your own `.pt`/`.wav`.
|
||||
- Run the browser client near the server. The 80 ms cadence is sensitive to
|
||||
network jitter; over a high-latency link playback can stutter regardless of
|
||||
engine speed. On localhost it is smooth.
|
||||
- Sliding window is 3000 frames (~4 min). The ring KV recycles beyond that;
|
||||
very long sessions keep running with a rolling context.
|
||||
@@ -46,6 +46,7 @@ recipes/
|
||||
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
|
||||
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
|
||||
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
|
||||
| [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
|
||||
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
|
||||
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
|
||||
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
|
||||
|
||||
@@ -164,10 +164,11 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
PR [#2495](https://github.com/vllm-project/vllm-omni/pull/2495) adds
|
||||
performance CI configs for the same DiT-only settings. The CI step is
|
||||
currently opt-in (gated by `RUN_HUNYUAN_IMAGE3_PERF=1`) with `soft_fail`
|
||||
enabled, intended for initial data collection. Performance assertions are
|
||||
skipped (`skip-performance-assertion: true`); the baseline values in the
|
||||
JSON configs are reference-only and will be promoted to regression gates
|
||||
once enough nightly data has been collected.
|
||||
enabled, intended for initial data collection. Perf runners no longer
|
||||
compare metrics against JSON baselines at assert time; the `baseline`
|
||||
blocks in the configs are reference-only and are copied into result JSON
|
||||
for reporting. Regression gates can be reintroduced once enough nightly
|
||||
data has been collected.
|
||||
|
||||
The user-facing equivalent is to launch one of the CLI commands above and
|
||||
generate 1024x1024 images with 50 denoising steps.
|
||||
|
||||
@@ -423,6 +423,156 @@ ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_
|
||||
once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`;
|
||||
unknown keys for the model are dropped.
|
||||
|
||||
## ROCm
|
||||
|
||||
### 1x MI350X (Online serving)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04+
|
||||
- Python: 3.12+
|
||||
- Driver / runtime: ROCm, HIP 7.2, **gfx942 or gfx950** (validated on gfx950 / MI350X,
|
||||
which exposes 252 GiB of HBM to PyTorch)
|
||||
- PyTorch: 2.11.0 (ROCm build)
|
||||
- vLLM version: match the repository requirements from your current checkout
|
||||
- vLLM-Omni version or commit: use the commit you are deploying from
|
||||
|
||||
No ROCm-specific patch or flag is needed: Cosmos3 runs on the stock code path.
|
||||
|
||||
#### Command
|
||||
|
||||
The NVIDIA command works unchanged. Guardrails need the gated
|
||||
`nvidia/Cosmos-1.0-Guardrail` repo, so the quick path is to disable them (you are
|
||||
then responsible for license compliance):
|
||||
|
||||
```bash
|
||||
vllm serve nvidia/Cosmos3-Nano \
|
||||
--omni \
|
||||
--no-guardrails \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--init-timeout 1800
|
||||
```
|
||||
|
||||
For 720p video, add `--vae-use-tiling`: it cuts peak VRAM by ~68% for ~13%
|
||||
latency, which is the difference between needing a 120 GiB card and a 40 GiB one.
|
||||
`--quantization fp8` and `--enable-layerwise-offload` are also supported; see the
|
||||
Notes for the measured cost of each. For extra GPUs use `--ulysses-degree N` or
|
||||
`--tensor-parallel-size N` (not validated on ROCm yet).
|
||||
|
||||
#### Verification
|
||||
|
||||
Use the same `curl` requests as the CUDA section above. The ROCm images do not
|
||||
ship `ffprobe`, so the CUDA section's `ffprobe` line does not work here. `imageio`
|
||||
and `pyav` are already installed and report the same thing:
|
||||
|
||||
```bash
|
||||
python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print(im.size, im.mode)"
|
||||
python -c "import imageio.v3 as iio; print(iio.improps('cosmos3_t2v.mp4', plugin='pyav').shape)"
|
||||
```
|
||||
|
||||
Expect `(1024, 1024) RGB` and `(189, 720, 1280, 3)` for the requests above.
|
||||
Re-running the same request with the same seed reproduces a byte-identical mp4,
|
||||
so an md5 comparison across two runs also works as a smoke check.
|
||||
|
||||
#### Notes
|
||||
|
||||
- **Warm up once per output shape before timing anything.** The first generation
|
||||
of a shape pays a one-time cost that is cached on disk and reused by later
|
||||
processes in the same container: a ~250 s aiter attention-kernel JIT build plus
|
||||
a ~390 s build of the VAE decode path *for that shape*, so warming 1024² images
|
||||
does nothing for 189-frame video. At 720p / 189 frames / 35 steps this is
|
||||
**540 s cold vs 161 s warm, with byte-identical output**.
|
||||
- **Measured on 1x MI350X (bf16, guardrails off, warm):** T2I 1024² @ 50 steps
|
||||
**~2.7 s**; T2V 1280×720 / 189 frames @ 35 steps **~161 s**, of which ~92% is
|
||||
the DiT denoise loop and ~4% VAE decode, so optimization effort belongs in the
|
||||
denoise loop. The optional flags act on different terms of the memory bill and
|
||||
are therefore complementary rather than redundant:
|
||||
|
||||
| Flag | Latency | Peak reserved | Peak allocated | Acts on |
|
||||
|---|---|---|---|---|
|
||||
| *(none)* | 161 s | 120 GiB | 95 GiB | — |
|
||||
| `--vae-use-tiling` | 183 s (+13.5%) | **38 GiB** (−68%) | **36 GiB** (−62%) | decode activations |
|
||||
| `--enable-layerwise-offload` | 162 s (+0.8%) | 84 GiB (−30%) | 69 GiB (−27%) | weights |
|
||||
| `--quantization fp8` | 150 s (−6.6%) | 107 GiB (−11%) | 82 GiB (−14%) | weights (online, no calibration) |
|
||||
|
||||
Reserved is the caching-allocator high water mark and runs ~2–25 GiB above
|
||||
allocated; allocated is the figure to compare against another platform.
|
||||
- **`--vae-use-tiling` is by far the biggest lever for video.** Of the 95 GiB
|
||||
allocated by default only ~35 GiB is weights and fixed overhead; the remaining
|
||||
**60 GiB is activations from decoding all 189 frames in one piece**, which is
|
||||
why fp8 and offload cannot touch it and why tiling, which leaves under 5 GiB of
|
||||
them, wins so much. `--vae-use-slicing` is a no-op for a single video request:
|
||||
peak memory is bit-identical to the default.
|
||||
- **`--enable-layerwise-offload` is close to free on this part:** per-block DiT
|
||||
compute is long enough to hide the weight transfer, so on a card this size it
|
||||
buys headroom for concurrency or higher resolution rather than the
|
||||
speed-for-memory trade it is on small cards.
|
||||
- **Attention backend:** the numbers above are aiter FlashAttention, not SDPA.
|
||||
ROCm picks aiter when it is installed and the encoded device capability
|
||||
satisfies `90 < major * 10 + minor < 100` (gfx942 and gfx950); otherwise it
|
||||
falls back to `TORCH_SDPA` and says so only at debug level, which is easy to
|
||||
miss. If you are reporting numbers, assert the choice instead of trusting the
|
||||
default: `RocmOmniPlatform.get_diffusion_attn_backend_cls(None, 128)` returns
|
||||
the resolved path.
|
||||
- **Determinism is per-configuration:** same seed and same flags reproduce a
|
||||
byte-identical mp4 across processes, but different configurations do not match
|
||||
each other — including offload, which nominally only relocates weights, so do
|
||||
not use it as a bit-exact regression baseline. Quality was not evaluated for
|
||||
any configuration.
|
||||
|
||||
### 1x MI350X (Offline generation)
|
||||
|
||||
#### Environment
|
||||
|
||||
Same as the online serving section above.
|
||||
|
||||
#### Command
|
||||
|
||||
The standard task examples run unchanged:
|
||||
|
||||
```bash
|
||||
# Text-to-image
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--prompt "A photorealistic red sports car at golden hour, cinematic lighting." \
|
||||
--negative-prompt "blurry, distorted, low quality" \
|
||||
--height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \
|
||||
--extra-body '{"flow_shift": 3.0, "guardrails": false}' \
|
||||
--output cosmos3_t2i.png
|
||||
|
||||
# Text-to-video
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--prompt "A robot arm is cleaning a plate in the kitchen." \
|
||||
--negative-prompt "blurry, distorted, low quality" \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 --flow-shift 10.0 \
|
||||
--extra-body '{"max_sequence_length": 4096, "guardrails": false,
|
||||
"use_resolution_template": false, "use_duration_template": false}' \
|
||||
--output cosmos3_t2v.mp4
|
||||
```
|
||||
|
||||
Both scripts print `Total generation time` and `Worker peak GPU memory
|
||||
(reserved)`, which is where the latency and reserved figures in the Notes above
|
||||
come from. Add `--vae-use-tiling`, `--quantization fp8` or
|
||||
`--enable-layerwise-offload` to reproduce the other rows. The allocated column
|
||||
needs `VLLM_LOGGING_LEVEL=DEBUG`, which enables a per-request line reporting
|
||||
reserved, allocated and pool overhead together.
|
||||
|
||||
#### Verification
|
||||
|
||||
Same as the online serving section above.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Run each command twice and read the second run; the first pays the one-time
|
||||
compilation described above. If you run in a container, keep the container
|
||||
alive between runs (`docker run --rm` throws the JIT cache away and you pay
|
||||
the ~250 s kernel build every time).
|
||||
- Latency and memory figures are identical to the online serving section — they
|
||||
were taken from this path, because a single process has no client/server
|
||||
timing ambiguity.
|
||||
|
||||
## NPU
|
||||
|
||||
### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving
|
||||
|
||||
@@ -81,6 +81,21 @@ def _from_pipeline_key(
|
||||
)
|
||||
|
||||
|
||||
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
|
||||
omni_config = _from_pipeline_key("personaplex")
|
||||
|
||||
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
|
||||
|
||||
|
||||
def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
|
||||
deploy_path = tmp_path / "personaplex-turn.yaml"
|
||||
deploy_path.write_text("session_mode: turn\n", encoding="utf-8")
|
||||
|
||||
omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
|
||||
|
||||
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
|
||||
def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str):
|
||||
pipeline = _resolve_pipeline_or_skip(model_type)
|
||||
@@ -105,6 +120,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
|
||||
assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1)
|
||||
|
||||
engine_args = legacy_stage.yaml_engine_args
|
||||
assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
|
||||
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
|
||||
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
|
||||
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
|
||||
@@ -379,6 +395,7 @@ def test_runtime_config_fields_match_rfc_runtime_scope():
|
||||
def test_sub_config_fields_match_rfc_scopes():
|
||||
assert {f.name for f in fields(OmniStageModelConfig)} == {
|
||||
"active_stream_window",
|
||||
"duplex_max_sessions",
|
||||
"enable_sleep_mode",
|
||||
"default_sampling_params",
|
||||
"subtalker_sampling_params",
|
||||
|
||||
@@ -3,7 +3,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from types import SimpleNamespace
|
||||
from collections.abc import Callable
|
||||
from types import MethodType, SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
@@ -12,9 +13,22 @@ from vllm.v1.outputs import LogprobsLists
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
# Shared update_from_output helpers moved onto OmniSchedulerMixin; bind them on
|
||||
# the lightweight stub so direct method calls still exercise the real paths.
|
||||
_MIXIN_UPDATE_HELPERS = (
|
||||
"_remove_stopped_requests_from_queues",
|
||||
"_handle_failed_kv_load_outputs",
|
||||
"_aggregate_kv_connector_stats",
|
||||
"_publish_kv_cache_events",
|
||||
"_attach_finished_request_sets",
|
||||
"_attach_scheduler_stats",
|
||||
"_capture_omni_connector_output",
|
||||
)
|
||||
|
||||
|
||||
def _logprobs(
|
||||
token_ids: list[list[int]],
|
||||
@@ -105,18 +119,14 @@ class _RequestQueue(list):
|
||||
self.remove(request)
|
||||
|
||||
|
||||
def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
|
||||
"""Spec decode can sample tokens past a stop; the trailing tokens are
|
||||
trimmed by _update_request_with_output, and the emitted logprob rows must
|
||||
be trimmed with them (upstream slices logprobs after the trim)."""
|
||||
request = _Request("req")
|
||||
request.num_computed_tokens = 8
|
||||
def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
|
||||
"""Build a minimal stub for OmniARScheduler.update_from_output()."""
|
||||
scheduler = SimpleNamespace(
|
||||
perf_metrics=None,
|
||||
connector=None,
|
||||
chunk_transfer_adapter=None,
|
||||
requests={request.request_id: request},
|
||||
running=[request],
|
||||
requests={request.request_id: request for request in requests},
|
||||
running=list(requests),
|
||||
waiting=_RequestQueue(),
|
||||
skipped_waiting=_RequestQueue(),
|
||||
structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
|
||||
@@ -129,7 +139,42 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
|
||||
finished_req_ids_dict=defaultdict(set),
|
||||
kv_cache_manager=SimpleNamespace(take_events=lambda: None),
|
||||
kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
|
||||
recompute_kv_load_failures=False,
|
||||
)
|
||||
for name in _MIXIN_UPDATE_HELPERS:
|
||||
setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler))
|
||||
scheduler._cleanup_kv_tracking = MethodType(OmniARScheduler._cleanup_kv_tracking, scheduler)
|
||||
scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
|
||||
scheduler.make_stats = lambda *args, **kwargs: None
|
||||
return scheduler
|
||||
|
||||
|
||||
def _bind_request_lifecycle(
|
||||
scheduler: SimpleNamespace,
|
||||
*,
|
||||
update_request: Callable,
|
||||
handle_stopped: Callable | None = None,
|
||||
) -> None:
|
||||
def free_request(request):
|
||||
scheduler.requests.pop(request.request_id, None)
|
||||
scheduler.finished_req_ids.add(request.request_id)
|
||||
scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
|
||||
# vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
|
||||
return None, None
|
||||
|
||||
scheduler._update_request_with_output = update_request
|
||||
scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
|
||||
scheduler._handle_stopped_request = handle_stopped or (lambda _request: True)
|
||||
scheduler._free_request = free_request
|
||||
|
||||
|
||||
def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
|
||||
"""Spec decode can sample tokens past a stop; the trailing tokens are
|
||||
trimmed by _update_request_with_output, and the emitted logprob rows must
|
||||
be trimmed with them (upstream slices logprobs after the trim)."""
|
||||
request = _Request("req")
|
||||
request.num_computed_tokens = 8
|
||||
scheduler = _make_scheduler_stub([request])
|
||||
|
||||
def update_request_trimming(req, token_ids):
|
||||
# Mirror vllm Scheduler._update_request_with_output: the first sampled
|
||||
@@ -139,19 +184,7 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
|
||||
req.status = RequestStatus.FINISHED_STOPPED
|
||||
return token_ids, True
|
||||
|
||||
def free_request(req):
|
||||
scheduler.requests.pop(req.request_id, None)
|
||||
scheduler.finished_req_ids.add(req.request_id)
|
||||
scheduler.finished_req_ids_dict[req.client_index].add(req.request_id)
|
||||
return None, None
|
||||
|
||||
scheduler._update_request_with_output = update_request_trimming
|
||||
scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
|
||||
scheduler._handle_stopped_request = lambda _request: True
|
||||
scheduler._free_request = free_request
|
||||
scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
|
||||
scheduler.make_stats = lambda *args, **kwargs: None
|
||||
scheduler._capture_omni_connector_output = lambda _output: None
|
||||
_bind_request_lifecycle(scheduler, update_request=update_request_trimming)
|
||||
|
||||
scheduler_output = SimpleNamespace(
|
||||
num_scheduled_tokens={"req": 3},
|
||||
@@ -190,45 +223,14 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
|
||||
"""A bad runner response must not bring down unrelated batch requests."""
|
||||
bad = _Request("bad")
|
||||
good = _Request("good")
|
||||
scheduler = SimpleNamespace(
|
||||
perf_metrics=None,
|
||||
connector=None,
|
||||
chunk_transfer_adapter=None,
|
||||
requests={bad.request_id: bad, good.request_id: good},
|
||||
running=[bad, good],
|
||||
waiting=_RequestQueue(),
|
||||
skipped_waiting=_RequestQueue(),
|
||||
structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
|
||||
transfer_triggered_requests=set(),
|
||||
active_kv_transfers=set(),
|
||||
pending_stop_after_extraction=set(),
|
||||
waiting_for_transfer_free=set(),
|
||||
_new_prompt_len_snapshot={},
|
||||
finished_req_ids=set(),
|
||||
finished_req_ids_dict=defaultdict(set),
|
||||
kv_cache_manager=SimpleNamespace(take_events=lambda: None),
|
||||
kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
|
||||
)
|
||||
scheduler = _make_scheduler_stub([bad, good])
|
||||
update_calls: list[str] = []
|
||||
|
||||
def update_request(request, token_ids):
|
||||
update_calls.append(request.request_id)
|
||||
return token_ids, False
|
||||
|
||||
def free_request(request):
|
||||
scheduler.requests.pop(request.request_id, None)
|
||||
scheduler.finished_req_ids.add(request.request_id)
|
||||
scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
|
||||
# vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
|
||||
return None, None
|
||||
|
||||
scheduler._update_request_with_output = update_request
|
||||
scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
|
||||
scheduler._handle_stopped_request = lambda _request: True
|
||||
scheduler._free_request = free_request
|
||||
scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
|
||||
scheduler.make_stats = lambda *args, **kwargs: None
|
||||
scheduler._capture_omni_connector_output = lambda _output: None
|
||||
_bind_request_lifecycle(scheduler, update_request=update_request)
|
||||
|
||||
scheduler_output = SimpleNamespace(
|
||||
num_scheduled_tokens={"bad": 1, "good": 1},
|
||||
|
||||
@@ -110,6 +110,9 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
||||
session.status = RequestStatus.RUNNING
|
||||
session.resumable = True
|
||||
session.num_computed_tokens = len(session.prompt_token_ids)
|
||||
# schedule() has counted this step's token as in flight; update_from_output()
|
||||
# must settle it before handling the segment boundary.
|
||||
session.num_in_flight_tokens = 1
|
||||
|
||||
sched = MagicMock()
|
||||
sched.requests = {session.request_id: session}
|
||||
@@ -118,12 +121,20 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
||||
is_done_receiving_chunks=lambda _request_id: True,
|
||||
segment_finished_requests={session.request_id},
|
||||
)
|
||||
sched._handle_stopped_request.return_value = False
|
||||
|
||||
def rearm_resumable_request(request: Request) -> bool:
|
||||
request.status = RequestStatus.WAITING
|
||||
return False
|
||||
|
||||
sched._handle_stopped_request.side_effect = rearm_resumable_request
|
||||
sched.running = [session]
|
||||
sched.waiting = MagicMock()
|
||||
sched.skipped_waiting = MagicMock()
|
||||
sched.structured_output_manager.should_advance.return_value = False
|
||||
sched._pending_finish_reqs = []
|
||||
# Async scheduling can observe the segment boundary in the next schedule()
|
||||
# before this model output is applied and defer the same request for finish.
|
||||
# update_from_output() must not re-arm the resumable request twice.
|
||||
sched._pending_finish_reqs = [session]
|
||||
sched.recompute_kv_load_failures = False
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
@@ -155,6 +166,8 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
||||
output = outputs[session.client_index].outputs[0]
|
||||
assert output.finish_reason is not None
|
||||
assert output.is_segment_finished is True
|
||||
sched._handle_stopped_request.assert_called_once_with(session)
|
||||
assert sched._pending_finish_reqs == []
|
||||
|
||||
|
||||
def test_async_chunk_resumable_stop_rearms_connector_polling() -> None:
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
from collections import defaultdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from vllm.v1.engine import FinishReason
|
||||
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.core.sched.output import OmniChunkRecvHandle
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _Scheduler(OmniSchedulerMixin):
|
||||
pass
|
||||
|
||||
|
||||
def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
calls = []
|
||||
scheduler = _Scheduler()
|
||||
scheduler.waiting = ["waiting"]
|
||||
scheduler.running = ["running"]
|
||||
scheduler.requests = {"request": object()}
|
||||
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
|
||||
scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
|
||||
scheduler.chunk_transfer_adapter = SimpleNamespace(
|
||||
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("process", waiting, running, scheduler_requests)
|
||||
),
|
||||
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("restore-chunks", waiting, running, scheduler_requests)
|
||||
),
|
||||
)
|
||||
scheduler.input_coordinator = SimpleNamespace(
|
||||
restore_queues=lambda waiting: calls.append(("restore-full", waiting))
|
||||
)
|
||||
|
||||
scheduler._process_pending_omni_inputs("ar")
|
||||
scheduler._restore_omni_wait_queues()
|
||||
|
||||
assert calls == [
|
||||
("consume", "ar"),
|
||||
("timeouts",),
|
||||
("process", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
("restore-full", scheduler.waiting),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("synthesize_abort_outputs", "expected_finish_reason"),
|
||||
[(False, None), (True, FinishReason.ABORT)],
|
||||
)
|
||||
def test_finished_request_attachment_keeps_ar_abort_policy_explicit(
|
||||
synthesize_abort_outputs,
|
||||
expected_finish_reason,
|
||||
):
|
||||
scheduler = _Scheduler()
|
||||
scheduler.finished_req_ids_dict = defaultdict(set, {2: {"req-finished"}})
|
||||
outputs = {}
|
||||
|
||||
scheduler._attach_finished_request_sets(
|
||||
outputs,
|
||||
synthesize_abort_outputs=synthesize_abort_outputs,
|
||||
)
|
||||
|
||||
assert outputs[2].finished_requests == {"req-finished"}
|
||||
if expected_finish_reason is None:
|
||||
assert outputs[2].outputs == []
|
||||
else:
|
||||
assert outputs[2].outputs[0].finish_reason == expected_finish_reason
|
||||
assert scheduler.finished_req_ids_dict == {}
|
||||
|
||||
|
||||
def test_chunk_receive_handle_carries_minimal_registration_fields():
|
||||
handle = OmniChunkRecvHandle(request_id="req", external_req_id="external")
|
||||
assert handle.request_id == "req"
|
||||
assert handle.external_req_id == "external"
|
||||
|
||||
|
||||
def test_output_helper_preserves_required_nan_counter_default():
|
||||
scheduler = _Scheduler()
|
||||
request = SimpleNamespace(
|
||||
request_id="req-output",
|
||||
trace_headers=None,
|
||||
take_events=lambda: [],
|
||||
)
|
||||
|
||||
output = scheduler._make_omni_engine_output(request, new_token_ids=[])
|
||||
|
||||
assert output.num_nans_in_logits == 0
|
||||
@@ -1,3 +1,6 @@
|
||||
import math
|
||||
from typing import NamedTuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -317,6 +320,103 @@ def test_deferred_multimodal_cache_write_commits_on_completion():
|
||||
assert torch.equal(rows[8:16], expected)
|
||||
|
||||
|
||||
def test_discard_deferred_mm_outputs_drops_pending_chunks():
|
||||
"""discard_deferred_mm_outputs is the abort path: a request that
|
||||
leaves without a cache commit must not leak its staged chunks into
|
||||
the cache if commit_deferred_mm_outputs is later called for it."""
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
block_table = torch.tensor(
|
||||
[
|
||||
[2, 3, 4, 5],
|
||||
[6, 7, 8, 9],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
input_batch = MockInputBatch(
|
||||
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
|
||||
block_table=block_table,
|
||||
)
|
||||
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)},
|
||||
num_tokens_unpadded=4,
|
||||
slot_mapping=torch.arange(8, 12),
|
||||
skip_mm_cache_keys={mm_key},
|
||||
)
|
||||
|
||||
chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: chunk},
|
||||
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
assert "req1" in cache._deferred_mm_outputs
|
||||
|
||||
cache.discard_deferred_mm_outputs("req1")
|
||||
assert "req1" not in cache._deferred_mm_outputs
|
||||
|
||||
# A later commit call for the now-discarded req_id must be a no-op,
|
||||
# not a resurrection of the dropped chunk.
|
||||
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
|
||||
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
|
||||
assert torch.all(rows[8:12] == 0)
|
||||
|
||||
|
||||
def test_deferred_mm_outputs_request_id_reused_after_discard():
|
||||
"""A request ID reused after discard_deferred_mm_outputs must start
|
||||
clean: the new occupant's committed data must not be mixed with the
|
||||
previous (discarded) occupant's staged chunks."""
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
block_table = torch.tensor(
|
||||
[
|
||||
[2, 3, 4, 5],
|
||||
[6, 7, 8, 9],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
input_batch = MockInputBatch(
|
||||
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
|
||||
block_table=block_table,
|
||||
)
|
||||
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs={mm_key: torch.arange(8, dtype=torch.long).reshape(4, 2)},
|
||||
num_tokens_unpadded=4,
|
||||
slot_mapping=torch.arange(8, 12),
|
||||
skip_mm_cache_keys={mm_key},
|
||||
)
|
||||
|
||||
stale_chunk = torch.full((4, 2), -1, dtype=torch.long)
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: stale_chunk},
|
||||
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
cache.discard_deferred_mm_outputs("req1")
|
||||
|
||||
# "req1" is reused by a brand-new request with different data.
|
||||
fresh_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: fresh_chunk},
|
||||
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
|
||||
|
||||
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
|
||||
assert torch.equal(rows[8:12], fresh_chunk)
|
||||
|
||||
|
||||
def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
@@ -599,3 +699,219 @@ def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded, mocker):
|
||||
|
||||
assert req2_merged_mm_outputs.shape == torch.Size([num_new_toks_req2, curr_feat_dim])
|
||||
assert torch.all(req2_merged_mm_outputs == req2_new_mm_outputs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"passthrough_shape",
|
||||
[(5,), (5, 3)],
|
||||
ids=["1d_token_metadata", "2d_feature_tensor"],
|
||||
)
|
||||
def test_get_merged_multimodal_outputs_slices_passthrough_tensor_per_request(passthrough_shape, mocker):
|
||||
"""A passthrough key that is itself a real tensor (not the string/list
|
||||
passthrough already covered above) must be sliced per request through
|
||||
to_payload_element's tensor branch, under the same mixed cache
|
||||
hit/miss + uneven scheduled-length batch as test_get_merged_multimodal_outputs.
|
||||
|
||||
Covers both a 1D token-aligned metadata tensor (e.g. ref_code_len,
|
||||
codec_streaming -- shape (total_tokens,)) and a 2D per-token feature
|
||||
tensor (shape (total_tokens, feat_dim)), since W2's batching could
|
||||
easily special-case one and silently mishandle the other.
|
||||
"""
|
||||
feat_dims = {"foo": 100, "bar": 100}
|
||||
cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN)
|
||||
|
||||
orig_num_tokens_unpadded = 8
|
||||
slot_offset = 8
|
||||
orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
|
||||
feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
|
||||
orig_mm_outputs = {
|
||||
key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||
}
|
||||
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs=orig_mm_outputs,
|
||||
num_tokens_unpadded=orig_num_tokens_unpadded,
|
||||
slot_mapping=orig_slot_mapping,
|
||||
num_tokens_padded=None,
|
||||
)
|
||||
|
||||
num_new_toks_req1 = 3
|
||||
num_new_toks_req2 = 2
|
||||
cache.add_prefix_cached_new_req_id("req1")
|
||||
|
||||
num_scheduled_tokens = {
|
||||
"req1": num_new_toks_req1,
|
||||
"req2": num_new_toks_req2,
|
||||
}
|
||||
|
||||
new_mm_outputs = {}
|
||||
for mm_key in cache.mm_cache_keys:
|
||||
new_mm_outputs[mm_key] = torch.rand(
|
||||
(num_new_toks_req1 + num_new_toks_req2, feature_dims[mm_key]),
|
||||
dtype=DTYPE,
|
||||
)
|
||||
total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2
|
||||
assert passthrough_shape[0] == total_scheduled_tokens
|
||||
passthrough_tensor = torch.arange(int(torch.prod(torch.tensor(passthrough_shape))), dtype=DTYPE).reshape(
|
||||
passthrough_shape
|
||||
)
|
||||
new_mm_outputs["passthrough_tensor"] = passthrough_tensor
|
||||
|
||||
input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
|
||||
|
||||
mocker.patch(
|
||||
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
|
||||
new=fake_get_cached_block_ids,
|
||||
)
|
||||
merged_mm_outputs = cache.get_merged_multimodal_states(
|
||||
query_start_loc=[0, num_new_toks_req1],
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs=new_mm_outputs,
|
||||
num_scheduled_tokens=num_scheduled_tokens,
|
||||
)
|
||||
|
||||
assert "passthrough_tensor" not in cache.mm_cache_keys
|
||||
|
||||
passthrough_out = merged_mm_outputs["passthrough_tensor"]
|
||||
# req1 (cache hit) gets tokens [0:3), req2 (cache miss) gets tokens [3:5) --
|
||||
# a passthrough tensor is sliced by scheduled position, not merged with
|
||||
# any cached history (unlike the mm_cache_keys tensors above).
|
||||
assert torch.equal(passthrough_out["req1"], passthrough_tensor[0:num_new_toks_req1])
|
||||
assert torch.equal(passthrough_out["req2"], passthrough_tensor[num_new_toks_req1:total_scheduled_tokens])
|
||||
|
||||
|
||||
class MergeHarness(NamedTuple):
|
||||
"""A primed cache plus the arguments get_merged_multimodal_states needs."""
|
||||
|
||||
cache: OmniTensorPrefixCache
|
||||
input_batch: MockInputBatch
|
||||
num_scheduled_tokens: dict[str, int]
|
||||
query_start_loc: list[int]
|
||||
total_scheduled_tokens: int
|
||||
new_mm_outputs: dict[str, object]
|
||||
|
||||
|
||||
def build_mixed_hit_miss_harness(mocker, num_new_toks_req1: int = 3, num_new_toks_req2: int = 2) -> MergeHarness:
|
||||
"""Two requests over one cache key: req1 hits 8 cached tokens, req2 misses,
|
||||
and their scheduled lengths differ -- the same batch shape as
|
||||
test_get_merged_multimodal_outputs, which is what makes a per-request
|
||||
boundary error observable. Callers add their passthrough keys to
|
||||
``new_mm_outputs`` before calling get_merged_multimodal_states.
|
||||
"""
|
||||
cache = get_omni_pcache_with_mm_tensors({"foo": 100}, seq_len=DEFAULT_SEQ_LEN)
|
||||
|
||||
orig_num_tokens_unpadded = 8
|
||||
feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()}
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs={
|
||||
key: torch.rand((orig_num_tokens_unpadded, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||
},
|
||||
num_tokens_unpadded=orig_num_tokens_unpadded,
|
||||
slot_mapping=torch.arange(8, 8 + orig_num_tokens_unpadded),
|
||||
num_tokens_padded=None,
|
||||
)
|
||||
cache.add_prefix_cached_new_req_id("req1")
|
||||
|
||||
total_scheduled_tokens = num_new_toks_req1 + num_new_toks_req2
|
||||
mocker.patch(
|
||||
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
|
||||
new=fake_get_cached_block_ids,
|
||||
)
|
||||
return MergeHarness(
|
||||
cache=cache,
|
||||
input_batch=MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])),
|
||||
num_scheduled_tokens={"req1": num_new_toks_req1, "req2": num_new_toks_req2},
|
||||
query_start_loc=[0, num_new_toks_req1],
|
||||
total_scheduled_tokens=total_scheduled_tokens,
|
||||
new_mm_outputs={
|
||||
key: torch.rand((total_scheduled_tokens, feature_dims[key]), dtype=DTYPE) for key in cache.mm_cache_keys
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_get_merged_multimodal_outputs_passes_per_request_tensor_list_through(mocker):
|
||||
"""The passthrough shape qwen3-tts actually emits: a nested ``codes``/``meta``
|
||||
dict whose ``codes.ref`` holds one tensor per request
|
||||
(qwen3_tts_talker.make_omni_output), carried through to_payload_element's
|
||||
pass_lists_through branch beside 1D int metadata that IS sliced per request.
|
||||
|
||||
Each request must receive the list *whole*. Per-request indexing happens
|
||||
later, in the runner's _unwrap_lists; a W2 that split the list here would
|
||||
make that second index select a different request's reference codes. The
|
||||
existing list coverage passes a list of strings, which cannot catch a
|
||||
tensor-specific regression, and no other test covers an int dtype.
|
||||
"""
|
||||
harness = build_mixed_hit_miss_harness(mocker)
|
||||
|
||||
ref_codes = [torch.tensor([11, 12], dtype=torch.long), torch.tensor([21], dtype=torch.long)]
|
||||
ref_code_len = torch.arange(harness.total_scheduled_tokens, dtype=torch.int32)
|
||||
harness.new_mm_outputs["codes"] = {"ref": ref_codes}
|
||||
harness.new_mm_outputs["meta"] = {"ref_code_len": ref_code_len}
|
||||
|
||||
merged = harness.cache.get_merged_multimodal_states(
|
||||
query_start_loc=harness.query_start_loc,
|
||||
input_batch=harness.input_batch,
|
||||
multimodal_outputs=harness.new_mm_outputs,
|
||||
num_scheduled_tokens=harness.num_scheduled_tokens,
|
||||
)
|
||||
|
||||
assert "codes" not in harness.cache.mm_cache_keys
|
||||
for req_id in ("req1", "req2"):
|
||||
req_refs = merged["codes"][req_id]["ref"]
|
||||
assert len(req_refs) == len(ref_codes)
|
||||
for got, want in zip(req_refs, ref_codes):
|
||||
assert torch.equal(got, want)
|
||||
assert got.dtype == torch.long
|
||||
|
||||
# Per-request clones: the branch exists so one request's downstream mutation
|
||||
# cannot reach another request's payload or the source list.
|
||||
merged["codes"]["req1"]["ref"][0] += 100
|
||||
assert torch.equal(merged["codes"]["req2"]["ref"][0], torch.tensor([11, 12], dtype=torch.long))
|
||||
assert torch.equal(ref_codes[0], torch.tensor([11, 12], dtype=torch.long))
|
||||
|
||||
# Token-aligned 1D metadata is still split per request, unlike the list.
|
||||
req1_len = harness.num_scheduled_tokens["req1"]
|
||||
assert torch.equal(merged["meta"]["req1"]["ref_code_len"], ref_code_len[:req1_len])
|
||||
assert torch.equal(merged["meta"]["req2"]["ref_code_len"], ref_code_len[req1_len:])
|
||||
assert merged["meta"]["req1"]["ref_code_len"].dtype == torch.int32
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"passthrough_shape",
|
||||
[(3, 4), (1, 4), (7,)],
|
||||
ids=["shorter_2d", "single_row_2d", "longer_1d"],
|
||||
)
|
||||
def test_get_merged_multimodal_outputs_broadcasts_non_token_aligned_passthrough(passthrough_shape, mocker):
|
||||
"""A passthrough tensor whose first dimension is not the scheduled-token
|
||||
count is request-invariant: to_payload_element hands every request the whole
|
||||
tensor, cloned, rather than slicing it.
|
||||
|
||||
W2 batching most plausibly breaks by assuming every passthrough tensor is
|
||||
token-aligned; that assumption would slice these and give each request a
|
||||
truncated view, or an empty one once start exceeds the first dimension.
|
||||
"""
|
||||
harness = build_mixed_hit_miss_harness(mocker)
|
||||
assert passthrough_shape[0] != harness.total_scheduled_tokens
|
||||
|
||||
passthrough = torch.arange(math.prod(passthrough_shape), dtype=DTYPE).reshape(passthrough_shape)
|
||||
expected = passthrough.clone()
|
||||
harness.new_mm_outputs["speaker_embedding"] = passthrough
|
||||
|
||||
merged = harness.cache.get_merged_multimodal_states(
|
||||
query_start_loc=harness.query_start_loc,
|
||||
input_batch=harness.input_batch,
|
||||
multimodal_outputs=harness.new_mm_outputs,
|
||||
num_scheduled_tokens=harness.num_scheduled_tokens,
|
||||
)
|
||||
|
||||
assert "speaker_embedding" not in harness.cache.mm_cache_keys
|
||||
out = merged["speaker_embedding"]
|
||||
for req_id in ("req1", "req2"):
|
||||
assert out[req_id].shape == passthrough_shape
|
||||
assert torch.equal(out[req_id], expected)
|
||||
|
||||
out["req1"] += 100
|
||||
assert torch.equal(out["req2"], expected)
|
||||
assert torch.equal(passthrough, expected)
|
||||
@@ -0,0 +1,119 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
|
||||
|
||||
# `cpu` is the lane selector, not a hardware claim: the only Buildkite steps that
|
||||
# reach tests/ select `core_model and cpu`, and they run on GPU hardware (l4_1).
|
||||
# Marking this `cuda` collects it nowhere. Real hardware need is the skipif below.
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.cpu,
|
||||
pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device"),
|
||||
]
|
||||
|
||||
NUM_BLOCKS = 10
|
||||
BLOCK_SIZE = 4
|
||||
HIDDEN_SIZE = 2
|
||||
DTYPE = torch.float32
|
||||
|
||||
|
||||
def get_omni_pcache() -> OmniTensorPrefixCache:
|
||||
"""Build an OmniTensorPrefixCache, but don't init mm tensors."""
|
||||
return OmniTensorPrefixCache(
|
||||
num_blocks=NUM_BLOCKS,
|
||||
block_size=BLOCK_SIZE,
|
||||
hidden_size=HIDDEN_SIZE,
|
||||
hs_dtype=DTYPE,
|
||||
)
|
||||
|
||||
|
||||
def test_drain_ready_async_writes_is_a_noop_before_the_event_is_ready(monkeypatch):
|
||||
"""drain_ready_async_writes() must not touch the cache while the
|
||||
scheduled D2H's event has not completed, and must consume it exactly
|
||||
once after it has.
|
||||
|
||||
Real GPU timing (e.g. via torch.cuda._sleep to stall the copy stream)
|
||||
was tried and found unreliable on this box: cross-stream wait_stream/
|
||||
wait_event did not reliably delay event.query() from reporting ready
|
||||
immediately, likely a virtualization/driver quirk of this cloud GPU.
|
||||
Forcing the event's query() result directly tests the branching logic
|
||||
in drain_ready_async_writes/_consume_pending_write deterministically,
|
||||
independent of real CUDA scheduling behavior.
|
||||
"""
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4)
|
||||
|
||||
num_tokens = 4
|
||||
slot_mapping_gpu = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda")
|
||||
mm_gpu = {
|
||||
mm_key: torch.arange(num_tokens * HIDDEN_SIZE, dtype=torch.float32, device="cuda").reshape(
|
||||
num_tokens, HIDDEN_SIZE
|
||||
)
|
||||
}
|
||||
|
||||
cache.schedule_async_write(
|
||||
hidden_states_gpu=None,
|
||||
multimodal_outputs_gpu=mm_gpu,
|
||||
slot_mapping_gpu=slot_mapping_gpu,
|
||||
num_tokens_unpadded=num_tokens,
|
||||
num_tokens_padded=num_tokens,
|
||||
)
|
||||
|
||||
pending_event = cache._pending_write.event
|
||||
rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE)
|
||||
|
||||
monkeypatch.setattr(pending_event, "query", lambda: False)
|
||||
assert cache.drain_ready_async_writes() == 0
|
||||
assert torch.all(rows[8 : 8 + num_tokens] == 0)
|
||||
assert cache._pending_write is not None
|
||||
|
||||
monkeypatch.setattr(pending_event, "query", lambda: True)
|
||||
assert cache.drain_ready_async_writes() == 1
|
||||
assert torch.equal(rows[8 : 8 + num_tokens], mm_gpu[mm_key].cpu())
|
||||
assert cache._pending_write is None
|
||||
|
||||
# Nothing pending anymore: a second drain is a no-op, not a re-scatter.
|
||||
assert cache.drain_ready_async_writes() == 0
|
||||
|
||||
|
||||
def test_schedule_async_write_consumes_previous_pending_write_first():
|
||||
"""Scheduling step N+1's write must first consume (scatter) step N's
|
||||
pending write, even if the caller never explicitly called
|
||||
drain_ready_async_writes() in between -- otherwise step N's data would
|
||||
be silently dropped when overwritten by step N+1's pending write."""
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
cache.maybe_init_missing_mm_cache_keys({mm_key: torch.zeros(4, HIDDEN_SIZE)}, seq_len=4)
|
||||
|
||||
num_tokens = 4
|
||||
first_slots = torch.arange(8, 8 + num_tokens, dtype=torch.int64, device="cuda")
|
||||
first_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 1.0, device="cuda")}
|
||||
second_slots = torch.arange(12, 12 + num_tokens, dtype=torch.int64, device="cuda")
|
||||
second_gpu = {mm_key: torch.full((num_tokens, HIDDEN_SIZE), 2.0, device="cuda")}
|
||||
|
||||
cache.schedule_async_write(
|
||||
hidden_states_gpu=None,
|
||||
multimodal_outputs_gpu=first_gpu,
|
||||
slot_mapping_gpu=first_slots,
|
||||
num_tokens_unpadded=num_tokens,
|
||||
num_tokens_padded=num_tokens,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# This call's internal _consume_pending_write() must flush the FIRST
|
||||
# write before the second one is scheduled; the caller never drained.
|
||||
cache.schedule_async_write(
|
||||
hidden_states_gpu=None,
|
||||
multimodal_outputs_gpu=second_gpu,
|
||||
slot_mapping_gpu=second_slots,
|
||||
num_tokens_unpadded=num_tokens,
|
||||
num_tokens_padded=num_tokens,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
cache.drain_ready_async_writes()
|
||||
|
||||
rows = cache.mm_outputs_cache[mm_key].view(-1, HIDDEN_SIZE)
|
||||
assert torch.all(rows[8:12] == 1.0)
|
||||
assert torch.all(rows[12:16] == 2.0)
|
||||
+8
-146
@@ -1,16 +1,15 @@
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import modify_stage_config
|
||||
@@ -119,32 +118,6 @@ def create_unique_server_pytest_params(
|
||||
]
|
||||
|
||||
|
||||
def create_benchmark_pytest_params(
|
||||
benchmark_configs: list[dict[str, Any]],
|
||||
server_to_benchmark_mapping: dict[str, dict],
|
||||
) -> list[Any]:
|
||||
"""Like :func:`create_benchmark_indices`, but wrap each index in ``pytest.param`` with JSON marks."""
|
||||
marks_by_name = _marks_by_test_name(benchmark_configs)
|
||||
params: list[Any] = []
|
||||
seen: set[str] = set()
|
||||
for config in benchmark_configs:
|
||||
test_name = config["test_name"]
|
||||
if test_name in seen:
|
||||
continue
|
||||
seen.add(test_name)
|
||||
params_list = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
|
||||
id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
|
||||
for idx, id_suffix in enumerate(id_suffixes):
|
||||
params.append(
|
||||
pytest.param(
|
||||
(test_name, idx),
|
||||
marks=marks_by_name.get(test_name, []),
|
||||
id=f"{test_name}-{id_suffix}",
|
||||
)
|
||||
)
|
||||
return params
|
||||
|
||||
|
||||
def create_paired_benchmark_pytest_params(
|
||||
server_entries: list[tuple[Any, str]],
|
||||
params_by_test_name: dict[str, list[Any]],
|
||||
@@ -357,10 +330,6 @@ def supports_video_generation(model_name: str) -> bool:
|
||||
return any(key in lower for key in ("wan", "video", "i2v", "t2v"))
|
||||
|
||||
|
||||
def supports_chat_generation(model_name: str) -> bool:
|
||||
return not supports_video_generation(model_name)
|
||||
|
||||
|
||||
def parse_stage_devices(stage_config_path: str) -> str:
|
||||
text = Path(stage_config_path).read_text(encoding="utf-8")
|
||||
raw_devices: list[str] = re.findall(r"^\s*devices:\s*\"?([0-9,\s]+)\"?\s*$", text, flags=re.MULTILINE)
|
||||
@@ -389,48 +358,6 @@ def assert_fault_exception(exc: Exception, error_keywords: tuple[str, ...]) -> N
|
||||
assert any(key in text for key in error_keywords), f"unexpected error under fault injection: {exc}"
|
||||
|
||||
|
||||
def wait_chat_request_ready(host: str, port: int, model: str, timeout_sec: int = 180) -> None:
|
||||
"""Poll a minimal chat request until success."""
|
||||
deadline = time.time() + timeout_sec
|
||||
payload = json.dumps(
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "Say hello in one short sentence."}],
|
||||
"stream": False,
|
||||
"modalities": ["text"],
|
||||
}
|
||||
)
|
||||
last_error: str | None = None
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
status, body = post_chat_completions_raw(host, port, payload)
|
||||
if status == 200:
|
||||
return
|
||||
last_error = f"http={status} body={body[:200]!r}"
|
||||
except Exception as exc: # noqa: BLE001
|
||||
last_error = str(exc)
|
||||
time.sleep(2)
|
||||
raise TimeoutError(f"runtime-teardown warmup request did not succeed within {timeout_sec}s: {last_error}")
|
||||
|
||||
|
||||
def assert_no_extra_worker_processes(
|
||||
baseline_pids: set[int],
|
||||
worker_pattern: str,
|
||||
timeout_sec: int = 60,
|
||||
) -> None:
|
||||
"""Ensure no extra worker PIDs remain after teardown."""
|
||||
deadline = time.time() + timeout_sec
|
||||
while time.time() < deadline:
|
||||
current = set(list_remote_process_pids_by_pattern(worker_pattern))
|
||||
extra = current - baseline_pids
|
||||
if not extra:
|
||||
return
|
||||
time.sleep(2)
|
||||
current = set(list_remote_process_pids_by_pattern(worker_pattern))
|
||||
extra = sorted(current - baseline_pids)
|
||||
assert not extra, f"orphan worker processes remain after container teardown: {extra}"
|
||||
|
||||
|
||||
def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]:
|
||||
mapping = {}
|
||||
for config in configs:
|
||||
@@ -637,57 +564,6 @@ def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str:
|
||||
return get_runtime_resource_label()
|
||||
|
||||
|
||||
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."""
|
||||
if baseline_raw is None:
|
||||
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)):
|
||||
if sweep_index is None:
|
||||
raise ValueError("list baseline requires sweep_index")
|
||||
if not (0 <= sweep_index < len(baseline_raw)):
|
||||
raise IndexError(f"baseline list len={len(baseline_raw)} has no index {sweep_index}")
|
||||
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 baseline config to one threshold per metric for this iteration."""
|
||||
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 run_benchmark(
|
||||
args: list[str],
|
||||
test_name: str,
|
||||
@@ -696,18 +572,17 @@ def run_benchmark(
|
||||
num_prompt: int,
|
||||
*,
|
||||
baseline_config: dict[str, Any] | None = None,
|
||||
sweep_index: int | None = None,
|
||||
request_rate: Any | None = None,
|
||||
max_concurrency: Any | None = None,
|
||||
random_input_len: Any | None = None,
|
||||
random_output_len: Any | None = None,
|
||||
resource_label: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""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.
|
||||
After ``vllm bench`` writes the JSON, ``result["baseline"]`` stores the raw
|
||||
``baseline`` object from the test config (when ``baseline_config`` is provided),
|
||||
including any hardware-nested map (e.g. ``{"H100": {...}}``). Downstream report
|
||||
tooling can select a hardware bucket when needed. 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)
|
||||
@@ -774,12 +649,8 @@ def run_benchmark(
|
||||
result = json.load(f)
|
||||
|
||||
if baseline_config:
|
||||
result["baseline"] = _baseline_thresholds_for_step(
|
||||
baseline_config,
|
||||
sweep_index=sweep_index,
|
||||
request_rate=request_rate,
|
||||
max_concurrency=max_concurrency,
|
||||
)
|
||||
# Persist the JSON baseline as-is (no hardware / sweep resolution).
|
||||
result["baseline"] = copy.deepcopy(baseline_config)
|
||||
else:
|
||||
result["baseline"] = {}
|
||||
if random_input_len is not None:
|
||||
@@ -800,12 +671,3 @@ 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)."
|
||||
),
|
||||
)
|
||||
@@ -15,15 +15,10 @@ from tests.dfx.conftest import (
|
||||
get_runtime_resource_label,
|
||||
is_diffusion_perf_config,
|
||||
load_benchmark_configs,
|
||||
resolve_baseline_value,
|
||||
run_benchmark,
|
||||
)
|
||||
from tests.helpers.runtime import OmniServer
|
||||
|
||||
# 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.
|
||||
#
|
||||
# Optional JSON field ``mark`` is applied as pytest marks via
|
||||
# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
|
||||
# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).
|
||||
@@ -162,36 +157,8 @@ def benchmark_params(request):
|
||||
}
|
||||
|
||||
|
||||
def assert_result(
|
||||
result,
|
||||
params,
|
||||
num_prompt,
|
||||
*,
|
||||
assert_baseline: bool,
|
||||
sweep_index: int | None = None,
|
||||
max_concurrency: Any | None = None,
|
||||
request_rate: Any | None = None,
|
||||
) -> None:
|
||||
def assert_result(result, num_prompt) -> 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_raw,
|
||||
sweep_index=sweep_index,
|
||||
max_concurrency=max_concurrency,
|
||||
request_rate=request_rate,
|
||||
)
|
||||
if "throughput" in metric_name:
|
||||
if current_value <= baseline_value:
|
||||
print(
|
||||
f"ERROR: Throughput test results were below baseline: {metric_name}: {current_value} > {baseline_value}"
|
||||
)
|
||||
else:
|
||||
if current_value >= baseline_value:
|
||||
print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}")
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@@ -200,7 +167,7 @@ def assert_result(
|
||||
paired_benchmark_params,
|
||||
indirect=["omni_server", "benchmark_params"],
|
||||
)
|
||||
def test_performance_benchmark(omni_server, benchmark_params, request):
|
||||
def test_performance_benchmark(omni_server, benchmark_params):
|
||||
test_name = benchmark_params["test_name"]
|
||||
params = benchmark_params["params"]
|
||||
dataset_name = params.get("dataset_name", "")
|
||||
@@ -212,7 +179,6 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
|
||||
print(f"Running benchmark for model: {model}")
|
||||
print(f"Benchmark parameters: {benchmark_params}")
|
||||
|
||||
assert_baseline = request.config.getoption("--assert-baseline", default=False)
|
||||
resource_label = get_runtime_resource_label()
|
||||
|
||||
def to_list(value, default=None):
|
||||
@@ -271,7 +237,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
|
||||
break
|
||||
|
||||
# QPS / request-rate sweep
|
||||
for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
|
||||
for qps, num_prompt in zip(qps_list, num_prompt_list):
|
||||
args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
|
||||
result = run_benchmark(
|
||||
args=args,
|
||||
@@ -280,24 +246,14 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
|
||||
dataset_name=dataset_name,
|
||||
num_prompt=num_prompt,
|
||||
baseline_config=params.get("baseline"),
|
||||
sweep_index=i,
|
||||
request_rate=qps,
|
||||
max_concurrency=None,
|
||||
random_input_len=params.get("random_input_len"),
|
||||
random_output_len=params.get("random_output_len"),
|
||||
resource_label=resource_label,
|
||||
)
|
||||
assert_result(
|
||||
result,
|
||||
params,
|
||||
num_prompt,
|
||||
assert_baseline=assert_baseline,
|
||||
sweep_index=i,
|
||||
request_rate=qps,
|
||||
)
|
||||
assert_result(result, num_prompt)
|
||||
|
||||
# concurrency test
|
||||
for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
|
||||
for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list):
|
||||
args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
|
||||
result = run_benchmark(
|
||||
args=args,
|
||||
@@ -306,18 +262,8 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
|
||||
dataset_name=dataset_name,
|
||||
num_prompt=num_prompt,
|
||||
baseline_config=params.get("baseline"),
|
||||
sweep_index=i,
|
||||
request_rate=None,
|
||||
max_concurrency=concurrency,
|
||||
random_input_len=params.get("random_input_len"),
|
||||
random_output_len=params.get("random_output_len"),
|
||||
resource_label=resource_label,
|
||||
)
|
||||
assert_result(
|
||||
result,
|
||||
params,
|
||||
num_prompt,
|
||||
assert_baseline=assert_baseline,
|
||||
sweep_index=i,
|
||||
max_concurrency=concurrency,
|
||||
)
|
||||
assert_result(result, num_prompt)
|
||||
@@ -13,8 +13,6 @@ A config JSON file may be passed via --test-config-file. If omitted, every ``*.j
|
||||
pytest run_diffusion_benchmark.py -m "diffusion"
|
||||
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).
|
||||
|
||||
Optional JSON field ``mark`` is applied as pytest marks on that case via
|
||||
``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``).
|
||||
|
||||
@@ -25,6 +23,7 @@ of all runs from cases in that file). Bulk load without ``--test-config-file`` u
|
||||
the same per-file aggregation; ``-m`` only selects which cases run.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
@@ -274,11 +273,6 @@ def _write_result_record(record: dict[str, Any]) -> Path:
|
||||
return target
|
||||
|
||||
|
||||
def _append_to_aggregated_file(record: dict[str, Any]) -> None:
|
||||
"""Backward-compatible wrapper; prefer :func:`_write_result_record`."""
|
||||
_write_result_record(record)
|
||||
|
||||
|
||||
_server_lock = threading.Lock()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -678,7 +672,7 @@ def run_benchmark(
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
|
||||
tmp_result_file = Path(tmp.name)
|
||||
|
||||
exclude_keys = {"baseline", "dataset", "task", "name", "skip-performance-assertion"}
|
||||
exclude_keys = {"baseline", "dataset", "task", "name"}
|
||||
|
||||
cmd = [
|
||||
sys.executable,
|
||||
@@ -772,12 +766,17 @@ def run_benchmark(
|
||||
completed = metrics.get("completed_requests", metrics.get("completed", 0))
|
||||
failed = metrics.get("failed_requests", metrics.get("failed", 0))
|
||||
|
||||
# Persist the JSON baseline as-is (no hardware / sweep resolution).
|
||||
baseline = copy.deepcopy(params.get("baseline") or {})
|
||||
metrics["baseline"] = baseline
|
||||
|
||||
record: dict[str, Any] = {
|
||||
"test_name": test_name,
|
||||
"endpoint": endpoint,
|
||||
"timestamp": timestamp,
|
||||
"server_params": server_cfg.get("server_params"),
|
||||
"benchmark_params": params,
|
||||
"baseline": baseline,
|
||||
"result": metrics,
|
||||
"log_file": str(log_file),
|
||||
"Model": model,
|
||||
@@ -827,70 +826,14 @@ def run_benchmark(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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."""
|
||||
if baseline_raw is None:
|
||||
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)):
|
||||
if sweep_index is None:
|
||||
raise ValueError("sweep_index is required when baseline is a list or tuple")
|
||||
return baseline_raw[sweep_index]
|
||||
return baseline_raw
|
||||
|
||||
|
||||
def assert_result(
|
||||
result: dict[str, Any],
|
||||
params: dict[str, Any],
|
||||
num_prompts: int,
|
||||
*,
|
||||
sweep_index: int | None = None,
|
||||
max_concurrency: Any = None,
|
||||
request_rate: Any = None,
|
||||
assert_baseline: bool = True,
|
||||
) -> None:
|
||||
"""Assert that benchmark metrics satisfy the configured baselines."""
|
||||
"""Assert that the expected number of requests completed."""
|
||||
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
|
||||
|
||||
for metric, baseline_raw in params.get("baseline", {}).items():
|
||||
current = result.get(metric)
|
||||
assert current is not None, f"Metric '{metric}' not found in result: {list(result.keys())}"
|
||||
threshold = _resolve_baseline_value(
|
||||
baseline_raw,
|
||||
sweep_index=sweep_index,
|
||||
max_concurrency=max_concurrency,
|
||||
request_rate=request_rate,
|
||||
)
|
||||
if "throughput" in metric:
|
||||
assert current >= threshold, f"{metric}: {current:.4f} < baseline {threshold}"
|
||||
else:
|
||||
assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}"
|
||||
|
||||
|
||||
def _default_benchmark_endpoint_for_task(task: str) -> str:
|
||||
"""Return the default client-side benchmark endpoint for a diffusion task."""
|
||||
@@ -919,7 +862,6 @@ def _build_run_params(
|
||||
params: dict[str, Any],
|
||||
*,
|
||||
num_prompts: int,
|
||||
sweep_index: int | None = None,
|
||||
request_rate: Any | None = None,
|
||||
max_concurrency: Any | None = None,
|
||||
) -> dict[str, Any]:
|
||||
@@ -931,16 +873,6 @@ def _build_run_params(
|
||||
run_params["request-rate"] = request_rate
|
||||
if max_concurrency is not None:
|
||||
run_params["max-concurrency"] = max_concurrency
|
||||
if "baseline" in params:
|
||||
run_params["baseline"] = {
|
||||
metric: _resolve_baseline_value(
|
||||
baseline_raw,
|
||||
sweep_index=sweep_index,
|
||||
max_concurrency=max_concurrency,
|
||||
request_rate=request_rate,
|
||||
)
|
||||
for metric, baseline_raw in params["baseline"].items()
|
||||
}
|
||||
return run_params
|
||||
|
||||
|
||||
@@ -963,36 +895,28 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
|
||||
sweep_runs: list[dict[str, Any]] = []
|
||||
|
||||
for i, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
|
||||
for request_rate, num_prompts in zip(request_rate_list, num_prompt_list):
|
||||
sweep_runs.append(
|
||||
{
|
||||
"params": _build_run_params(
|
||||
params,
|
||||
request_rate=request_rate,
|
||||
num_prompts=num_prompts,
|
||||
sweep_index=i,
|
||||
),
|
||||
"num_prompts": num_prompts,
|
||||
"sweep_index": i,
|
||||
"request_rate": request_rate,
|
||||
"max_concurrency": None,
|
||||
}
|
||||
)
|
||||
|
||||
for i, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
|
||||
for max_concurrency, num_prompts in zip(max_concurrency_list, num_prompt_list):
|
||||
sweep_runs.append(
|
||||
{
|
||||
"params": _build_run_params(
|
||||
params,
|
||||
max_concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
sweep_index=i,
|
||||
request_rate="inf",
|
||||
),
|
||||
"num_prompts": num_prompts,
|
||||
"sweep_index": i,
|
||||
"request_rate": None,
|
||||
"max_concurrency": max_concurrency,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -1005,9 +929,6 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
num_prompts=default_num_prompts,
|
||||
),
|
||||
"num_prompts": default_num_prompts,
|
||||
"sweep_index": None,
|
||||
"request_rate": None,
|
||||
"max_concurrency": None,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -1023,14 +944,12 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
paired_benchmark_params,
|
||||
indirect=["diffusion_server", "benchmark_params"],
|
||||
)
|
||||
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
|
||||
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
|
||||
"""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; metrics are recorded to the aggregated result file.
|
||||
|
||||
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"]
|
||||
@@ -1073,14 +992,4 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
|
||||
print(f"\n Aggregated results: {_aggregated_result_file_for_source(str(source))}")
|
||||
print("=" * 60)
|
||||
|
||||
assert_baseline = request.config.getoption("--assert-baseline", default=False)
|
||||
|
||||
assert_result(
|
||||
result,
|
||||
params,
|
||||
sweep_run["num_prompts"],
|
||||
sweep_index=sweep_run["sweep_index"],
|
||||
max_concurrency=sweep_run["max_concurrency"],
|
||||
request_rate=sweep_run["request_rate"],
|
||||
assert_baseline=assert_baseline,
|
||||
)
|
||||
assert_result(result, sweep_run["num_prompts"])
|
||||
@@ -40,10 +40,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.495,
|
||||
"latency_mean": 2.0238,
|
||||
"peak_memory_mb_max": 29393.4286,
|
||||
"peak_memory_mb_mean": 29392.7143
|
||||
"H100": {
|
||||
"throughput_qps": 0.495,
|
||||
"latency_mean": 2.0238,
|
||||
"peak_memory_mb_max": 29393.4286,
|
||||
"peak_memory_mb_mean": 29392.7143
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -90,10 +92,12 @@
|
||||
"enable-negative-prompt": true,
|
||||
"num-input-images": 1,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.3451,
|
||||
"latency_mean": 2.9031,
|
||||
"peak_memory_mb_max": 30480.2857,
|
||||
"peak_memory_mb_mean": 30479.2857
|
||||
"H100": {
|
||||
"throughput_qps": 0.3451,
|
||||
"latency_mean": 2.9031,
|
||||
"peak_memory_mb_max": 30480.2857,
|
||||
"peak_memory_mb_mean": 30479.2857
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -139,10 +143,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.5048,
|
||||
"latency_mean": 1.987,
|
||||
"peak_memory_mb_max": 29380.0,
|
||||
"peak_memory_mb_mean": 29380.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.5048,
|
||||
"latency_mean": 1.987,
|
||||
"peak_memory_mb_max": 29380.0,
|
||||
"peak_memory_mb_mean": 29380.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -189,10 +195,12 @@
|
||||
"enable-negative-prompt": true,
|
||||
"num-input-images": 1,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.149,
|
||||
"latency_mean": 6.7157,
|
||||
"peak_memory_mb_max": 30574.0,
|
||||
"peak_memory_mb_mean": 30574.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.149,
|
||||
"latency_mean": 6.7157,
|
||||
"peak_memory_mb_max": 30574.0,
|
||||
"peak_memory_mb_mean": 30574.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -24,10 +24,12 @@
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
|
||||
"latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
|
||||
"peak_memory_mb_max": 40500,
|
||||
"peak_memory_mb_mean": 40500
|
||||
"H100": {
|
||||
"throughput_qps": {"1": 0.0303, "2": 0.0302, "4": 0.0302},
|
||||
"latency_mean": {"1": 32.6, "2": 62.2, "4": 111.7},
|
||||
"peak_memory_mb_max": 40500,
|
||||
"peak_memory_mb_mean": 40500
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -24,10 +24,12 @@
|
||||
"warmup-num-inference-steps": 28,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
|
||||
"latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
|
||||
"peak_memory_mb_max": 40500,
|
||||
"peak_memory_mb_mean": 40500
|
||||
"H100": {
|
||||
"throughput_qps": {"1": 0.0549, "2": 0.0547, "4": 0.0549},
|
||||
"latency_mean": {"1": 18.0, "2": 34.3, "4": 61.4},
|
||||
"peak_memory_mb_max": 40500,
|
||||
"peak_memory_mb_mean": 40500
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -48,9 +48,11 @@
|
||||
"guardrails": false
|
||||
},
|
||||
"baseline": {
|
||||
"throughput_qps": 1.107,
|
||||
"latency_mean": 0.8929,
|
||||
"peak_memory_mb_mean": 80000.0
|
||||
"H100": {
|
||||
"throughput_qps": 1.107,
|
||||
"latency_mean": 0.8929,
|
||||
"peak_memory_mb_mean": 80000.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -115,9 +117,11 @@
|
||||
"use_duration_template": false
|
||||
},
|
||||
"baseline": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 36.6824,
|
||||
"peak_memory_mb_mean": 30155.4
|
||||
"H100": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 36.6824,
|
||||
"peak_memory_mb_mean": 30155.4
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -181,9 +185,11 @@
|
||||
"guardrails": false
|
||||
},
|
||||
"baseline": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 36.968,
|
||||
"peak_memory_mb_mean": 30443.6
|
||||
"H100": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 36.968,
|
||||
"peak_memory_mb_mean": 30443.6
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -251,9 +257,11 @@
|
||||
"condition_video_keep": "first"
|
||||
},
|
||||
"baseline": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 37.2117,
|
||||
"peak_memory_mb_mean": 30135.6
|
||||
"H100": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 37.2117,
|
||||
"peak_memory_mb_mean": 30135.6
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -46,12 +46,14 @@
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
223.3584
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.4393
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
223.3584
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.4393
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -70,15 +72,17 @@
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
276.575
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.5073
|
||||
],
|
||||
"audio_throughput": [
|
||||
27.3728
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
276.575
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.5073
|
||||
],
|
||||
"audio_throughput": [
|
||||
27.3728
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -133,13 +133,14 @@
|
||||
"bot-task": "think_recaption",
|
||||
"max-concurrency": 16,
|
||||
"num-prompts": 96,
|
||||
"skip-performance-assertion": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.1874,
|
||||
"latency_p99": 101.852,
|
||||
"latency_mean": 79.0385,
|
||||
"peak_memory_mb_max": 70000,
|
||||
"peak_memory_mb_mean": 70000
|
||||
"H100": {
|
||||
"throughput_qps": 0.1874,
|
||||
"latency_p99": 101.852,
|
||||
"latency_mean": 79.0385,
|
||||
"peak_memory_mb_max": 70000,
|
||||
"peak_memory_mb_mean": 70000
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -37,12 +37,13 @@
|
||||
"num-inference-steps": 8,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"skip-performance-assertion": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.21,
|
||||
"latency_p99": 4.7469,
|
||||
"peak_memory_mb_max": 101100,
|
||||
"peak_memory_mb_mean": 101100
|
||||
"H100": {
|
||||
"throughput_qps": 0.21,
|
||||
"latency_p99": 4.7469,
|
||||
"peak_memory_mb_max": 101100,
|
||||
"peak_memory_mb_mean": 101100
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -37,12 +37,13 @@
|
||||
"num-inference-steps": 8,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"skip-performance-assertion": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.20,
|
||||
"latency_p99": 5.1025,
|
||||
"peak_memory_mb_max": 97402,
|
||||
"peak_memory_mb_mean": 97402
|
||||
"H100": {
|
||||
"throughput_qps": 0.20,
|
||||
"latency_p99": 5.1025,
|
||||
"peak_memory_mb_max": 97402,
|
||||
"peak_memory_mb_mean": 97402
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -36,13 +36,14 @@
|
||||
"num-inference-steps": 8,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"skip-performance-assertion": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.213,
|
||||
"latency_mean": 4.6953,
|
||||
"latency_p99": 4.8601,
|
||||
"peak_memory_mb_max": 56912.0,
|
||||
"peak_memory_mb_mean": 56912.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.213,
|
||||
"latency_mean": 4.6953,
|
||||
"latency_p99": 4.8601,
|
||||
"peak_memory_mb_max": 56912.0,
|
||||
"peak_memory_mb_mean": 56912.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -26,8 +26,7 @@
|
||||
"warmup-requests": 1,
|
||||
"warmup-concurrency": 1,
|
||||
"warmup-num-inference-steps": 2,
|
||||
"enable-negative-prompt": true,
|
||||
"skip-performance-assertion": true
|
||||
"enable-negative-prompt": true
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -44,27 +44,29 @@
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
96.4061,
|
||||
140.8147,
|
||||
271.9025,
|
||||
362.2916,
|
||||
507.776
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
18507.4582,
|
||||
28365.4956,
|
||||
31907.9613,
|
||||
48161.5829,
|
||||
72630.8601
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1175,
|
||||
0.18,
|
||||
0.2359,
|
||||
0.3278,
|
||||
0.6296
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
96.4061,
|
||||
140.8147,
|
||||
271.9025,
|
||||
362.2916,
|
||||
507.776
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
18507.4582,
|
||||
28365.4956,
|
||||
31907.9613,
|
||||
48161.5829,
|
||||
72630.8601
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1175,
|
||||
0.18,
|
||||
0.2359,
|
||||
0.3278,
|
||||
0.6296
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -91,15 +93,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
354.0951
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
3952.5615
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.117
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
354.0951
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
3952.5615
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.117
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -128,15 +132,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
259.6317
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4253.3009
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1448
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
259.6317
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4253.3009
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1448
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -167,15 +173,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
457.4453
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
5912.5623
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1967
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
457.4453
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
5912.5623
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1967
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -206,20 +214,22 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
90.8033,
|
||||
202.1779,
|
||||
326.5864,
|
||||
544.8914,
|
||||
961.3929
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4517.3411,
|
||||
7857.9711,
|
||||
9861.1687,
|
||||
13283.3363,
|
||||
19541.0736
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
90.8033,
|
||||
202.1779,
|
||||
326.5864,
|
||||
544.8914,
|
||||
961.3929
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4517.3411,
|
||||
7857.9711,
|
||||
9861.1687,
|
||||
13283.3363,
|
||||
19541.0736
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -251,12 +261,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
318.048
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
807.6114
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
318.048
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
807.6114
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -290,12 +302,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
186.168
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
726.1137
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
186.168
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
726.1137
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -331,12 +345,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
375.5282
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1006.1663
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
375.5282
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1006.1663
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -35,10 +35,12 @@
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085],
|
||||
"mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485],
|
||||
"mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275],
|
||||
"mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085],
|
||||
"mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485],
|
||||
"mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275],
|
||||
"mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -61,9 +63,11 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [340.5122],
|
||||
"mean_audio_ttfp_ms": [463.0752],
|
||||
"mean_audio_rtf": [0.1153]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [340.5122],
|
||||
"mean_audio_ttfp_ms": [463.0752],
|
||||
"mean_audio_rtf": [0.1153]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -88,9 +92,11 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [236.926],
|
||||
"mean_audio_ttfp_ms": [395.2182],
|
||||
"mean_audio_rtf": [0.1259]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [236.926],
|
||||
"mean_audio_ttfp_ms": [395.2182],
|
||||
"mean_audio_rtf": [0.1259]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -117,9 +123,11 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [470.6126],
|
||||
"mean_audio_ttfp_ms": [665.5474],
|
||||
"mean_audio_rtf": [0.166]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [470.6126],
|
||||
"mean_audio_ttfp_ms": [665.5474],
|
||||
"mean_audio_rtf": [0.166]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -44,27 +44,29 @@
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
115.8348,
|
||||
173.1823,
|
||||
446.672,
|
||||
464.5985,
|
||||
3420.0634
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
25070.9523,
|
||||
31992.0646,
|
||||
37619.7669,
|
||||
50393.8257,
|
||||
136619.2743
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.169,
|
||||
0.2111,
|
||||
0.2475,
|
||||
0.3551,
|
||||
0.9266
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
115.8348,
|
||||
173.1823,
|
||||
446.672,
|
||||
464.5985,
|
||||
3420.0634
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
25070.9523,
|
||||
31992.0646,
|
||||
37619.7669,
|
||||
50393.8257,
|
||||
136619.2743
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.169,
|
||||
0.2111,
|
||||
0.2475,
|
||||
0.3551,
|
||||
0.9266
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -91,15 +93,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
343.1377
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4811.1123
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1463
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
343.1377
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4811.1123
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1463
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -128,15 +132,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
245.5446
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
5019.1256
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1713
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
245.5446
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
5019.1256
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1713
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -167,15 +173,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
506.1407
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
6814.9656
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.2213
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
506.1407
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
6814.9656
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.2213
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -51,20 +51,22 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
69.8081,
|
||||
173.0651,
|
||||
271.8108,
|
||||
413.4079,
|
||||
497.2546
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4720.8225,
|
||||
7937.2532,
|
||||
9962.3566,
|
||||
13724.5479,
|
||||
18324.091
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
69.8081,
|
||||
173.0651,
|
||||
271.8108,
|
||||
413.4079,
|
||||
497.2546
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4720.8225,
|
||||
7937.2532,
|
||||
9962.3566,
|
||||
13724.5479,
|
||||
18324.091
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -97,12 +99,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
25.9477
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
517.3155
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
25.9477
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
517.3155
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -137,12 +141,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
26.7056
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
551.308
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
26.7056
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
551.308
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -179,12 +185,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [
|
||||
27.6906
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
596.5284
|
||||
]
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
27.6906
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
596.5284
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -34,10 +34,12 @@
|
||||
"num-input-images": 2,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0721,
|
||||
"latency_mean": 13.8398,
|
||||
"peak_memory_mb_max": 57610.0,
|
||||
"peak_memory_mb_mean": 57610.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.0721,
|
||||
"latency_mean": 13.8398,
|
||||
"peak_memory_mb_max": 57610.0,
|
||||
"peak_memory_mb_mean": 57610.0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -52,10 +54,12 @@
|
||||
"num-input-images": 2,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0186,
|
||||
"latency_mean": 53.6942,
|
||||
"peak_memory_mb_max": 67790.0,
|
||||
"peak_memory_mb_mean": 67790.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.0186,
|
||||
"latency_mean": 53.6942,
|
||||
"peak_memory_mb_max": 67790.0,
|
||||
"peak_memory_mb_mean": 67790.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -100,10 +104,12 @@
|
||||
"num-input-images": 2,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0538,
|
||||
"latency_mean": 18.5591,
|
||||
"peak_memory_mb_max": 56449.4286,
|
||||
"peak_memory_mb_mean": 56449.4286
|
||||
"H100": {
|
||||
"throughput_qps": 0.0538,
|
||||
"latency_mean": 18.5591,
|
||||
"peak_memory_mb_max": 56449.4286,
|
||||
"peak_memory_mb_mean": 56449.4286
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -157,10 +163,12 @@
|
||||
"num-input-images": 2,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.2776,
|
||||
"latency_mean": 3.5938,
|
||||
"peak_memory_mb_max": 58113.7143,
|
||||
"peak_memory_mb_mean": 58075.1429
|
||||
"H100": {
|
||||
"throughput_qps": 0.2776,
|
||||
"latency_mean": 3.5938,
|
||||
"peak_memory_mb_max": 58113.7143,
|
||||
"peak_memory_mb_mean": 58075.1429
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -175,10 +183,12 @@
|
||||
"num-input-images": 2,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.1051,
|
||||
"latency_mean": 9.4918,
|
||||
"peak_memory_mb_max": 68107.4286,
|
||||
"peak_memory_mb_mean": 68107.4286
|
||||
"H100": {
|
||||
"throughput_qps": 0.1051,
|
||||
"latency_mean": 9.4918,
|
||||
"peak_memory_mb_max": 68107.4286,
|
||||
"peak_memory_mb_mean": 68107.4286
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -33,10 +33,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.1129,
|
||||
"latency_mean": 8.8352,
|
||||
"peak_memory_mb_max": 58572.0,
|
||||
"peak_memory_mb_mean": 58572.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.1129,
|
||||
"latency_mean": 8.8352,
|
||||
"peak_memory_mb_max": 58572.0,
|
||||
"peak_memory_mb_mean": 58572.0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -50,10 +52,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.024,
|
||||
"latency_mean": 41.5686,
|
||||
"peak_memory_mb_max": 66368.0,
|
||||
"peak_memory_mb_mean": 66368.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.024,
|
||||
"latency_mean": 41.5686,
|
||||
"peak_memory_mb_max": 66368.0,
|
||||
"peak_memory_mb_mean": 66368.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -96,10 +100,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0672,
|
||||
"latency_mean": 14.8517,
|
||||
"peak_memory_mb_max": 57076.0,
|
||||
"peak_memory_mb_mean": 57076.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.0672,
|
||||
"latency_mean": 14.8517,
|
||||
"peak_memory_mb_max": 57076.0,
|
||||
"peak_memory_mb_mean": 57076.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -152,10 +158,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.2879,
|
||||
"latency_mean": 3.4701,
|
||||
"peak_memory_mb_max": 58758.0,
|
||||
"peak_memory_mb_mean": 58758.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.2879,
|
||||
"latency_mean": 3.4701,
|
||||
"peak_memory_mb_max": 58758.0,
|
||||
"peak_memory_mb_mean": 58758.0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -169,10 +177,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.1113,
|
||||
"latency_mean": 8.9656,
|
||||
"peak_memory_mb_max": 67424.0,
|
||||
"peak_memory_mb_mean": 67424.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.1113,
|
||||
"latency_mean": 8.9656,
|
||||
"peak_memory_mb_max": 67424.0,
|
||||
"peak_memory_mb_mean": 67424.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -33,10 +33,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0703,
|
||||
"latency_mean": 14.208,
|
||||
"peak_memory_mb_max": 63225.1429,
|
||||
"peak_memory_mb_mean": 63225.1429
|
||||
"H100": {
|
||||
"throughput_qps": 0.0703,
|
||||
"latency_mean": 14.208,
|
||||
"peak_memory_mb_max": 63225.1429,
|
||||
"peak_memory_mb_mean": 63225.1429
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -50,10 +52,12 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0417,
|
||||
"latency_mean": 23.914,
|
||||
"peak_memory_mb_max": 63225.1429,
|
||||
"peak_memory_mb_mean": 63225.1429
|
||||
"H100": {
|
||||
"throughput_qps": 0.0417,
|
||||
"latency_mean": 23.914,
|
||||
"peak_memory_mb_max": 63225.1429,
|
||||
"peak_memory_mb_mean": 63225.1429
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -33,9 +33,11 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.4153,
|
||||
"latency_mean": 2.4065,
|
||||
"peak_memory_mb_mean": 55121.4286
|
||||
"H100": {
|
||||
"throughput_qps": 0.4153,
|
||||
"latency_mean": 2.4065,
|
||||
"peak_memory_mb_mean": 55121.4286
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -49,9 +51,11 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0423,
|
||||
"latency_mean": 23.5991,
|
||||
"peak_memory_mb_mean": 65118.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.0423,
|
||||
"latency_mean": 23.5991,
|
||||
"peak_memory_mb_mean": 65118.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -91,9 +95,11 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.4175,
|
||||
"latency_mean": 2.3953,
|
||||
"peak_memory_mb_mean": 55088.5714
|
||||
"H100": {
|
||||
"throughput_qps": 0.4175,
|
||||
"latency_mean": 2.3953,
|
||||
"peak_memory_mb_mean": 55088.5714
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -107,9 +113,11 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0423,
|
||||
"latency_mean": 23.5768,
|
||||
"peak_memory_mb_mean": 66012.2857
|
||||
"H100": {
|
||||
"throughput_qps": 0.0423,
|
||||
"latency_mean": 23.5768,
|
||||
"peak_memory_mb_mean": 66012.2857
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -152,9 +160,11 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.1214,
|
||||
"latency_mean": 8.2211,
|
||||
"peak_memory_mb_mean": 54501.1429
|
||||
"H100": {
|
||||
"throughput_qps": 0.1214,
|
||||
"latency_mean": 8.2211,
|
||||
"peak_memory_mb_mean": 54501.1429
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -207,9 +217,11 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.559,
|
||||
"latency_mean": 1.794,
|
||||
"peak_memory_mb_mean": 55290.5714
|
||||
"H100": {
|
||||
"throughput_qps": 0.559,
|
||||
"latency_mean": 1.794,
|
||||
"peak_memory_mb_mean": 55290.5714
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -223,9 +235,11 @@
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.1918,
|
||||
"latency_mean": 5.2022,
|
||||
"peak_memory_mb_mean": 64838.8571
|
||||
"H100": {
|
||||
"throughput_qps": 0.1918,
|
||||
"latency_mean": 5.2022,
|
||||
"peak_memory_mb_mean": 64838.8571
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -269,9 +283,11 @@
|
||||
"warmup-num-inference-steps": 20,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
|
||||
"latency_mean": [2.4215, 2.992, 5.9293, 9.656],
|
||||
"peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821]
|
||||
"H100": {
|
||||
"throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
|
||||
"latency_mean": [2.4215, 2.992, 5.9293, 9.656],
|
||||
"peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
+137
-111
@@ -39,12 +39,14 @@
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
141.2791
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1591
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
141.2791
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1591
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -67,21 +69,23 @@
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
226.747,
|
||||
647.1272,
|
||||
4777.2714
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.3313,
|
||||
0.3642,
|
||||
1.2282
|
||||
],
|
||||
"audio_throughput": [
|
||||
24.2296,
|
||||
40.3552,
|
||||
43.2572
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
226.747,
|
||||
647.1272,
|
||||
4777.2714
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.3313,
|
||||
0.3642,
|
||||
1.2282
|
||||
],
|
||||
"audio_throughput": [
|
||||
24.2296,
|
||||
40.3552,
|
||||
43.2572
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -102,9 +106,11 @@
|
||||
"seed_tts_wer_eval": true,
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_audio_rtf": [
|
||||
0.45
|
||||
]
|
||||
"H100": {
|
||||
"mean_audio_rtf": [
|
||||
0.45
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -153,12 +159,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
46.8434
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1367
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
46.8434
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1367
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -186,21 +194,23 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
75.3341,
|
||||
713.7727,
|
||||
5941.7179
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1859,
|
||||
0.3148,
|
||||
1.1662
|
||||
],
|
||||
"audio_throughput": [
|
||||
38.3858,
|
||||
55.3523,
|
||||
68.7032
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
75.3341,
|
||||
713.7727,
|
||||
5941.7179
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1859,
|
||||
0.3148,
|
||||
1.1662
|
||||
],
|
||||
"audio_throughput": [
|
||||
38.3858,
|
||||
55.3523,
|
||||
68.7032
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -226,9 +236,11 @@
|
||||
"seed_tts_wer_eval": true,
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_audio_rtf": [
|
||||
0.35
|
||||
]
|
||||
"H100": {
|
||||
"mean_audio_rtf": [
|
||||
0.35
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -252,15 +264,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
56.746
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1571
|
||||
],
|
||||
"audio_throughput": [
|
||||
11.1226
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
56.746
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1571
|
||||
],
|
||||
"audio_throughput": [
|
||||
11.1226
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -283,12 +297,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
46.3479
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1352
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
46.3479
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1352
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -315,21 +331,23 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
75.5967,
|
||||
587.5877,
|
||||
4960.3726
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1871,
|
||||
0.3143,
|
||||
1.1703
|
||||
],
|
||||
"audio_throughput": [
|
||||
38.9378,
|
||||
55.1846,
|
||||
71.5708
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
75.5967,
|
||||
587.5877,
|
||||
4960.3726
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1871,
|
||||
0.3143,
|
||||
1.1703
|
||||
],
|
||||
"audio_throughput": [
|
||||
38.9378,
|
||||
55.1846,
|
||||
71.5708
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -352,15 +370,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
55.5029
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1508
|
||||
],
|
||||
"audio_throughput": [
|
||||
9.4885
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
55.5029
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1508
|
||||
],
|
||||
"audio_throughput": [
|
||||
9.4885
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -402,12 +422,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
82.1886
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1241
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
82.1886
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1241
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -429,15 +451,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
160.9168
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1907
|
||||
],
|
||||
"audio_throughput": [
|
||||
39.2332
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
160.9168
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1907
|
||||
],
|
||||
"audio_throughput": [
|
||||
39.2332
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -459,15 +483,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
100.1265
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1437
|
||||
],
|
||||
"audio_throughput": [
|
||||
11.7605
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
100.1265
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1437
|
||||
],
|
||||
"audio_throughput": [
|
||||
11.7605
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -37,12 +37,14 @@
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
137.645
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.0497
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
137.645
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.0497
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -61,15 +63,17 @@
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
475.3357
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1243
|
||||
],
|
||||
"audio_throughput": [
|
||||
36.6358
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
475.3357
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1243
|
||||
],
|
||||
"audio_throughput": [
|
||||
36.6358
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -90,9 +94,11 @@
|
||||
"seed_tts_wer_eval": true,
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_audio_rtf": [
|
||||
0.45
|
||||
]
|
||||
"H100": {
|
||||
"mean_audio_rtf": [
|
||||
0.45
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -114,12 +120,14 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
79.6927
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.0423
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
79.6927
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.0423
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -141,15 +149,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
262.6212
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.0753
|
||||
],
|
||||
"audio_throughput": [
|
||||
103.5843
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
262.6212
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.0753
|
||||
],
|
||||
"audio_throughput": [
|
||||
103.5843
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -171,15 +181,17 @@
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
80.1095
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.0479
|
||||
],
|
||||
"audio_throughput": [
|
||||
17.2918
|
||||
]
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
80.1095
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.0479
|
||||
],
|
||||
"audio_throughput": [
|
||||
17.2918
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -203,9 +215,11 @@
|
||||
"seed_tts_wer_eval": true,
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_audio_rtf": [
|
||||
0.35
|
||||
]
|
||||
"H100": {
|
||||
"mean_audio_rtf": [
|
||||
0.35
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -42,9 +42,11 @@
|
||||
}
|
||||
],
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0434,
|
||||
"latency_mean": 23.0654,
|
||||
"peak_memory_mb_mean": 76360.0
|
||||
"H100": {
|
||||
"throughput_qps": 0.0434,
|
||||
"latency_mean": 23.0654,
|
||||
"peak_memory_mb_mean": 76360.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -97,9 +99,11 @@
|
||||
}
|
||||
],
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0606,
|
||||
"latency_mean": 16.4928,
|
||||
"peak_memory_mb_mean": 47164.8
|
||||
"H100": {
|
||||
"throughput_qps": 0.0606,
|
||||
"latency_mean": 16.4928,
|
||||
"peak_memory_mb_mean": 47164.8
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
@@ -122,9 +126,11 @@
|
||||
}
|
||||
],
|
||||
"baseline": {
|
||||
"throughput_qps": 0.0124,
|
||||
"latency_mean": 80.7784,
|
||||
"peak_memory_mb_mean": 55374.3429
|
||||
"H100": {
|
||||
"throughput_qps": 0.0124,
|
||||
"latency_mean": 80.7784,
|
||||
"peak_memory_mb_mean": 55374.3429
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Omni sleep/wakeup JSON validation (live ``Qwen/Qwen-Image`` process)."""
|
||||
"""Omni sleep/wakeup JSON validation and level-2 wake contract (live ``Qwen/Qwen-Image``)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -15,9 +15,17 @@ 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")
|
||||
|
||||
# One module-scoped server for this file. ``--enable-sleep-mode`` is required by the
|
||||
# level-2 wake contract test; invalid-JSON cases only hit request validation and are
|
||||
# unaffected. Do not mix ``omni_server`` (module) with ``omni_server_function`` here:
|
||||
# ``iter_omni_server`` holds ``omni_fixture_lock`` across the module yield, so a
|
||||
# function-scoped second server would deadlock / contend for the GPU.
|
||||
_QWEN_IMAGE = [
|
||||
pytest.param(
|
||||
OmniServerParams(model="Qwen/Qwen-Image"),
|
||||
OmniServerParams(
|
||||
model="Qwen/Qwen-Image",
|
||||
server_args=["--enable-sleep-mode"],
|
||||
),
|
||||
id="qwen_image",
|
||||
marks=hardware_marks(res={"cuda": "H100"}),
|
||||
),
|
||||
@@ -93,3 +101,37 @@ def test_omni_sleep_wakeup_invalid_requests(
|
||||
openai_client.send_omni_sleep_http_request(cfg)
|
||||
else:
|
||||
openai_client.send_omni_wakeup_http_request(cfg)
|
||||
|
||||
|
||||
# Keep this after the invalid-JSON cases: sleep(level=2) discards weights and the
|
||||
# shared module server is not usable afterward; module teardown then stops it.
|
||||
@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
|
||||
def test_wakeup_after_level2_sleep_fails(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
) -> None:
|
||||
"""Regression for #4473 Repro A via HTTP: wake after sleep(level=2) → 501.
|
||||
|
||||
Moved from ``tests/entrypoints/test_omni_sleep_mode.py`` (BAGEL cold start).
|
||||
Live serve surfaces ``NotImplementedError`` as OpenAI-style JSON with code 501.
|
||||
"""
|
||||
assert omni_server is not None
|
||||
sleep_resps = openai_client.send_omni_sleep_http_request(
|
||||
{
|
||||
"json": {"stage_ids": [0], "level": 2},
|
||||
"timeout": 180,
|
||||
}
|
||||
)
|
||||
sleep_resp = sleep_resps[0]
|
||||
assert sleep_resp.status_code == 200, sleep_resp
|
||||
assert sleep_resp.json_body is not None
|
||||
assert sleep_resp.json_body.get("status") == "SUCCESS", sleep_resp.json_body
|
||||
|
||||
openai_client.send_omni_wakeup_http_request(
|
||||
{
|
||||
"json": {"stage_ids": [0]},
|
||||
"timeout": 120,
|
||||
"err_code": 501,
|
||||
"err_message": ("sleep(level=2)", "not yet implemented"),
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,64 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from contextlib import nullcontext
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
|
||||
observed = {}
|
||||
|
||||
def fake_sdpa(query, key, value, **kwargs):
|
||||
observed.update(query=query, key=key, value=value, kwargs=kwargs)
|
||||
return query
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
|
||||
lambda _backends: nullcontext(),
|
||||
)
|
||||
monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa)
|
||||
impl = CuDNNAttentionImpl(
|
||||
num_heads=2,
|
||||
head_size=4,
|
||||
softmax_scale=0.5,
|
||||
)
|
||||
query = torch.randn(1, 8, 2, 4)
|
||||
key = torch.randn_like(query)
|
||||
value = torch.randn_like(query)
|
||||
|
||||
output = impl.forward_cuda(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
AttentionMetadata(extra={"valid_kv_length": 5}),
|
||||
)
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert observed["query"].shape == (1, 2, 8, 4)
|
||||
assert observed["key"].shape == (1, 2, 5, 4)
|
||||
assert observed["value"].shape == (1, 2, 5, 4)
|
||||
assert observed["kwargs"]["attn_mask"] is None
|
||||
|
||||
|
||||
def test_cudnn_rejects_invalid_valid_kv_length():
|
||||
impl = CuDNNAttentionImpl(
|
||||
num_heads=2,
|
||||
head_size=4,
|
||||
softmax_scale=0.5,
|
||||
)
|
||||
query = torch.randn(1, 8, 2, 4)
|
||||
|
||||
with pytest.raises(ValueError, match="valid_kv_length"):
|
||||
impl.forward_cuda(
|
||||
query,
|
||||
query,
|
||||
query,
|
||||
AttentionMetadata(extra={"valid_kv_length": 9}),
|
||||
)
|
||||
@@ -30,6 +30,13 @@ def test_pipeline_import_registry_and_component_discovery():
|
||||
assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
|
||||
|
||||
|
||||
def test_dlo_offload_plan_includes_token_refiner():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
assert MiniMaxH3Pipeline._offload_plan.offload_submodules == {"token_refiner": "blocks"}
|
||||
assert MiniMaxH3Pipeline._offload_plan.resident_dit_paths == frozenset({"transformer"})
|
||||
|
||||
|
||||
def test_joint_postprocess_is_multiprocessing_picklable():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import (
|
||||
get_minimax_h3_post_process_func,
|
||||
@@ -87,6 +94,87 @@ def test_shifted_sigma_schedule_matches_reference_values():
|
||||
)
|
||||
|
||||
|
||||
def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||
MiniMaxH3Attention,
|
||||
)
|
||||
|
||||
class FakeBackend:
|
||||
supports_prefix_kv_slicing = True
|
||||
|
||||
class FakeAttention(torch.nn.Module):
|
||||
attn_backend = FakeBackend
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.metadata = None
|
||||
|
||||
def forward(self, query, key, value, metadata):
|
||||
self.metadata = metadata
|
||||
return query
|
||||
|
||||
attention = object.__new__(MiniMaxH3Attention)
|
||||
torch.nn.Module.__init__(attention)
|
||||
attention.attention = FakeAttention()
|
||||
# Model-side Ulysses shards rows before Attention gathers them again.
|
||||
# The Python packed_total must therefore remain global even though q is
|
||||
# local here.
|
||||
q = torch.randn(4, 2, 4)
|
||||
cu_seqlens = torch.tensor([0, 5, 8], dtype=torch.int32)
|
||||
|
||||
output = attention._run_packed_attention(
|
||||
q,
|
||||
q,
|
||||
q,
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=5,
|
||||
packed_total=8,
|
||||
)
|
||||
|
||||
assert output.shape == q.shape
|
||||
assert attention.attention.metadata.attn_mask is None
|
||||
assert attention.attention.metadata.extra["valid_kv_length"] == 5
|
||||
|
||||
|
||||
def test_packed_attention_keeps_padding_mask_for_other_backends():
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||
MiniMaxH3Attention,
|
||||
)
|
||||
|
||||
class FakeBackend:
|
||||
supports_prefix_kv_slicing = False
|
||||
|
||||
class FakeAttention(torch.nn.Module):
|
||||
attn_backend = FakeBackend
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.metadata = None
|
||||
|
||||
def forward(self, query, key, value, metadata):
|
||||
self.metadata = metadata
|
||||
return query
|
||||
|
||||
attention = object.__new__(MiniMaxH3Attention)
|
||||
torch.nn.Module.__init__(attention)
|
||||
attention.attention = FakeAttention()
|
||||
q = torch.randn(8, 2, 4)
|
||||
|
||||
attention._run_packed_attention(
|
||||
q,
|
||||
q,
|
||||
q,
|
||||
cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
|
||||
max_seqlen=5,
|
||||
packed_total=8,
|
||||
)
|
||||
|
||||
assert torch.equal(
|
||||
attention.attention.metadata.attn_mask,
|
||||
torch.tensor([[True, True, True, True, True, False, False, False]]),
|
||||
)
|
||||
|
||||
|
||||
def test_reference_image_resize_contract():
|
||||
from PIL import Image
|
||||
|
||||
@@ -102,6 +190,52 @@ def test_reference_image_resize_contract():
|
||||
_reference_image_shape(Image.new("RGB", (100, 501)))
|
||||
|
||||
|
||||
def test_fl2va_supports_first_last_and_explicit_frame_index_contracts():
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_resolve_fl2va_keyframe_indices,
|
||||
)
|
||||
|
||||
assert _resolve_fl2va_keyframe_indices({}, 1) == [0]
|
||||
assert _resolve_fl2va_keyframe_indices({}, 2) == [0, -1]
|
||||
assert _resolve_fl2va_keyframe_indices({"frame_index": -1}, 1) == [-1]
|
||||
assert _resolve_fl2va_keyframe_indices({"frame_indices": [0, -1]}, 2) == [0, -1]
|
||||
with pytest.raises(ValueError, match="frame_indices"):
|
||||
_resolve_fl2va_keyframe_indices({"frame_indices": [0, 1]}, 2)
|
||||
|
||||
|
||||
def test_minimax_h3_uses_the_official_output_canvas_policy():
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_resolve_output_canvas,
|
||||
)
|
||||
|
||||
assert _resolve_output_canvas(21 / 9, 768) == (672, 1536)
|
||||
assert _resolve_output_canvas(16 / 9, 768) == (768, 1344)
|
||||
assert _resolve_output_canvas(9 / 16, 768) == (1344, 768)
|
||||
with pytest.raises(ValueError, match="short_edge"):
|
||||
_resolve_output_canvas(16 / 9, 720)
|
||||
|
||||
|
||||
def test_minimax_h3_accepts_sglang_auto_aspect_ratio_alias():
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
sampling = SimpleNamespace(
|
||||
fps=24,
|
||||
num_frames=1,
|
||||
height=None,
|
||||
width=None,
|
||||
extra_args={"duration": 5.0, "aspect_ratio": "auto"},
|
||||
)
|
||||
height, width, *_ = pipeline._resolve_shape("ref2va", sampling, None)
|
||||
assert (height, width) == (768, 1344)
|
||||
|
||||
|
||||
def test_minimax_h3_advertises_the_official_ref2va_image_limit():
|
||||
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
|
||||
|
||||
assert get_diffusion_model_metadata("MiniMaxH3Pipeline").max_multimodal_image_inputs == 9
|
||||
|
||||
|
||||
def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
||||
MiniMaxH3Qwen3VLEncoder,
|
||||
@@ -247,6 +381,127 @@ def test_layerwise_offload_releases_text_encoder():
|
||||
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
|
||||
|
||||
|
||||
def test_distributed_layerwise_offload_releases_text_encoder():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=False,
|
||||
enable_layerwise_offload=False,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
)
|
||||
pipeline.text_encoder = Mock()
|
||||
expected = torch.ones(2, 3)
|
||||
pipeline.text_encoder.encode_ids.return_value = expected
|
||||
|
||||
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
|
||||
|
||||
assert actual is expected
|
||||
pipeline.text_encoder.load_to_device.assert_called_once_with()
|
||||
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
|
||||
|
||||
|
||||
def test_distributed_layerwise_offload_stages_vae_component():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_layerwise_offload=False,
|
||||
enable_distributed_layerwise_offload=True,
|
||||
)
|
||||
component = Mock()
|
||||
|
||||
with pipeline._component_on_device(component):
|
||||
component.load_to_device.assert_called_once_with()
|
||||
component.offload_to_cpu.assert_not_called()
|
||||
|
||||
component.offload_to_cpu.assert_called_once_with()
|
||||
|
||||
|
||||
def test_distributed_layerwise_resident_blocks_are_stage_scoped():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
controller = Mock()
|
||||
pipeline._dlo_residency_controller = controller
|
||||
|
||||
with pipeline._resident_dit_layers_on_device():
|
||||
controller.load_resident_layers.assert_called_once_with()
|
||||
controller.offload_resident_layers.assert_not_called()
|
||||
|
||||
controller.offload_resident_layers.assert_called_once_with()
|
||||
|
||||
|
||||
def test_distributed_layerwise_resident_blocks_release_on_failure():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
controller = Mock()
|
||||
pipeline._dlo_residency_controller = controller
|
||||
|
||||
with pytest.raises(RuntimeError, match="denoise failed"):
|
||||
with pipeline._resident_dit_layers_on_device():
|
||||
raise RuntimeError("denoise failed")
|
||||
|
||||
controller.load_resident_layers.assert_called_once_with()
|
||||
controller.offload_resident_layers.assert_called_once_with()
|
||||
|
||||
|
||||
def test_encoder_layerwise_offload_keeps_tp_blocks_rank_local(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3.encoder import (
|
||||
MiniMaxH3Qwen3VLEncoder,
|
||||
)
|
||||
|
||||
class Stack(torch.nn.Module):
|
||||
def __init__(self, count):
|
||||
super().__init__()
|
||||
self.blocks = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
|
||||
|
||||
class TextStack(torch.nn.Module):
|
||||
def __init__(self, count):
|
||||
super().__init__()
|
||||
self.layers = torch.nn.ModuleList([torch.nn.Linear(2, 2) for _ in range(count)])
|
||||
|
||||
hooks = []
|
||||
|
||||
def fake_apply(block, next_block, device, stream, pin_memory):
|
||||
hook = SimpleNamespace(
|
||||
block=block,
|
||||
next_block=next_block,
|
||||
device=device,
|
||||
stream=stream,
|
||||
pin_memory=pin_memory,
|
||||
_prev_hook=None,
|
||||
offload_layer=Mock(),
|
||||
)
|
||||
hooks.append(hook)
|
||||
return hook
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.offloader.layerwise_backend.apply_block_hook",
|
||||
fake_apply,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.platforms.current_omni_platform.Stream",
|
||||
Mock(return_value="copy-stream"),
|
||||
)
|
||||
encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
|
||||
torch.nn.Module.__init__(encoder)
|
||||
encoder.device_target = torch.device("cpu")
|
||||
encoder.vision = Stack(2)
|
||||
encoder.text_model = TextStack(3)
|
||||
|
||||
encoder.enable_omni_layerwise_offload(pin_memory=False)
|
||||
|
||||
assert len(hooks) == 5
|
||||
assert all(hook._prev_hook is not None for hook in hooks)
|
||||
assert all(hook.device == torch.device("cpu") for hook in hooks)
|
||||
|
||||
|
||||
def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
|
||||
|
||||
@@ -278,6 +533,82 @@ def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
|
||||
assert next(video_vae.parameters()).dtype == torch.float32
|
||||
|
||||
|
||||
def test_video_vae_can_load_on_cpu_for_staged_gpu_residency(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
|
||||
|
||||
class FakeRemote(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.model = torch.nn.Linear(1, 1)
|
||||
|
||||
monkeypatch.setattr(
|
||||
vae_module,
|
||||
"_load_component_config",
|
||||
lambda _path: {
|
||||
"latent_channels": 1,
|
||||
"latents_mean": [0.0],
|
||||
"latents_std": [1.0],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
vae_module,
|
||||
"_load_remote_component",
|
||||
lambda _path, _config: FakeRemote(),
|
||||
)
|
||||
|
||||
video_vae = vae_module.MiniMaxH3VideoVAE(
|
||||
"unused",
|
||||
device=torch.device("meta"),
|
||||
load_device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
assert next(video_vae.parameters()).device.type == "cpu"
|
||||
assert video_vae._device_target.type == "meta"
|
||||
|
||||
|
||||
def test_video_vae_uses_snapshot_stager_for_accelerator_target(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
|
||||
|
||||
class FakeRemote(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.model = torch.nn.Linear(1, 1)
|
||||
|
||||
stager = Mock()
|
||||
stager_cls = Mock(return_value=stager)
|
||||
monkeypatch.setattr(
|
||||
vae_module,
|
||||
"_load_component_config",
|
||||
lambda _path: {
|
||||
"latent_channels": 1,
|
||||
"latents_mean": [0.0],
|
||||
"latents_std": [1.0],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
vae_module,
|
||||
"_load_remote_component",
|
||||
lambda _path, _config: FakeRemote(),
|
||||
)
|
||||
monkeypatch.setattr(vae_module, "PinnedModuleStager", stager_cls)
|
||||
|
||||
video_vae = vae_module.MiniMaxH3VideoVAE(
|
||||
"unused",
|
||||
device=torch.device("cuda"),
|
||||
load_device=torch.device("cpu"),
|
||||
)
|
||||
video_vae.load_to_device()
|
||||
video_vae.offload_to_cpu()
|
||||
|
||||
stager_cls.assert_called_once_with(
|
||||
video_vae.remote,
|
||||
torch.device("cuda"),
|
||||
pin_memory=True,
|
||||
)
|
||||
stager.load.assert_called_once_with()
|
||||
stager.offload.assert_called_once_with()
|
||||
|
||||
|
||||
def test_video_vae_encode_uses_configured_parallel_tiling():
|
||||
from vllm_omni.diffusion.models.minimax_h3.vae import (
|
||||
MiniMaxH3VideoVAE,
|
||||
@@ -373,3 +704,360 @@ def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
|
||||
torch.tensor([[1.0, 1.0], [2.0, 2.0]]),
|
||||
)
|
||||
assert shapes == [(1, 2, 3), (2, 2, 3)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("case", "extra", "image_count", "expected"),
|
||||
[
|
||||
("F1", {"frame_index": -1}, 1, [-1]),
|
||||
("F2", {"frame_indices": [0, -1]}, 2, [0, -1]),
|
||||
],
|
||||
)
|
||||
def test_f1_f2_official_fl2va_keyframe_matrix(case, extra, image_count, expected):
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_resolve_fl2va_keyframe_indices,
|
||||
)
|
||||
|
||||
assert case in {"F1", "F2"}
|
||||
assert _resolve_fl2va_keyframe_indices(extra, image_count) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("case", "counts"),
|
||||
[
|
||||
("R1", (3, 0, 0)),
|
||||
("R2", (1, 0, 0)),
|
||||
("R3", (1, 1, 0)),
|
||||
("R4", (0, 1, 1)),
|
||||
("R5", (1, 1, 1)),
|
||||
("R6", (1, 0, 2)),
|
||||
],
|
||||
)
|
||||
def test_r1_r6_ref2va_reference_count_matrix(case, counts):
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_validate_ref2va_reference_counts,
|
||||
)
|
||||
|
||||
assert case.startswith("R")
|
||||
_validate_ref2va_reference_counts(*counts)
|
||||
|
||||
|
||||
def test_ref2va_reference_count_validation_preserves_client_error_metadata():
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_validate_ref2va_reference_counts,
|
||||
)
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
with pytest.raises(OmniClientError, match="at least one image or video"):
|
||||
_validate_ref2va_reference_counts(0, 0, 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("case", "start_time", "expected_duration"),
|
||||
[
|
||||
("R7", None, 10.0),
|
||||
("R8", 4.0, 6.0),
|
||||
],
|
||||
)
|
||||
def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_time, expected_duration):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
|
||||
|
||||
source = tmp_path / f"{case}.mp4"
|
||||
source.touch()
|
||||
metadata = {
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"fps": 24.0,
|
||||
"frame_count": 240,
|
||||
"duration": 10.0,
|
||||
"format_names": ("mp4",),
|
||||
"video_codec": "h264",
|
||||
"audio_codecs": (),
|
||||
"file_size": 1024,
|
||||
}
|
||||
transcode_calls = []
|
||||
monkeypatch.setattr(reference_video_module, "_probe_video", lambda _path: metadata)
|
||||
monkeypatch.setattr(
|
||||
reference_video_module,
|
||||
"_transcode_reference_video",
|
||||
lambda source, **kwargs: transcode_calls.append((source, kwargs)) or "prepared.mp4",
|
||||
)
|
||||
|
||||
prepared = reference_video_module.prepare_reference_videos(
|
||||
[str(source)],
|
||||
target_frame_count=209,
|
||||
workdir=str(tmp_path / "work"),
|
||||
start_time_seconds=start_time,
|
||||
)
|
||||
|
||||
assert prepared[0]["duration_seconds"] == pytest.approx(expected_duration)
|
||||
assert prepared[0]["start_time_seconds"] == pytest.approx(start_time or 0.0)
|
||||
assert transcode_calls[0][1]["duration_seconds"] == pytest.approx(expected_duration)
|
||||
|
||||
|
||||
def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
|
||||
|
||||
first = tmp_path / "first.mp4"
|
||||
second = tmp_path / "second.mp4"
|
||||
first.touch()
|
||||
second.touch()
|
||||
metadata = {
|
||||
str(first): {
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"fps": 24.0,
|
||||
"frame_count": 241,
|
||||
"duration": 10.041667,
|
||||
"format_names": ("mp4",),
|
||||
"video_codec": "h264",
|
||||
"audio_codecs": (),
|
||||
"file_size": 1024,
|
||||
},
|
||||
str(second): {
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"fps": 24.0,
|
||||
"frame_count": 120,
|
||||
"duration": 4.966667,
|
||||
"format_names": ("mp4",),
|
||||
"video_codec": "h264",
|
||||
"audio_codecs": (),
|
||||
"file_size": 1024,
|
||||
},
|
||||
}
|
||||
transcode_calls = []
|
||||
monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)])
|
||||
monkeypatch.setattr(
|
||||
reference_video_module,
|
||||
"_transcode_reference_video",
|
||||
lambda source, **kwargs: transcode_calls.append((source, kwargs)) or f"{source}.prepared.mp4",
|
||||
)
|
||||
|
||||
prepared = reference_video_module.prepare_reference_videos(
|
||||
[str(first), str(second)],
|
||||
target_frame_count=124,
|
||||
workdir=str(tmp_path / "work"),
|
||||
)
|
||||
|
||||
assert [item["duration_seconds"] for item in prepared] == pytest.approx([10.041667, 4.958333])
|
||||
assert sum(item["duration_seconds"] for item in prepared) == pytest.approx(15.0)
|
||||
assert transcode_calls[1][1]["duration_seconds"] == pytest.approx(4.958333)
|
||||
|
||||
|
||||
def test_ref2va_two_video_recipe_rejects_real_duration_overflow(monkeypatch, tmp_path):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
first = tmp_path / "first.mp4"
|
||||
second = tmp_path / "second.mp4"
|
||||
first.touch()
|
||||
second.touch()
|
||||
metadata = {
|
||||
str(first): {
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"fps": 24.0,
|
||||
"frame_count": 240,
|
||||
"duration": 10.0,
|
||||
"format_names": ("mp4",),
|
||||
"video_codec": "h264",
|
||||
"audio_codecs": (),
|
||||
"file_size": 1024,
|
||||
},
|
||||
str(second): {
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"fps": 24.0,
|
||||
"frame_count": 120,
|
||||
"duration": 5.02,
|
||||
"format_names": ("mp4",),
|
||||
"video_codec": "h264",
|
||||
"audio_codecs": (),
|
||||
"file_size": 1024,
|
||||
},
|
||||
}
|
||||
monkeypatch.setattr(reference_video_module, "_probe_video", lambda path: metadata[str(path)])
|
||||
monkeypatch.setattr(reference_video_module, "_transcode_reference_video", lambda source, **kwargs: "prepared.mp4")
|
||||
|
||||
with pytest.raises(OmniClientError, match="15 seconds"):
|
||||
reference_video_module.prepare_reference_videos(
|
||||
[str(first), str(second)],
|
||||
target_frame_count=124,
|
||||
workdir=str(tmp_path / "work"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("duration", [4.0, 15.0])
|
||||
def test_g2_output_duration_accepts_official_boundaries(duration):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
sampling = SimpleNamespace(
|
||||
fps=24,
|
||||
num_frames=1,
|
||||
height=None,
|
||||
width=None,
|
||||
extra_args={"duration": duration},
|
||||
)
|
||||
|
||||
height, width, *_ = pipeline._resolve_shape("ref2va", sampling, None)
|
||||
assert height > 0 and width > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("duration", [3.99, 15.01, float("nan"), "not-a-duration"])
|
||||
def test_g2_output_duration_rejects_out_of_contract_values(duration):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
sampling = SimpleNamespace(
|
||||
fps=24,
|
||||
num_frames=1,
|
||||
height=None,
|
||||
width=None,
|
||||
extra_args={"duration": duration},
|
||||
)
|
||||
with pytest.raises(ValueError, match="duration"):
|
||||
pipeline._resolve_shape("ref2va", sampling, None)
|
||||
|
||||
|
||||
def test_g1_fanout_uses_incrementing_output_seeds():
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_minimax_h3_output_seeds,
|
||||
_resolve_minimax_h3_num_outputs,
|
||||
)
|
||||
|
||||
assert _resolve_minimax_h3_num_outputs(3) == 3
|
||||
assert _minimax_h3_output_seeds(100, 3) == [100, 101, 102]
|
||||
with pytest.raises(ValueError, match="num_outputs_per_prompt"):
|
||||
_resolve_minimax_h3_num_outputs(11)
|
||||
|
||||
|
||||
def test_g3_task_specific_aspect_ratio_policy():
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_resolve_minimax_h3_aspect_ratio,
|
||||
)
|
||||
|
||||
image = Image.new("RGB", (1280, 720))
|
||||
assert _resolve_minimax_h3_aspect_ratio("fl2va", "9:16", image) == pytest.approx(16 / 9)
|
||||
assert _resolve_minimax_h3_aspect_ratio("ref2va", None, None) == pytest.approx(16 / 9)
|
||||
assert _resolve_minimax_h3_aspect_ratio("ref2va", "auto", None) == pytest.approx(16 / 9)
|
||||
with pytest.raises(ValueError, match="requires an explicit"):
|
||||
_resolve_minimax_h3_aspect_ratio("t2va", None, None)
|
||||
with pytest.raises(ValueError, match="one of"):
|
||||
_resolve_minimax_h3_aspect_ratio("t2va", "2:1", None)
|
||||
|
||||
|
||||
def test_g3_t2va_shape_requires_a_named_ratio_and_fl2va_ignores_override():
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
t2va_sampling = SimpleNamespace(
|
||||
fps=24,
|
||||
num_frames=1,
|
||||
height=None,
|
||||
width=None,
|
||||
extra_args={"duration": 5.0},
|
||||
)
|
||||
with pytest.raises(ValueError, match="requires an explicit aspect_ratio"):
|
||||
pipeline._resolve_shape("t2va", t2va_sampling, None)
|
||||
|
||||
fl2va_sampling = SimpleNamespace(
|
||||
fps=24,
|
||||
num_frames=1,
|
||||
height=None,
|
||||
width=None,
|
||||
extra_args={"duration": 5.0, "aspect_ratio": "9:16"},
|
||||
)
|
||||
height, width, *_ = pipeline._resolve_shape("fl2va", fl2va_sampling, Image.new("RGB", (1280, 720)))
|
||||
assert (height, width) == (768, 1344)
|
||||
|
||||
|
||||
def test_g4_reference_image_boundaries_and_aspect_ratio():
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_validate_reference_image,
|
||||
)
|
||||
|
||||
_validate_reference_image(Image.new("RGB", (256, 256)))
|
||||
_validate_reference_image(Image.new("RGB", (5760, 2304)))
|
||||
with pytest.raises(ValueError, match="dimensions"):
|
||||
_validate_reference_image(Image.new("RGB", (255, 255)))
|
||||
with pytest.raises(ValueError, match="aspect ratio"):
|
||||
_validate_reference_image(Image.new("RGB", (256, 641)))
|
||||
|
||||
|
||||
def test_g4_reference_image_file_format_and_size_contract(tmp_path):
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_load_image,
|
||||
)
|
||||
|
||||
valid_path = tmp_path / "reference.png"
|
||||
Image.new("RGB", (256, 256)).save(valid_path)
|
||||
assert _load_image(valid_path).size == (256, 256)
|
||||
|
||||
invalid_path = tmp_path / "reference.bmp"
|
||||
Image.new("RGB", (256, 256)).save(invalid_path)
|
||||
with pytest.raises(ValueError, match="must use"):
|
||||
_load_image(invalid_path)
|
||||
|
||||
|
||||
def test_g4_standalone_audio_duration_and_total_duration_contract():
|
||||
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
|
||||
validate_reference_audio_waveforms,
|
||||
)
|
||||
|
||||
sample_rate = 16000
|
||||
validate_reference_audio_waveforms(
|
||||
[
|
||||
(torch.zeros(1, 2 * sample_rate), sample_rate),
|
||||
(torch.zeros(1, 13 * sample_rate), sample_rate),
|
||||
]
|
||||
)
|
||||
with pytest.raises(ValueError, match="duration"):
|
||||
validate_reference_audio_waveforms([(torch.zeros(1, int(1.9 * sample_rate)), sample_rate)])
|
||||
with pytest.raises(ValueError, match="at most 15 seconds in total"):
|
||||
validate_reference_audio_waveforms(
|
||||
[
|
||||
(torch.zeros(1, 8 * sample_rate), sample_rate),
|
||||
(torch.zeros(1, 8 * sample_rate), sample_rate),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value", "message"),
|
||||
[
|
||||
("fps", 23.0, "FPS"),
|
||||
("duration", 1.0, "duration"),
|
||||
("format_names", ("avi",), "container"),
|
||||
("video_codec", "vp9", "H.264"),
|
||||
("audio_codecs", ("opus",), "AAC"),
|
||||
("file_size", 50 * 1024 * 1024 + 1, "size"),
|
||||
],
|
||||
)
|
||||
def test_g4_reference_video_metadata_validation(field, value, message, tmp_path):
|
||||
from vllm_omni.diffusion.models.minimax_h3.reference_video import (
|
||||
_validate_reference_video_metadata,
|
||||
)
|
||||
|
||||
metadata = {
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"fps": 24.0,
|
||||
"duration": 10.0,
|
||||
"format_names": ("mp4",),
|
||||
"video_codec": "h264",
|
||||
"audio_codecs": ("aac",),
|
||||
"file_size": 1024,
|
||||
}
|
||||
metadata[field] = value
|
||||
with pytest.raises(ValueError, match=message):
|
||||
_validate_reference_video_metadata(metadata, index=0, source=str(tmp_path / "reference.mp4"))
|
||||
@@ -16,6 +16,17 @@ pytestmark = [
|
||||
MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
|
||||
|
||||
|
||||
def _assert_joint_output(outputs):
|
||||
assert len(outputs) == 1
|
||||
frames = np.asarray(outputs[0].images[0])
|
||||
assert frames.shape == (39, 256, 448, 3)
|
||||
multimodal = outputs[0].multimodal_output
|
||||
assert multimodal is not None
|
||||
assert np.asarray(multimodal["audio"]).shape[1] == 2
|
||||
assert multimodal["audio_sample_rate"] == 32000
|
||||
assert multimodal["fps"] == 24
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get(MODEL_ENV),
|
||||
reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
|
||||
@@ -50,6 +61,8 @@ def test_minimax_h3_t2va_ulysses8_smoke():
|
||||
output_type="np",
|
||||
extra_args={
|
||||
"task": "t2va",
|
||||
"duration": 4.0,
|
||||
"aspect_ratio": "16:9",
|
||||
"flow_shift": 12.0,
|
||||
"audio_flow_shift": 3.0,
|
||||
},
|
||||
@@ -61,9 +74,61 @@ def test_minimax_h3_t2va_ulysses8_smoke():
|
||||
|
||||
assert len(outputs) == 1
|
||||
frames = np.asarray(outputs[0].images[0])
|
||||
assert frames.shape == (39, 256, 448, 3)
|
||||
assert frames.shape == (107, 256, 448, 3)
|
||||
multimodal = outputs[0].multimodal_output
|
||||
assert multimodal is not None
|
||||
assert np.asarray(multimodal["audio"]).shape[1] == 2
|
||||
assert multimodal["audio_sample_rate"] == 32000
|
||||
assert multimodal["fps"] == 24
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.environ.get(MODEL_ENV),
|
||||
reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path",
|
||||
)
|
||||
def test_minimax_h3_t2va_fp8_single_gpu_smoke():
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
if torch.accelerator.device_count() < 1:
|
||||
pytest.skip("MiniMax H3 FP8 smoke requires a CUDA device")
|
||||
|
||||
engine = Omni(
|
||||
model=os.environ[MODEL_ENV],
|
||||
quantization_config={
|
||||
"transformer": {"method": "fp8"},
|
||||
"text_encoder": None,
|
||||
"video_vae": None,
|
||||
"audio_vae": None,
|
||||
},
|
||||
parallel_config=DiffusionParallelConfig(ulysses_degree=1),
|
||||
trust_remote_code=True,
|
||||
enable_cpu_offload=True,
|
||||
enforce_eager=True,
|
||||
)
|
||||
try:
|
||||
outputs = engine.generate(
|
||||
"A quiet cinematic night scene with matching ambient sound.",
|
||||
OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=448,
|
||||
num_frames=29,
|
||||
fps=24,
|
||||
num_inference_steps=2,
|
||||
seed=42,
|
||||
output_type="np",
|
||||
extra_args={
|
||||
"task": "t2va",
|
||||
"flow_shift": 12.0,
|
||||
"audio_flow_shift": 3.0,
|
||||
},
|
||||
),
|
||||
use_tqdm=False,
|
||||
)
|
||||
finally:
|
||||
engine.close()
|
||||
|
||||
_assert_joint_output(outputs)
|
||||
@@ -144,3 +144,20 @@ def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one():
|
||||
noise_aug=0.25,
|
||||
)
|
||||
torch.testing.assert_close(first, second)
|
||||
|
||||
|
||||
def test_condition_noise_accepts_a_reference_video_longer_than_the_target():
|
||||
from vllm_omni.diffusion.models.minimax_h3.condition_noise import (
|
||||
minimax_h3_imgvid_cond_noise_aug_rows,
|
||||
)
|
||||
|
||||
rows = torch.zeros(4 * 2 * 2, 96)
|
||||
result = minimax_h3_imgvid_cond_noise_aug_rows(
|
||||
rows,
|
||||
condition_shapes=[(4, 4, 4)],
|
||||
target_latent_t=2,
|
||||
imgvid_cond_num_frames=1,
|
||||
seed=7,
|
||||
noise_aug=0.5,
|
||||
)
|
||||
assert result.shape == rows.shape
|
||||
@@ -0,0 +1,213 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class _FakeLinear(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
bias=False,
|
||||
params_dtype=None,
|
||||
quant_config=None,
|
||||
prefix="",
|
||||
total_num_heads=None,
|
||||
total_num_kv_heads=None,
|
||||
**kwargs,
|
||||
):
|
||||
del args, kwargs
|
||||
super().__init__()
|
||||
self.prefix = prefix
|
||||
self.quant_config = quant_config
|
||||
self.weight = nn.Parameter(torch.empty(1, dtype=params_dtype))
|
||||
if bias:
|
||||
self.bias = nn.Parameter(torch.empty(1, dtype=params_dtype))
|
||||
else:
|
||||
self.register_parameter("bias", None)
|
||||
self.num_heads = total_num_heads
|
||||
self.num_kv_heads = total_num_kv_heads
|
||||
|
||||
|
||||
class _FakeAttention(nn.Module):
|
||||
def __init__(self, **kwargs):
|
||||
del kwargs
|
||||
super().__init__()
|
||||
|
||||
|
||||
def _small_od_config():
|
||||
arch = {
|
||||
"num_layers": 1,
|
||||
"token_refiner_num_layers": 1,
|
||||
"hidden_size": 8,
|
||||
"num_attention_heads": 2,
|
||||
"attention_head_dim": 4,
|
||||
"ffn_hidden_size": 16,
|
||||
"latents_dim": 2,
|
||||
"audio_latents_dim": 2,
|
||||
"patch_size": (1, 2, 2),
|
||||
"text_dim": 6,
|
||||
"timestep_input_dim": 4,
|
||||
"time_embed_hidden_size": 8,
|
||||
"time_embed_dim": 4,
|
||||
"adaln_out_features": 18 * 8,
|
||||
"final_adaln_out_features": 2 * 8,
|
||||
"rope_inv_freq_len": 2,
|
||||
}
|
||||
return SimpleNamespace(
|
||||
tf_model_config=arch,
|
||||
parallel_config=SimpleNamespace(ulysses_degree=1),
|
||||
)
|
||||
|
||||
|
||||
def test_fp8_scope_and_prefix_propagation(monkeypatch):
|
||||
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
|
||||
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
is_layer_skipped,
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer as h3
|
||||
|
||||
monkeypatch.setattr(h3, "ColumnParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "MergedColumnParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "QKVParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "RowParallelLinear", _FakeLinear)
|
||||
monkeypatch.setattr(h3, "Attention", _FakeAttention)
|
||||
monkeypatch.setattr(h3, "get_tensor_model_parallel_world_size", lambda: 1)
|
||||
|
||||
ignored_layers = {
|
||||
"token_refiner.blocks.0.mlp.fc2",
|
||||
"blocks.0.attn.qkv_proj",
|
||||
"condition_proj",
|
||||
"blocks.0.adaln_proj.linear",
|
||||
"final_layer.adaln_proj.linear",
|
||||
}
|
||||
fp8_config = Fp8Config(ignored_layers=sorted(ignored_layers))
|
||||
model = h3.MiniMaxH3DiTModel(
|
||||
_small_od_config(),
|
||||
quant_config=fp8_config,
|
||||
)
|
||||
linears = {module.prefix: module.quant_config for module in model.modules() if isinstance(module, _FakeLinear)}
|
||||
|
||||
quantized = {
|
||||
"token_refiner.blocks.0.attn.qkv_proj",
|
||||
"token_refiner.blocks.0.attn.out_proj",
|
||||
"token_refiner.blocks.0.mlp.fc1",
|
||||
"token_refiner.blocks.0.mlp.fc2",
|
||||
"blocks.0.attn.qkv_proj",
|
||||
"blocks.0.attn.out_proj",
|
||||
"blocks.0.mlp.fc1",
|
||||
"blocks.0.mlp.fc2",
|
||||
"condition_proj",
|
||||
"blocks.0.adaln_proj.linear",
|
||||
"final_layer.adaln_proj.linear",
|
||||
}
|
||||
full_precision = {
|
||||
"video_patch_proj",
|
||||
"audio_patch_proj",
|
||||
"time_embedder.proj_in",
|
||||
"time_embedder.proj_out",
|
||||
"final_layer.video_out",
|
||||
"final_layer.audio_out",
|
||||
}
|
||||
|
||||
assert quantized <= linears.keys()
|
||||
assert full_precision <= linears.keys()
|
||||
assert all(linears[prefix] is fp8_config for prefix in quantized)
|
||||
assert all(linears[prefix] is None for prefix in full_precision)
|
||||
assert {prefix for prefix in quantized if is_layer_skipped(prefix, fp8_config.ignored_layers)} == ignored_layers
|
||||
linear = Mock(spec=LinearBase)
|
||||
assert all(
|
||||
isinstance(
|
||||
fp8_config.get_quant_method(linear, prefix),
|
||||
UnquantizedLinearMethod,
|
||||
)
|
||||
for prefix in ignored_layers
|
||||
)
|
||||
|
||||
|
||||
class _WeightTarget(nn.Module):
|
||||
def __init__(self, loader):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty(1))
|
||||
self.weight.weight_loader = loader
|
||||
|
||||
|
||||
def test_model_load_weights_transforms_before_calling_vllm_loader():
|
||||
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
||||
MiniMaxH3DiTArchConfig,
|
||||
MiniMaxH3DiTModel,
|
||||
)
|
||||
|
||||
qkv_calls = []
|
||||
fc1_calls = []
|
||||
|
||||
def qkv_loader(param, loaded_weight):
|
||||
del param
|
||||
qkv_calls.append(loaded_weight.clone())
|
||||
|
||||
def fc1_loader(param, loaded_weight, shard_id):
|
||||
del param
|
||||
fc1_calls.append((shard_id, loaded_weight.clone()))
|
||||
|
||||
model = object.__new__(MiniMaxH3DiTModel)
|
||||
nn.Module.__init__(model)
|
||||
model.arch = MiniMaxH3DiTArchConfig(
|
||||
hidden_size=1,
|
||||
num_attention_heads=2,
|
||||
attention_head_dim=1,
|
||||
ffn_hidden_size=2,
|
||||
)
|
||||
model.blocks = nn.ModuleList([nn.Module()])
|
||||
model.blocks[0].attn = nn.Module()
|
||||
model.blocks[0].attn.qkv_proj = _WeightTarget(qkv_loader)
|
||||
model.blocks[0].mlp = nn.Module()
|
||||
model.blocks[0].mlp.fc1 = _WeightTarget(fc1_loader)
|
||||
|
||||
qkv = torch.arange(6, dtype=torch.float32).reshape(6, 1)
|
||||
fc1 = torch.arange(4, dtype=torch.float32).reshape(4, 1)
|
||||
loaded = model.load_weights(
|
||||
[
|
||||
("blocks.0.attn.qkv_proj.weight", qkv),
|
||||
("blocks.0.mlp.fc1.weight", fc1),
|
||||
]
|
||||
)
|
||||
|
||||
assert loaded == {
|
||||
"blocks.0.attn.qkv_proj.weight",
|
||||
"blocks.0.mlp.fc1.weight",
|
||||
}
|
||||
assert qkv_calls[0][:, 0].tolist() == [0, 3, 1, 4, 2, 5]
|
||||
assert [(shard_id, tensor[:, 0].tolist()) for shard_id, tensor in fc1_calls] == [
|
||||
(0, [0, 1]),
|
||||
(1, [2, 3]),
|
||||
]
|
||||
|
||||
|
||||
def test_pipeline_resolves_transformer_component_quant_config():
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_resolve_component_quant_config,
|
||||
)
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
ignored_layers = ["blocks.0.attn.qkv_proj"]
|
||||
component_config = build_quant_config(
|
||||
{
|
||||
"transformer": {
|
||||
"method": "fp8",
|
||||
"ignored_layers": ignored_layers,
|
||||
}
|
||||
}
|
||||
)
|
||||
transformer_config = component_config.resolve("transformer")
|
||||
|
||||
assert transformer_config.ignored_layers == ignored_layers
|
||||
assert _resolve_component_quant_config(component_config, "transformer") is transformer_config
|
||||
assert _resolve_component_quant_config(transformer_config, "transformer") is transformer_config
|
||||
@@ -0,0 +1,297 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.diffusion.quantization.test_quantization_quality import (
|
||||
_OUTPUT_DIR,
|
||||
QualityTestConfig,
|
||||
_compute_lpips,
|
||||
_compute_psnr_and_mae,
|
||||
_free_gpu_memory,
|
||||
_maybe_save_output,
|
||||
)
|
||||
from tests.helpers.mark import hardware_marks
|
||||
|
||||
_MINIMAX_H3_REPO = "MiniMaxAI/MiniMax-H3"
|
||||
_MINIMAX_H3_REVISION = "48d93ede732756e404a3b1b2f3b3a9b5a22f6cfc"
|
||||
_MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000
|
||||
_MIN_AUDIO_SPECTRAL_COSINE = 0.80
|
||||
_MIN_AUDIO_RMS_RATIO = 0.50
|
||||
_MAX_AUDIO_RMS_RATIO = 2.00
|
||||
|
||||
_QUALITY_CONFIG = QualityTestConfig(
|
||||
id="fp8_minimax_h3_fl2va",
|
||||
model=_MINIMAX_H3_REPO,
|
||||
quantization={"transformer": {"method": "fp8"}},
|
||||
task="t2v",
|
||||
prompt=(
|
||||
"A bright daytime cinematic tracking shot of a golden retriever running through a sunlit meadow "
|
||||
"filled with colorful wildflowers, vivid blue sky, crisp natural colors."
|
||||
),
|
||||
max_lpips=0.20,
|
||||
height=384,
|
||||
width=672,
|
||||
num_frames=107,
|
||||
num_inference_steps=10,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_fl2va_model_ref() -> str:
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
repo_root = Path(
|
||||
snapshot_download(
|
||||
repo_id=_MINIMAX_H3_REPO,
|
||||
revision=_MINIMAX_H3_REVISION,
|
||||
allow_patterns=["FL2VA/**"],
|
||||
)
|
||||
)
|
||||
fl2va_root = repo_root / "FL2VA"
|
||||
if not (fl2va_root / "model_index.json").is_file():
|
||||
raise FileNotFoundError(f"MiniMax H3 FL2VA model_index.json not found under {repo_root}")
|
||||
return str(fl2va_root)
|
||||
|
||||
|
||||
def _canonical_audio(audio: object) -> np.ndarray:
|
||||
"""Return H3 audio as float32 [channels, samples]."""
|
||||
array = np.asarray(audio, dtype=np.float32)
|
||||
while array.ndim > 2 and array.shape[0] == 1:
|
||||
array = array[0]
|
||||
if array.ndim == 1:
|
||||
array = array[None, :]
|
||||
if array.ndim != 2:
|
||||
raise ValueError(f"Expected mono/stereo audio, got shape {array.shape}")
|
||||
if array.shape[0] > array.shape[1] and array.shape[1] <= 8:
|
||||
array = array.T
|
||||
if not 1 <= array.shape[0] <= 8:
|
||||
raise ValueError(f"Expected channel-first audio, got shape {array.shape}")
|
||||
if not np.isfinite(array).all():
|
||||
raise ValueError("Audio contains non-finite samples")
|
||||
return np.ascontiguousarray(array)
|
||||
|
||||
|
||||
def _audio_quality_metrics(baseline: object, quantized: object) -> tuple[float, float]:
|
||||
"""Compare phase-tolerant spectral structure and overall signal energy."""
|
||||
baseline_array = _canonical_audio(baseline)
|
||||
quantized_array = _canonical_audio(quantized)
|
||||
if baseline_array.shape != quantized_array.shape:
|
||||
raise ValueError(
|
||||
"Audio shapes do not match for metric computation: "
|
||||
f"baseline={baseline_array.shape}, quantized={quantized_array.shape}"
|
||||
)
|
||||
|
||||
baseline_tensor = torch.from_numpy(baseline_array)
|
||||
quantized_tensor = torch.from_numpy(quantized_array)
|
||||
similarities: list[float] = []
|
||||
for n_fft in (512, 1024, 2048):
|
||||
window = torch.hann_window(n_fft)
|
||||
baseline_magnitude = torch.stft(
|
||||
baseline_tensor,
|
||||
n_fft=n_fft,
|
||||
hop_length=n_fft // 4,
|
||||
window=window,
|
||||
return_complex=True,
|
||||
).abs()
|
||||
quantized_magnitude = torch.stft(
|
||||
quantized_tensor,
|
||||
n_fft=n_fft,
|
||||
hop_length=n_fft // 4,
|
||||
window=window,
|
||||
return_complex=True,
|
||||
).abs()
|
||||
cosine = torch.nn.functional.cosine_similarity(
|
||||
baseline_magnitude.flatten(1),
|
||||
quantized_magnitude.flatten(1),
|
||||
dim=1,
|
||||
eps=1e-12,
|
||||
)
|
||||
similarities.extend(float(value) for value in cosine)
|
||||
|
||||
baseline_rms = float(torch.sqrt(torch.mean(baseline_tensor.square())))
|
||||
quantized_rms = float(torch.sqrt(torch.mean(quantized_tensor.square())))
|
||||
rms_ratio = quantized_rms / max(baseline_rms, 1e-12)
|
||||
return float(np.mean(similarities)), rms_ratio
|
||||
|
||||
|
||||
def _generate_joint_output(omni, config: QualityTestConfig):
|
||||
"""Generate H3 video/audio, returning both payloads and peak GPU memory."""
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(config.seed)
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
outputs = omni.generate(
|
||||
{"prompt": config.prompt, "negative_prompt": config.negative_prompt},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=config.height,
|
||||
width=config.width,
|
||||
generator=generator,
|
||||
num_inference_steps=config.num_inference_steps,
|
||||
num_frames=config.num_frames,
|
||||
guidance_scale=config.guidance_scale,
|
||||
sigmas=config.sigmas,
|
||||
),
|
||||
)
|
||||
main_process_peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
|
||||
|
||||
first = outputs[0]
|
||||
worker_peak_memory_mb = float(getattr(first, "peak_memory_mb", 0.0) or 0.0)
|
||||
if hasattr(first, "request_output") and isinstance(first.request_output, list):
|
||||
first = first.request_output[0]
|
||||
worker_peak_memory_mb = max(
|
||||
worker_peak_memory_mb,
|
||||
float(getattr(first, "peak_memory_mb", 0.0) or 0.0),
|
||||
)
|
||||
if not hasattr(first, "images") or not first.images:
|
||||
raise ValueError("Could not extract video frames from H3 output")
|
||||
frames = first.images[0]
|
||||
if isinstance(frames, torch.Tensor):
|
||||
frames = frames.detach().cpu()
|
||||
if frames.dim() == 5:
|
||||
frames = frames[0]
|
||||
if frames.dim() == 4 and frames.shape[0] in (3, 4):
|
||||
frames = frames.permute(1, 2, 3, 0)
|
||||
if frames.is_floating_point():
|
||||
frames = frames.clamp(-1, 1) * 0.5 + 0.5
|
||||
video = frames.float().numpy()
|
||||
else:
|
||||
video = np.asarray(frames)
|
||||
if video.ndim == 5:
|
||||
video = video[0]
|
||||
|
||||
multimodal = getattr(first, "multimodal_output", None)
|
||||
if not multimodal or "audio" not in multimodal:
|
||||
raise ValueError("Could not extract audio from H3 output")
|
||||
sample_rate = int(multimodal.get("audio_sample_rate", 0))
|
||||
peak_mem = worker_peak_memory_mb / 1024 if worker_peak_memory_mb > 0 else main_process_peak_mem
|
||||
return video, _canonical_audio(multimodal["audio"]), sample_rate, peak_mem
|
||||
|
||||
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"config",
|
||||
[
|
||||
pytest.param(
|
||||
_QUALITY_CONFIG,
|
||||
id=_QUALITY_CONFIG.id,
|
||||
marks=hardware_marks(res={"cuda": _QUALITY_CONFIG.gpu}, num_cards=2),
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_minimax_h3_quantization_quality(config: QualityTestConfig):
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
model_ref = _resolve_fl2va_model_ref()
|
||||
common_kwargs = {
|
||||
"model": model_ref,
|
||||
"enforce_eager": True,
|
||||
"tensor_parallel_size": 2,
|
||||
"text_encoder_tp_size": 2,
|
||||
"vae_use_tiling": True,
|
||||
}
|
||||
|
||||
omni_bl = Omni(**common_kwargs)
|
||||
baseline_out, baseline_audio, baseline_sample_rate, bl_mem = _generate_joint_output(omni_bl, config)
|
||||
omni_bl.shutdown()
|
||||
del omni_bl
|
||||
_free_gpu_memory()
|
||||
_maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out)
|
||||
|
||||
quantization = config.quantization_ref()
|
||||
omni_qt = Omni(**common_kwargs, quantization_config=quantization)
|
||||
quant_out, quant_audio, quant_sample_rate, qt_mem = _generate_joint_output(omni_qt, config)
|
||||
omni_qt.shutdown()
|
||||
del omni_qt
|
||||
_free_gpu_memory()
|
||||
_maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out)
|
||||
|
||||
lpips_score = _compute_lpips(baseline_out, quant_out, config.task)
|
||||
psnr_score, mae_score = _compute_psnr_and_mae(baseline_out, quant_out, config.task)
|
||||
if baseline_sample_rate != _MINIMAX_H3_AUDIO_SAMPLE_RATE or quant_sample_rate != baseline_sample_rate:
|
||||
raise AssertionError(
|
||||
f"Unexpected H3 audio sample rate: baseline={baseline_sample_rate}, quantized={quant_sample_rate}"
|
||||
)
|
||||
audio_spectral_cosine, audio_rms_ratio = _audio_quality_metrics(baseline_audio, quant_audio)
|
||||
assert lpips_score <= config.max_lpips, (
|
||||
f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} "
|
||||
f"for {config.quantization_ref()} on {config.quantized_ref()}"
|
||||
)
|
||||
assert audio_spectral_cosine >= _MIN_AUDIO_SPECTRAL_COSINE, (
|
||||
f"Audio spectral cosine {audio_spectral_cosine:.4f} is below threshold "
|
||||
f"{_MIN_AUDIO_SPECTRAL_COSINE} for {config.id}"
|
||||
)
|
||||
assert _MIN_AUDIO_RMS_RATIO <= audio_rms_ratio <= _MAX_AUDIO_RMS_RATIO, (
|
||||
f"Audio RMS ratio {audio_rms_ratio:.4f} is outside "
|
||||
f"[{_MIN_AUDIO_RMS_RATIO}, {_MAX_AUDIO_RMS_RATIO}] for {config.id}"
|
||||
)
|
||||
|
||||
mem_reduction = (bl_mem - qt_mem) / bl_mem * 100 if bl_mem > 0 else 0
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"Quantization Quality: {config.id}")
|
||||
print(f"{'=' * 60}")
|
||||
print(f" Baseline: {config.baseline_ref()}")
|
||||
print(f" Quantized: {config.quantized_ref()}")
|
||||
print(f" Method: {config.quantization_ref()}")
|
||||
print(f" LPIPS: {lpips_score:.4f} (threshold: {config.max_lpips})")
|
||||
print(f" PSNR: {psnr_score:.4f} dB (higher is better)")
|
||||
print(f" MAE: {mae_score:.6f} (lower is better)")
|
||||
print(f" Audio cosine: {audio_spectral_cosine:.4f} (threshold: {_MIN_AUDIO_SPECTRAL_COSINE})")
|
||||
print(f" Audio RMS: {audio_rms_ratio:.4f}x (range: {_MIN_AUDIO_RMS_RATIO}-{_MAX_AUDIO_RMS_RATIO})")
|
||||
print(f" BF16 memory: {bl_mem:.2f} GiB")
|
||||
print(f" Quant memory: {qt_mem:.2f} GiB ({mem_reduction:.0f}% reduction)")
|
||||
print(" Result: PASS")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
assert np.isfinite(psnr_score) or np.isinf(psnr_score), f"PSNR is invalid for {config.id}: {psnr_score}"
|
||||
assert np.isfinite(mae_score), f"MAE is not finite for {config.id}: {mae_score}"
|
||||
|
||||
|
||||
def test_resolve_fl2va_model_ref(tmp_path, monkeypatch):
|
||||
fl2va_root = tmp_path / "FL2VA"
|
||||
fl2va_root.mkdir()
|
||||
(fl2va_root / "model_index.json").write_text("{}", encoding="utf-8")
|
||||
|
||||
def fake_snapshot_download(*, repo_id, revision, allow_patterns):
|
||||
assert repo_id == _MINIMAX_H3_REPO
|
||||
assert revision == _MINIMAX_H3_REVISION
|
||||
assert allow_patterns == ["FL2VA/**"]
|
||||
return str(tmp_path)
|
||||
|
||||
monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download)
|
||||
assert _resolve_fl2va_model_ref() == str(fl2va_root)
|
||||
|
||||
|
||||
def test_audio_quality_metrics_ignore_phase_but_reject_spectral_drift():
|
||||
sample_rate = _MINIMAX_H3_AUDIO_SAMPLE_RATE
|
||||
time = np.arange(sample_rate, dtype=np.float32) / sample_rate
|
||||
baseline = np.stack(
|
||||
[
|
||||
np.sin(2 * np.pi * 440 * time),
|
||||
np.sin(2 * np.pi * 880 * time),
|
||||
]
|
||||
)
|
||||
phase_shifted = np.stack(
|
||||
[
|
||||
np.sin(2 * np.pi * 440 * time + np.pi / 2),
|
||||
np.sin(2 * np.pi * 880 * time + np.pi / 2),
|
||||
]
|
||||
)
|
||||
different = np.stack(
|
||||
[
|
||||
np.sin(2 * np.pi * 3000 * time),
|
||||
np.sin(2 * np.pi * 4000 * time),
|
||||
]
|
||||
)
|
||||
|
||||
phase_cosine, phase_rms_ratio = _audio_quality_metrics(baseline, phase_shifted)
|
||||
different_cosine, _ = _audio_quality_metrics(baseline, different)
|
||||
|
||||
assert phase_cosine >= _MIN_AUDIO_SPECTRAL_COSINE
|
||||
assert phase_rms_ratio == pytest.approx(1.0, abs=1e-5)
|
||||
assert different_cosine < _MIN_AUDIO_SPECTRAL_COSINE
|
||||
@@ -27,7 +27,9 @@ from vllm_omni.diffusion.offloader.block_discovery import (
|
||||
from vllm_omni.diffusion.offloader.distributed_layerwise_backend import (
|
||||
DistributedLayerwiseOffloadBackend,
|
||||
DistributedLayerwiseOffloadHook,
|
||||
PinnedResidentLayerGroup,
|
||||
)
|
||||
from vllm_omni.diffusion.offloader.module_residency import PinnedModuleStager
|
||||
from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan, get_offload_plan
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
@@ -97,6 +99,7 @@ def patched_offload_runtime(monkeypatch):
|
||||
monkeypatch.setattr(dist_backend_module.current_omni_platform, "Event", DummyEvent)
|
||||
monkeypatch.setattr(dist_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
|
||||
monkeypatch.setattr(dist_backend_module.current_omni_platform, "stream", dummy_stream)
|
||||
monkeypatch.setattr(dist_backend_module.current_omni_platform, "synchronize", lambda: None)
|
||||
|
||||
|
||||
class TinyBlock(nn.Module):
|
||||
@@ -308,6 +311,164 @@ class TestDistributedLayerwiseOffloadHook:
|
||||
assert shard1.numel() == 2
|
||||
assert torch.equal(shard1, torch.tensor([102.0, 0.0]))
|
||||
|
||||
def test_mmap_layout_transform_is_applied_before_sharding(self, dist_group, patched_offload_runtime):
|
||||
block = nn.Module()
|
||||
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
|
||||
next_block = nn.Module()
|
||||
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
|
||||
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
|
||||
next_block.weight.mmap_weight_transform_pending = True
|
||||
|
||||
hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=next_block,
|
||||
device=torch.device("cpu"),
|
||||
dp_group=None,
|
||||
dp_size=2,
|
||||
rank=0,
|
||||
pin_memory=False,
|
||||
)
|
||||
hook.initialize_hook(block)
|
||||
|
||||
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([3.0, 2.0]))
|
||||
|
||||
def test_regular_loader_weight_does_not_reapply_mmap_transform(self, dist_group, patched_offload_runtime):
|
||||
block = nn.Module()
|
||||
block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
|
||||
next_block = nn.Module()
|
||||
next_block.weight = nn.Parameter(torch.arange(4, dtype=torch.float32))
|
||||
next_block.weight.mmap_weight_transform = lambda tensor: tensor.flip(0)
|
||||
|
||||
hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=next_block,
|
||||
device=torch.device("cpu"),
|
||||
dp_group=None,
|
||||
dp_size=2,
|
||||
rank=0,
|
||||
pin_memory=False,
|
||||
)
|
||||
hook.initialize_hook(block)
|
||||
|
||||
assert torch.equal(hook.cpu_shards[torch.float32], torch.tensor([0.0, 1.0]))
|
||||
|
||||
|
||||
class TestPinnedResidentLayerGroup:
|
||||
def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
|
||||
blocks = [nn.Linear(2, 2), nn.Linear(2, 2)]
|
||||
expected = []
|
||||
for index, block in enumerate(blocks):
|
||||
block.weight.data.fill_(index + 1)
|
||||
block.bias.data.fill_(index + 3)
|
||||
expected.append((block.weight.detach().clone(), block.bias.detach().clone()))
|
||||
|
||||
group = PinnedResidentLayerGroup(
|
||||
blocks,
|
||||
device=torch.device("cpu"),
|
||||
copy_stream=DummyStream(),
|
||||
pin_memory=False,
|
||||
)
|
||||
|
||||
assert all(block.weight.numel() == 0 for block in blocks)
|
||||
|
||||
group.load()
|
||||
for block, (weight, bias) in zip(blocks, expected):
|
||||
assert torch.equal(block.weight, weight)
|
||||
assert torch.equal(block.bias, bias)
|
||||
|
||||
# Device-side changes must not overwrite the persistent host master.
|
||||
blocks[0].weight.data.zero_()
|
||||
group.offload()
|
||||
assert all(block.weight.numel() == 0 for block in blocks)
|
||||
|
||||
group.load()
|
||||
for block, (weight, bias) in zip(blocks, expected):
|
||||
assert torch.equal(block.weight, weight)
|
||||
assert torch.equal(block.bias, bias)
|
||||
|
||||
group.offload()
|
||||
|
||||
def test_mmap_loader_attrs_survive_to_empty_parameter_replacement(self):
|
||||
module = nn.Linear(2, 2, bias=False)
|
||||
|
||||
def transform(tensor):
|
||||
return tensor.flip(0)
|
||||
|
||||
module.weight.mmap_weight_transform = transform
|
||||
|
||||
backend = object.__new__(DistributedLayerwiseOffloadBackend)
|
||||
backend._remember_mmap_param_attrs(module)
|
||||
module.to_empty(device="meta")
|
||||
|
||||
replacement = nn.Parameter(torch.arange(4).reshape(2, 2).float())
|
||||
backend._attach_mmap_param_attrs("weight", replacement, module.weight)
|
||||
|
||||
assert replacement.mmap_weight_transform is transform
|
||||
assert replacement.mmap_weight_transform_pending is True
|
||||
|
||||
def test_shared_allgather_output_is_narrowed_to_current_block(
|
||||
self, dist_group, patched_offload_runtime, monkeypatch
|
||||
):
|
||||
block = nn.Module()
|
||||
block.weight = nn.Parameter(torch.zeros(3, dtype=torch.float32))
|
||||
next_block = nn.Module()
|
||||
next_block.weight = nn.Parameter(torch.arange(3, dtype=torch.float32))
|
||||
|
||||
hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=next_block,
|
||||
device=torch.device("cpu"),
|
||||
dp_group=object(),
|
||||
dp_size=2,
|
||||
rank=0,
|
||||
pin_memory=False,
|
||||
)
|
||||
hook.initialize_hook(block)
|
||||
# Simulate unified buffers sized for a larger heterogeneous block.
|
||||
hook.gpu_buffers = [
|
||||
{torch.float32: torch.empty(12)},
|
||||
{torch.float32: torch.empty(12)},
|
||||
]
|
||||
hook.gpu_shard_buffers = [
|
||||
{torch.float32: torch.empty(6)},
|
||||
{torch.float32: torch.empty(6)},
|
||||
]
|
||||
output_sizes = []
|
||||
|
||||
def fake_allgather(output, local_shard, *, group):
|
||||
del group
|
||||
output_sizes.append(output.numel())
|
||||
output.zero_()
|
||||
output[: local_shard.numel()].copy_(local_shard)
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
|
||||
|
||||
hook.prefetch_layer(slot=0, non_blocking=False)
|
||||
|
||||
assert output_sizes == [4]
|
||||
|
||||
|
||||
class TestPinnedModuleStager:
|
||||
def test_offload_rebinds_immutable_cpu_master(self, patched_offload_runtime):
|
||||
module = nn.Linear(2, 2)
|
||||
expected_weight = module.weight.detach().clone()
|
||||
expected_bias = module.bias.detach().clone()
|
||||
stager = PinnedModuleStager(
|
||||
module,
|
||||
torch.device("cpu"),
|
||||
pin_memory=False,
|
||||
copy_stream=DummyStream(),
|
||||
)
|
||||
|
||||
stager.load()
|
||||
module.weight.data.zero_()
|
||||
module.bias.data.zero_()
|
||||
stager.offload()
|
||||
|
||||
assert torch.equal(module.weight, expected_weight)
|
||||
assert torch.equal(module.bias, expected_bias)
|
||||
stager.load()
|
||||
assert torch.equal(module.weight, expected_weight)
|
||||
assert torch.equal(module.bias, expected_bias)
|
||||
stager.offload()
|
||||
|
||||
|
||||
class _DummyBlock(nn.Module):
|
||||
def __init__(self):
|
||||
@@ -504,6 +665,81 @@ class TestCrossGroupSharedBuffer:
|
||||
assert not sync_called[0], "Should skip sync-prefetch when slot_group matches (non-contaminated)"
|
||||
|
||||
|
||||
class TestAllGatherOutputSlicing:
|
||||
"""Regression test: with heterogeneous block sizes (e.g. MiniMax H3's
|
||||
token_refiner with blocks of ~1231MB and ~239MB), the shared AllGather
|
||||
output buffers are sized to the *largest* block across all groups.
|
||||
prefetch_layer must slice the output buffer down to each block's actual
|
||||
AllGather output size, otherwise output.numel() != dp_size * input.numel()
|
||||
and all_gather_into_tensor raises during enable()'s first prefetch.
|
||||
"""
|
||||
|
||||
def test_prefetch_slices_shared_output_buffer_to_block_size(self, dist_group, patched_offload_runtime, monkeypatch):
|
||||
dp_size = 2
|
||||
|
||||
# Heterogeneous block sizes: 8 vs 32 elements.
|
||||
small_block = nn.Module()
|
||||
small_block.weight = nn.Parameter(torch.arange(8, dtype=torch.float32))
|
||||
large_block = nn.Module()
|
||||
large_block.weight = nn.Parameter(torch.arange(100, 132, dtype=torch.float32))
|
||||
|
||||
def make_hook(next_block: nn.Module) -> DistributedLayerwiseOffloadHook:
|
||||
hook = DistributedLayerwiseOffloadHook(
|
||||
next_block=next_block,
|
||||
device=torch.device("cpu"),
|
||||
dp_group=dist.group.WORLD, # non-None -> AllGather path (mocked below)
|
||||
dp_size=dp_size,
|
||||
rank=0,
|
||||
pin_memory=False,
|
||||
)
|
||||
current_block = nn.Module()
|
||||
current_block.weight = nn.Parameter(torch.zeros(4, dtype=torch.float32))
|
||||
hook.initialize_hook(current_block)
|
||||
return hook
|
||||
|
||||
hook_small = make_hook(small_block)
|
||||
hook_large = make_hook(large_block)
|
||||
hooks = [hook_small, hook_large]
|
||||
|
||||
# Assign max-sized shared buffers exactly as backend.enable() does.
|
||||
shared_buffers = DistributedLayerwiseOffloadBackend._allocate_shared_buffers(hooks)
|
||||
shared_shard_buffers = DistributedLayerwiseOffloadBackend._allocate_shared_shard_buffers(hooks)
|
||||
for hook in hooks:
|
||||
hook.gpu_buffers = shared_buffers
|
||||
hook.gpu_shard_buffers = shared_shard_buffers
|
||||
|
||||
# Shared output buffer must be sized to the largest block (32).
|
||||
assert shared_buffers[0][torch.float32].numel() == 32
|
||||
|
||||
# Fake AllGather that enforces the real torch size contract and
|
||||
# emulates reconstruction of the full block from all ranks' shards.
|
||||
full_values = {
|
||||
8: torch.arange(8, dtype=torch.float32),
|
||||
32: torch.arange(100, 132, dtype=torch.float32),
|
||||
}
|
||||
ag_output_numels: list[int] = []
|
||||
|
||||
def fake_all_gather(output, input, group=None):
|
||||
assert output.numel() == dp_size * input.numel(), (
|
||||
f"AllGather contract violated: output.numel()={output.numel()} "
|
||||
f"!= {dp_size} * input.numel()={input.numel()}"
|
||||
)
|
||||
ag_output_numels.append(output.numel())
|
||||
output.copy_(full_values[output.numel()])
|
||||
|
||||
monkeypatch.setattr(dist, "all_gather_into_tensor", fake_all_gather)
|
||||
|
||||
# Small block: output must be sliced to dp_size * shard = 8, not 32.
|
||||
hook_small.prefetch_layer(slot=0, non_blocking=False)
|
||||
assert ag_output_numels[-1] == 8
|
||||
assert torch.equal(small_block.weight, full_values[8])
|
||||
|
||||
# Large block: uses the full max-sized buffer.
|
||||
hook_large.prefetch_layer(slot=1, non_blocking=False)
|
||||
assert ag_output_numels[-1] == 32
|
||||
assert torch.equal(large_block.weight, full_values[32])
|
||||
|
||||
|
||||
class TestOffloadPlan:
|
||||
"""Test declarative OffloadPlan metadata."""
|
||||
|
||||
@@ -517,6 +753,7 @@ class TestOffloadPlan:
|
||||
plan = OffloadPlan(
|
||||
block_attrs={"transformer": ("gen_layers",)},
|
||||
offload_submodules={"context_encoder": "layers"},
|
||||
resident_dit_paths=frozenset({"transformer"}),
|
||||
)
|
||||
|
||||
class PipelineWithPlan(nn.Module):
|
||||
@@ -527,12 +764,14 @@ class TestOffloadPlan:
|
||||
assert result is plan
|
||||
assert result.block_attrs == {"transformer": ("gen_layers",)}
|
||||
assert result.offload_submodules == {"context_encoder": "layers"}
|
||||
assert result.resident_dit_paths == frozenset({"transformer"})
|
||||
|
||||
def test_offload_plan_defaults_to_empty(self):
|
||||
"""OffloadPlan with no arguments should have empty dicts."""
|
||||
plan = OffloadPlan()
|
||||
assert plan.block_attrs == {}
|
||||
assert plan.offload_submodules == {}
|
||||
assert plan.resident_dit_paths == frozenset()
|
||||
|
||||
def test_offload_plan_is_frozen(self):
|
||||
"""OffloadPlan should be immutable (frozen=True)."""
|
||||
@@ -540,6 +779,50 @@ class TestOffloadPlan:
|
||||
with pytest.raises(Exception):
|
||||
plan.block_attrs = {"x": ("y",)} # type: ignore
|
||||
|
||||
def test_all_resident_dit_still_prepares_planned_submodules(self, patched_offload_runtime):
|
||||
class TokenRefiner(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
|
||||
|
||||
class Transformer(nn.Module):
|
||||
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.blocks = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)])
|
||||
self.token_refiner = TokenRefiner()
|
||||
|
||||
class Pipeline(nn.Module):
|
||||
_offload_plan = OffloadPlan(
|
||||
offload_submodules={"token_refiner": "blocks"},
|
||||
resident_dit_paths=frozenset({"transformer"}),
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = Transformer()
|
||||
|
||||
pipeline = Pipeline()
|
||||
backend = DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=False,
|
||||
dlo_use_allgather=False,
|
||||
dlo_resident_layers=2,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
)
|
||||
|
||||
backend.enable(pipeline)
|
||||
|
||||
assert len(backend._resident_blocks) == 2
|
||||
assert len(backend._all_hook_groups) == 1
|
||||
assert len(backend._all_hook_groups[0]) == 2
|
||||
assert backend.enabled
|
||||
|
||||
backend.disable()
|
||||
|
||||
|
||||
class TestMmapValidation:
|
||||
"""Tests for mmap loader validation: TP rejection, strict check, validate_loaded_weights."""
|
||||
@@ -746,7 +1029,7 @@ class TestConfigValidation:
|
||||
OffloadConfig.from_od_config(FakeODConfig())
|
||||
|
||||
def test_hsdp_without_allgather_allowed(self):
|
||||
"""HSDP + DLO + no-AllGather should be allowed (full weights per rank)."""
|
||||
"""HSDP + DLO + no-AllGather uses standard-loader rank-local weights."""
|
||||
from vllm_omni.diffusion.offloader.base import OffloadConfig
|
||||
|
||||
class FakePC:
|
||||
@@ -768,12 +1051,33 @@ class TestConfigValidation:
|
||||
enable_layerwise_offload = False
|
||||
enable_distributed_layerwise_offload = True
|
||||
dlo_use_allgather = False # no AllGather → should be allowed
|
||||
dlo_resident_layers = 20
|
||||
pin_cpu_memory = True
|
||||
parallel_config = FakePC()
|
||||
model = "/fake/path"
|
||||
|
||||
config = OffloadConfig.from_od_config(FakeODConfig())
|
||||
assert config.dp_size == 1 # forced to 1 when no AllGather
|
||||
assert config.dlo_resident_layers == 20
|
||||
|
||||
def test_resident_layers_with_allgather_rejected(self):
|
||||
class FakePC:
|
||||
data_parallel_size = 2
|
||||
use_hsdp = False
|
||||
sequence_parallel_size = 1
|
||||
|
||||
class FakeODConfig:
|
||||
enable_cpu_offload = False
|
||||
enable_layerwise_offload = False
|
||||
enable_distributed_layerwise_offload = True
|
||||
dlo_use_allgather = True
|
||||
dlo_resident_layers = 20
|
||||
pin_cpu_memory = True
|
||||
parallel_config = FakePC()
|
||||
model = "/fake/path"
|
||||
|
||||
with pytest.raises(ValueError, match="requires --dlo-no-use-allgather"):
|
||||
OffloadConfig.from_od_config(FakeODConfig())
|
||||
|
||||
def test_num_inference_steps_none_rejected(self):
|
||||
"""DP multi-concurrency should reject None num_inference_steps."""
|
||||
|
||||
@@ -247,6 +247,56 @@ class TestRequestModeDispatch:
|
||||
results = {ro.request_id: ro.result.error for ro in out.runner_outputs}
|
||||
assert results == {request_id: f"result_for_{request_id}" for request_id in request_ids}
|
||||
|
||||
def test_dlo_dp_routes_multiple_requests_without_pipeline_batching(self):
|
||||
executor, _, _ = _make_executor(num_gpus=2)
|
||||
executor.od_config = SimpleNamespace(
|
||||
step_execution=False,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True,
|
||||
)
|
||||
executor.execute_request = Mock(return_value="dlo-dp")
|
||||
executor.collective_rpc = Mock()
|
||||
scheduler_output = _make_sched_output("A", "B")
|
||||
|
||||
result = executor.execute_batch(scheduler_output)
|
||||
|
||||
assert result == "dlo-dp"
|
||||
executor.execute_request.assert_called_once_with(scheduler_output)
|
||||
executor.collective_rpc.assert_not_called()
|
||||
|
||||
def test_dlo_dp_multi_rank_reply_uses_synchronous_rpc_collection(self):
|
||||
executor, req_q, res_q = _make_executor(num_gpus=2)
|
||||
executor.od_config = SimpleNamespace(
|
||||
step_execution=False,
|
||||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||||
)
|
||||
executor._sync_result_buffer = res_q
|
||||
|
||||
def worker():
|
||||
request = req_q.get(timeout=2)
|
||||
assert "rpc_id" not in request
|
||||
wave_id = request["wave_id"]
|
||||
for dp_rank in range(2):
|
||||
res_q.put(
|
||||
{
|
||||
"dp_rank": dp_rank,
|
||||
"output": _tagged_output(str(dp_rank)),
|
||||
"wave_id": wave_id,
|
||||
}
|
||||
)
|
||||
|
||||
thread = threading.Thread(target=worker, daemon=True)
|
||||
thread.start()
|
||||
results = executor.collective_rpc(
|
||||
"execute_model",
|
||||
unique_reply_rank=None,
|
||||
exec_all_ranks=True,
|
||||
)
|
||||
thread.join(timeout=2)
|
||||
|
||||
assert [result.error for result in results] == ["0", "1"]
|
||||
|
||||
|
||||
# ───────────────── concurrent collective RPC ─────────────────
|
||||
|
||||
|
||||
@@ -390,6 +390,105 @@ def test_send_single_request_respects_processor_receiver_boundary(build_adapter,
|
||||
assert sent_payload.meta.is_segment_finished.item() is False
|
||||
|
||||
|
||||
def test_send_single_request_personaplex_pending_frame_is_not_segment_boundary(
|
||||
build_adapter,
|
||||
):
|
||||
from vllm_omni.model_executor.stage_input_processors.personaplex import (
|
||||
talker2code2wav_async_chunk,
|
||||
)
|
||||
|
||||
adapter, connector = build_adapter(
|
||||
stage_id=0,
|
||||
connector_extra={
|
||||
"initial_codec_chunk_frames": 1,
|
||||
"codec_chunk_frames": 5,
|
||||
},
|
||||
)
|
||||
request = _req(
|
||||
"req-personaplex",
|
||||
RequestStatus.WAITING,
|
||||
external_req_id="ext-personaplex",
|
||||
)
|
||||
request.resumable = True
|
||||
request.additional_information = {
|
||||
"codes": {
|
||||
"audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
|
||||
}
|
||||
}
|
||||
adapter.custom_process_next_stage_input_func = talker2code2wav_async_chunk
|
||||
|
||||
adapter._send_single_request(
|
||||
{
|
||||
"multimodal_output": None,
|
||||
"request": request,
|
||||
"is_finished": False,
|
||||
"is_segment_finished": True,
|
||||
}
|
||||
)
|
||||
|
||||
sent_payload = connector.put.call_args.kwargs["data"]
|
||||
assert sent_payload.codes is None
|
||||
assert sent_payload.meta.finished.item() is False
|
||||
assert sent_payload.meta.is_segment_finished.item() is False
|
||||
|
||||
|
||||
def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter):
|
||||
from vllm_omni.model_executor.stage_input_processors.personaplex import (
|
||||
talker2code2wav_async_chunk,
|
||||
)
|
||||
|
||||
adapter, _ = build_adapter(
|
||||
stage_id=0,
|
||||
connector_extra={
|
||||
"initial_codec_chunk_frames": 1,
|
||||
"codec_chunk_frames": 5,
|
||||
},
|
||||
)
|
||||
request = _req(
|
||||
"req-personaplex-first",
|
||||
RequestStatus.WAITING,
|
||||
external_req_id="ext-personaplex-reused",
|
||||
)
|
||||
request.resumable = True
|
||||
request.additional_information = {
|
||||
"codes": {
|
||||
"audio": torch.arange(8, dtype=torch.long).reshape(1, 8),
|
||||
}
|
||||
}
|
||||
|
||||
first = talker2code2wav_async_chunk(
|
||||
adapter,
|
||||
multimodal_output=None,
|
||||
request=request,
|
||||
is_finished=True,
|
||||
)
|
||||
assert first is not None
|
||||
assert first.codes is None
|
||||
|
||||
adapter.cleanup_sender(request.external_req_id)
|
||||
|
||||
replacement = _req(
|
||||
"req-personaplex-replacement",
|
||||
RequestStatus.WAITING,
|
||||
external_req_id=request.external_req_id,
|
||||
)
|
||||
replacement.resumable = True
|
||||
replacement.additional_information = {
|
||||
"codes": {
|
||||
"audio": torch.arange(8, 16, dtype=torch.long).reshape(1, 8),
|
||||
}
|
||||
}
|
||||
second = talker2code2wav_async_chunk(
|
||||
adapter,
|
||||
multimodal_output=None,
|
||||
request=replacement,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert second is not None
|
||||
assert second.codes is None
|
||||
|
||||
|
||||
def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1)
|
||||
request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream")
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""End-to-end accuracy tests."""
|
||||
@@ -0,0 +1 @@
|
||||
"""MiniMax H3 accuracy tests."""
|
||||
@@ -0,0 +1,216 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
assert_video_metadata,
|
||||
assert_video_similarity_metrics,
|
||||
probe_binary,
|
||||
probe_video,
|
||||
reset_artifact_dir,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServer
|
||||
|
||||
pytestmark = [pytest.mark.benchmark, pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL_REPO_ID = "MiniMaxAI/MiniMax-H3"
|
||||
# The previously pinned revision was removed from the Hugging Face repository.
|
||||
MODEL_REVISION = "main"
|
||||
MODEL_ENV_VAR = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL"
|
||||
# Keep the official asset link visible in the test report. Hugging Face's
|
||||
# ``blob`` URL is an HTML page, so resolve it to the raw asset for download.
|
||||
REFERENCE_VIDEO_URL = "https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/assets/t2va.mp4"
|
||||
REFERENCE_VIDEO_SHA256 = "d66903241362e224085cc93f7a5e70fba6ab378d0ac6ff6af83acf2559849a42"
|
||||
|
||||
# This is the official reproducible 768p T2VA prompt associated with
|
||||
# ``assets/t2va.mp4`` in the MiniMax H3 repository.
|
||||
PROMPT = """integrated_multimodal_description: [Shot 1] Cinematic, medium wide shot, pushing in slowly. In the cavernous, dimly lit bridge of a starship, sleek metallic consoles with glowing amber displays flank a massive, curved observation window. A female captain, in her late 40s with an athletic build and short silver-streaked black hair, stands in the center midground. She wears a structured, high-collared dark navy military tunic with silver chest insignias. Her back is to the camera, silhouetted against the cool, ambient starlight pouring through the thick glass. She stands perfectly still with her hands clasped tightly behind her back. Outside the window, a massive armada of jagged, dark grey dreadnoughts hovers in tight formation against a deep purple space nebula. The fleet's massive rear thrusters begin to glow with an intense, escalating bright blue light. [Shot 2] At 00:04.500, the camera cuts to a close-up of the captain's face and shakes strongly. The brilliant blue-white light from the fleet's gathering energy reflects vividly in her dark eyes. Suddenly, a blinding white flash floods through the window, completely washing out the background as the fleet jumps to hyperspace. The sheer spatial force violently jolts the bridge, causing the captain from Shot 1 to stagger slightly forward, her shoulders tensing as she visibly braces herself against the physical tremors. As the intense white light fades abruptly, leaving only the dim, empty expanse of the purple nebula reflected on her starkly lit skin, her jaw clenches, and she slowly closes her eyes in the newly emptied space.
|
||||
overall_soundscape: A low, resonant hum of the ship's ambient life support systems serves as the baseline, soon drowned out by an audible, escalating, high-pitched electronic whine as the fleet outside charges its hyperdrives. A massive, deafening, bass-heavy boom and sharp crackle erupts during the blinding flash, accompanied by the loud metallic creaking, rattling, and deep thuds of the bridge's bulkheads vibrating under immense physical stress. The intense roaring impact then cuts abruptly back to a hollow, echoing room tone, leaving only the faint, steady hum of the isolated bridge.
|
||||
non_diegetic_music: Cinematic space-opera orchestral score, slow tempo, featuring a solitary, mournful French horn melody over deep, sustained string dissonances that build rapidly in volume and intensity, swelling to a massive orchestral peak before snapping immediately into silence right after the jump."""
|
||||
|
||||
WIDTH = 1344
|
||||
HEIGHT = 768
|
||||
FPS = 24
|
||||
NUM_FRAMES = 243
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
FLOW_SHIFT = 12.0
|
||||
AUDIO_FLOW_SHIFT = 3.0
|
||||
DURATION_SECONDS = 10.0
|
||||
SEED = 0
|
||||
SSIM_THRESHOLD = 0.82
|
||||
PSNR_THRESHOLD = 20.0
|
||||
REQUEST_TIMEOUT_SECONDS = 60 * 60
|
||||
|
||||
|
||||
def _model_name() -> str:
|
||||
configured = os.environ.get(MODEL_ENV_VAR)
|
||||
if configured:
|
||||
return configured
|
||||
|
||||
snapshot_root = snapshot_download(
|
||||
repo_id=MODEL_REPO_ID,
|
||||
revision=MODEL_REVISION,
|
||||
allow_patterns=["FL2VA/**"],
|
||||
)
|
||||
return str(Path(snapshot_root) / "FL2VA")
|
||||
|
||||
|
||||
def _server_args() -> list[str]:
|
||||
return [
|
||||
"--trust-remote-code",
|
||||
"--num-gpus",
|
||||
"4",
|
||||
"--usp",
|
||||
"4",
|
||||
"--ring",
|
||||
"1",
|
||||
"--use-hsdp",
|
||||
"--hsdp-shard-size",
|
||||
"4",
|
||||
"--text-encoder-tp-size",
|
||||
"4",
|
||||
"--vae-patch-parallel-size",
|
||||
"4",
|
||||
"--vae-parallel-mode",
|
||||
"tile",
|
||||
"--vae-use-tiling",
|
||||
"--diffusion-attention-backend",
|
||||
"FLASH_ATTN",
|
||||
"--stage-init-timeout",
|
||||
"1800",
|
||||
"--init-timeout",
|
||||
"1800",
|
||||
]
|
||||
|
||||
|
||||
def _download_reference_video(output_dir: Path) -> Path:
|
||||
reference_path = output_dir / "reference.mp4"
|
||||
download_url = REFERENCE_VIDEO_URL.replace("/blob/", "/resolve/", 1)
|
||||
response = requests.get(download_url, timeout=120)
|
||||
response.raise_for_status()
|
||||
content = response.content
|
||||
digest = hashlib.sha256(content).hexdigest()
|
||||
assert digest == REFERENCE_VIDEO_SHA256, (
|
||||
f"Official MiniMax H3 T2VA reference changed: got {digest}, expected {REFERENCE_VIDEO_SHA256}"
|
||||
)
|
||||
reference_path.write_bytes(content)
|
||||
return reference_path
|
||||
|
||||
|
||||
def _probe_audio(path: Path) -> dict[str, str | int]:
|
||||
probe_binary("ffprobe")
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-select_streams",
|
||||
"a:0",
|
||||
"-show_entries",
|
||||
"stream=codec_name,sample_rate,channels",
|
||||
"-of",
|
||||
"json",
|
||||
str(path),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
streams = json.loads(result.stdout).get("streams", [])
|
||||
assert len(streams) == 1, f"Expected one audio stream in {path}, got {streams}"
|
||||
stream = streams[0]
|
||||
return {
|
||||
"codec_name": str(stream["codec_name"]),
|
||||
"sample_rate": int(stream["sample_rate"]),
|
||||
"channels": int(stream["channels"]),
|
||||
}
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
def test_minimax_h3_t2va_matches_official_reference(
|
||||
accuracy_artifact_root: Path,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("MiniMax H3 T2VA accuracy test requires CUDA.")
|
||||
|
||||
probe_binary("ffmpeg")
|
||||
probe_binary("ffprobe")
|
||||
output_dir = reset_artifact_dir(accuracy_artifact_root / "minimax_h3_t2va")
|
||||
reference_path = _download_reference_video(output_dir)
|
||||
generated_path = output_dir / "vllm_omni.mp4"
|
||||
|
||||
server_env = {
|
||||
"FLASHINFER_DISABLE_VERSION_CHECK": "1",
|
||||
"VLLM_WORKER_MULTIPROC_METHOD": "spawn",
|
||||
"VLLM_OMNI_VIDEO_SYNC_TIMEOUT": str(REQUEST_TIMEOUT_SECONDS),
|
||||
"VLLM_OMNI_STORAGE_PATH": str(output_dir / "storage"),
|
||||
}
|
||||
request_data = {
|
||||
"prompt": PROMPT,
|
||||
"width": str(WIDTH),
|
||||
"height": str(HEIGHT),
|
||||
"fps": str(FPS),
|
||||
"num_inference_steps": str(NUM_INFERENCE_STEPS),
|
||||
"flow_shift": str(FLOW_SHIFT),
|
||||
"seed": str(SEED),
|
||||
"extra_params": json.dumps(
|
||||
{
|
||||
"task": "t2va",
|
||||
"duration": DURATION_SECONDS,
|
||||
"aspect_ratio": "16:9",
|
||||
"audio_flow_shift": AUDIO_FLOW_SHIFT,
|
||||
}
|
||||
),
|
||||
}
|
||||
|
||||
with OmniServer(
|
||||
_model_name(),
|
||||
_server_args(),
|
||||
env_dict=server_env,
|
||||
use_omni=True,
|
||||
) as server:
|
||||
response = requests.post(
|
||||
f"http://{server.host}:{server.port}/v1/videos/sync",
|
||||
data=request_data,
|
||||
timeout=REQUEST_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
assert response.headers["content-type"].startswith("video/mp4")
|
||||
generated_path.write_bytes(response.content)
|
||||
|
||||
reference_metadata = probe_video(reference_path)
|
||||
generated_metadata = probe_video(generated_path)
|
||||
assert reference_metadata == generated_metadata
|
||||
assert_video_metadata(
|
||||
generated_metadata,
|
||||
width=WIDTH,
|
||||
height=HEIGHT,
|
||||
fps=FPS,
|
||||
frame_count=NUM_FRAMES,
|
||||
)
|
||||
|
||||
expected_audio = {"codec_name": "aac", "sample_rate": 32000, "channels": 2}
|
||||
assert _probe_audio(reference_path) == expected_audio
|
||||
assert _probe_audio(generated_path) == expected_audio
|
||||
|
||||
assert_video_similarity_metrics(
|
||||
label="minimax_h3_t2va_official_reference",
|
||||
online_path=generated_path,
|
||||
offline_path=reference_path,
|
||||
ssim_threshold=SSIM_THRESHOLD,
|
||||
psnr_threshold=PSNR_THRESHOLD,
|
||||
)
|
||||
@@ -25,9 +25,15 @@ from comfyui_vllm_omni.nodes import (
|
||||
VLLMOmniGenerateVideo,
|
||||
VLLMOmniTTS,
|
||||
VLLMOmniUnderstanding,
|
||||
VLLMOmniVideoReferences,
|
||||
VLLMOmniVoiceClone,
|
||||
)
|
||||
from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams, WanModelSpecificParams
|
||||
from comfyui_vllm_omni.utils.types import (
|
||||
AutoregressionSamplingParams,
|
||||
DiffusionSamplingParams,
|
||||
MiniMaxH3ModelSpecificParams,
|
||||
WanModelSpecificParams,
|
||||
)
|
||||
from PIL import Image
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm import SamplingParams
|
||||
@@ -68,6 +74,8 @@ class SamplingKind(str, Enum):
|
||||
TTS_DIFFUSION_SINGLE = "tts_diffusion_single"
|
||||
VIDEO_NONE = "video_none"
|
||||
VIDEO_DIFFUSION_SINGLE = "video_diffusion_single"
|
||||
VIDEO_REF2VA_IMAGE_AUDIO = "video_ref2va_image_audio"
|
||||
VIDEO_REF2VA_MULTI_VIDEO = "video_ref2va_multi_video"
|
||||
|
||||
|
||||
# Pre-defined arguments to be used in function calls during the tests
|
||||
@@ -127,6 +135,15 @@ AR_LIST_SAMPLING_PARAMS = [
|
||||
|
||||
VIDEO_MODEL_PARAMS = WanModelSpecificParams({"guidance_scale_2": 5.0, "boundary_ratio": 0.98, "flow_shift": 12.0})
|
||||
|
||||
# Matches MiniMax-H3 recipe Ref2VA defaults (task is auto-routed by the client).
|
||||
H3_MODEL_PARAMS = MiniMaxH3ModelSpecificParams(
|
||||
{
|
||||
"type": "minimax_h3",
|
||||
"audio_flow_shift": 3.0,
|
||||
"flow_shift": 12.0,
|
||||
}
|
||||
)
|
||||
|
||||
LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10}
|
||||
|
||||
|
||||
@@ -306,6 +323,7 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
|
||||
received_sampling_params_list: Sequence[OmniSamplingParams] | None = (
|
||||
args[2] if len(args) > 2 else kwargs.get("sampling_params_list")
|
||||
)
|
||||
prompt = args[0] if len(args) > 0 else kwargs.get("prompt")
|
||||
|
||||
assert received_sampling_params_list is not None, (
|
||||
"In the current codebase, the API layer always provides not-None sampling parameter list when calling `AsyncOmni.generate`"
|
||||
@@ -372,6 +390,59 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
|
||||
LORA_PARAMS,
|
||||
)
|
||||
_assert_model_param_values(received_sampling_params_list[0], VIDEO_MODEL_PARAMS)
|
||||
elif sampling_case.kind is SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO:
|
||||
assert len(received_sampling_params_list) == 1
|
||||
assert isinstance(prompt, dict)
|
||||
multi_modal_data = prompt.get("multi_modal_data")
|
||||
assert isinstance(multi_modal_data, dict)
|
||||
assert isinstance(multi_modal_data.get("image"), Image.Image)
|
||||
assert isinstance(multi_modal_data.get("audio"), str)
|
||||
assert "video" not in multi_modal_data
|
||||
_assert_sampling_param_values(
|
||||
received_sampling_params_list[0],
|
||||
{
|
||||
"width": VIDEO_WIDTH,
|
||||
"height": VIDEO_HEIGHT,
|
||||
"num_frames": VIDEO_NUM_FRAMES,
|
||||
"fps": VIDEO_FPS,
|
||||
},
|
||||
)
|
||||
_assert_model_param_values(
|
||||
received_sampling_params_list[0],
|
||||
{
|
||||
"flow_shift": 12.0,
|
||||
"task": "ref2va",
|
||||
"audio_flow_shift": 3.0,
|
||||
},
|
||||
)
|
||||
elif sampling_case.kind is SamplingKind.VIDEO_REF2VA_MULTI_VIDEO:
|
||||
assert len(received_sampling_params_list) == 1
|
||||
assert isinstance(prompt, dict)
|
||||
multi_modal_data = prompt.get("multi_modal_data")
|
||||
assert isinstance(multi_modal_data, dict)
|
||||
input_videos = multi_modal_data.get("video")
|
||||
assert isinstance(input_videos, list)
|
||||
assert len(input_videos) == 2
|
||||
assert all(isinstance(path, str) for path in input_videos)
|
||||
assert "image" not in multi_modal_data
|
||||
assert "audio" not in multi_modal_data
|
||||
_assert_sampling_param_values(
|
||||
received_sampling_params_list[0],
|
||||
{
|
||||
"width": VIDEO_WIDTH,
|
||||
"height": VIDEO_HEIGHT,
|
||||
"num_frames": VIDEO_NUM_FRAMES,
|
||||
"fps": VIDEO_FPS,
|
||||
},
|
||||
)
|
||||
_assert_model_param_values(
|
||||
received_sampling_params_list[0],
|
||||
{
|
||||
"flow_shift": 12.0,
|
||||
"task": "ref2va",
|
||||
"audio_flow_shift": 3.0,
|
||||
},
|
||||
)
|
||||
else:
|
||||
raise AssertionError(f"Unknown sampling case: {sampling_case.kind}")
|
||||
|
||||
@@ -775,7 +846,7 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
|
||||
"model_params": VIDEO_MODEL_PARAMS,
|
||||
}
|
||||
if image_input:
|
||||
kwargs["image"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
|
||||
kwargs["frame"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32)
|
||||
if sampling_case.sampling_params is not None:
|
||||
kwargs["sampling_params"] = sampling_case.sampling_params
|
||||
if sampling_case.lora:
|
||||
@@ -786,3 +857,79 @@ async def test_video_generation_node(api_server: str, model: str, image_input: b
|
||||
assert isinstance(result, tuple)
|
||||
assert len(result) == 1
|
||||
assert isinstance(result[0], VideoInput)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"server_case,sampling_case,ref_mode",
|
||||
[
|
||||
pytest.param(
|
||||
ServerCase(
|
||||
served_model="MiniMaxAI/MiniMax-H3",
|
||||
stage_list=["diffusion"],
|
||||
stage_configs=[{"stage_type": "diffusion"}],
|
||||
outputs=[_build_diffusion_video_output()],
|
||||
),
|
||||
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_IMAGE_AUDIO, sampling_params=None),
|
||||
"image_audio",
|
||||
id="ref2va-image-audio",
|
||||
),
|
||||
pytest.param(
|
||||
ServerCase(
|
||||
served_model="MiniMaxAI/MiniMax-H3",
|
||||
stage_list=["diffusion"],
|
||||
stage_configs=[{"stage_type": "diffusion"}],
|
||||
outputs=[_build_diffusion_video_output()],
|
||||
),
|
||||
SamplingCase(kind=SamplingKind.VIDEO_REF2VA_MULTI_VIDEO, sampling_params=None),
|
||||
"multi_video",
|
||||
id="ref2va-multi-video",
|
||||
),
|
||||
],
|
||||
indirect=["server_case", "sampling_case"],
|
||||
)
|
||||
async def test_video_generation_node_minimax_h3_ref2va(
|
||||
api_server: str,
|
||||
sampling_case: SamplingCase,
|
||||
ref_mode: str,
|
||||
):
|
||||
"""MiniMax-H3 Ref2VA paths matching recipes/MiniMaxAI/MiniMax-H3.md."""
|
||||
node = VLLMOmniGenerateVideo()
|
||||
refs_node = VLLMOmniVideoReferences()
|
||||
|
||||
if ref_mode == "image_audio":
|
||||
(references,) = refs_node.get_references(
|
||||
image_1=torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32),
|
||||
audio_1={"waveform": torch.zeros((1, 1, 24000), dtype=torch.float32), "sample_rate": 24000},
|
||||
)
|
||||
prompt = (
|
||||
"A white cat with black mustache and eyebrow markings sits on a beige couch, "
|
||||
"lip-syncing precisely to the complete reference audio."
|
||||
)
|
||||
elif ref_mode == "multi_video":
|
||||
(references,) = refs_node.get_references(
|
||||
video_1=VideoInput(b"subject-video"), # type: ignore[reportAbstractUsage]
|
||||
video_2=VideoInput(b"background-video"), # type: ignore[reportAbstractUsage]
|
||||
)
|
||||
prompt = (
|
||||
"Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2."
|
||||
)
|
||||
else:
|
||||
raise AssertionError(f"Unknown ref_mode: {ref_mode}")
|
||||
|
||||
result = await node.generate(
|
||||
url=api_server,
|
||||
model="MiniMaxAI/MiniMax-H3",
|
||||
prompt=prompt,
|
||||
negative_prompt="",
|
||||
width=VIDEO_WIDTH,
|
||||
height=VIDEO_HEIGHT,
|
||||
fps=VIDEO_FPS,
|
||||
num_frames=VIDEO_NUM_FRAMES,
|
||||
references=references,
|
||||
model_params=H3_MODEL_PARAMS,
|
||||
)
|
||||
|
||||
assert isinstance(result, tuple)
|
||||
assert len(result) == 1
|
||||
assert isinstance(result[0], VideoInput)
|
||||
@@ -0,0 +1,30 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
|
||||
def test_batched_runtime_import_does_not_require_optional_opus_dependency() -> None:
|
||||
script = """
|
||||
import builtins
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
def guarded_import(name, *args, **kwargs):
|
||||
if name == "sphn":
|
||||
raise ModuleNotFoundError("blocked optional dependency: sphn")
|
||||
return real_import(name, *args, **kwargs)
|
||||
|
||||
builtins.__import__ = guarded_import
|
||||
from vllm_omni.experimental.fullduplex.personaplex.serving.batched import BatchedSessionManager
|
||||
assert BatchedSessionManager is not None
|
||||
"""
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
assert result.returncode == 0, result.stderr
|
||||
@@ -0,0 +1,303 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.experimental.fullduplex.personaplex.policy import (
|
||||
SILENCE_TOKENS,
|
||||
SINE_TOKENS,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.personaplex.stage0 import (
|
||||
PersonaPlexStage0DuplexRuntime,
|
||||
)
|
||||
from vllm_omni.model_executor.models.personaplex.personaplex_talker import (
|
||||
PersonaPlexTalkerForConditionalGeneration,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _FakeCodec:
|
||||
def __init__(self) -> None:
|
||||
self.encode_calls = 0
|
||||
self.reset_calls = 0
|
||||
|
||||
def streaming_init(self, batch_size: int) -> None:
|
||||
assert batch_size == 1
|
||||
|
||||
def encode_frame(self, pcm: torch.Tensor) -> torch.Tensor:
|
||||
assert pcm.shape == (1, 1920)
|
||||
self.encode_calls += 1
|
||||
return torch.full((1, 8), self.encode_calls, dtype=torch.long)
|
||||
|
||||
def reset_streaming(self) -> None:
|
||||
self.reset_calls += 1
|
||||
|
||||
|
||||
class _FakeTalker:
|
||||
def __init__(self) -> None:
|
||||
self.device = torch.device("cpu")
|
||||
self.dtype = torch.float32
|
||||
self.frame_calls: list[dict[str, torch.Tensor | None]] = []
|
||||
|
||||
def _build_prefill_embed(
|
||||
self,
|
||||
tokens,
|
||||
offset,
|
||||
span,
|
||||
device,
|
||||
silence=None,
|
||||
user_sine=None,
|
||||
):
|
||||
del offset, silence, user_sine
|
||||
values = tokens[:span].to(device=device, dtype=torch.float32)
|
||||
return values[:, None].expand(-1, 4).contiguous()
|
||||
|
||||
def _build_frame_embed(
|
||||
self,
|
||||
text_token,
|
||||
last_agent,
|
||||
prev_agent,
|
||||
device,
|
||||
user_d0=None,
|
||||
user_d1=None,
|
||||
):
|
||||
self.frame_calls.append(
|
||||
{
|
||||
"text_token": text_token.clone(),
|
||||
"last_agent": None if last_agent is None else last_agent.clone(),
|
||||
"prev_agent": None if prev_agent is None else prev_agent.clone(),
|
||||
"user_d0": None if user_d0 is None else user_d0.clone(),
|
||||
"user_d1": None if user_d1 is None else user_d1.clone(),
|
||||
}
|
||||
)
|
||||
value = float(user_d0[0].item()) if user_d0 is not None else -1.0
|
||||
return torch.full((1, 4), value, device=device)
|
||||
|
||||
|
||||
def _duplex_info(*, seq: int, session_id: str = "session", incarnation: int = 1):
|
||||
pcm = np.zeros(1920, dtype="<f4")
|
||||
return {
|
||||
"data_plane": True,
|
||||
"session_id": session_id,
|
||||
"incarnation": incarnation,
|
||||
"epoch": 0,
|
||||
"seq": seq,
|
||||
"payload": {
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": 24000,
|
||||
"audio": base64.b64encode(pcm.tobytes()).decode("ascii"),
|
||||
},
|
||||
"runtime_config": {
|
||||
"personaplex_model_path": "/unused",
|
||||
"personaplex_voice_prompt": "NATF2.pt",
|
||||
"personaplex_persona": "Be concise.",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _runtime(*codecs: _FakeCodec) -> PersonaPlexStage0DuplexRuntime:
|
||||
voice_embeddings = torch.arange(8, dtype=torch.float32).reshape(2, 1, 1, 4)
|
||||
codec_args = {"codec": codecs[0]}
|
||||
if len(codecs) > 1:
|
||||
available = iter(codecs)
|
||||
codec_args = {
|
||||
"codec_factory": lambda: next(available),
|
||||
"max_sessions": len(codecs),
|
||||
}
|
||||
return PersonaPlexStage0DuplexRuntime(
|
||||
_FakeTalker(),
|
||||
model_path="/unused",
|
||||
device="cpu",
|
||||
**codec_args,
|
||||
tokenizer=lambda _text: [7, 8, 9],
|
||||
voice_loader=lambda _voice: {
|
||||
"embeddings": voice_embeddings,
|
||||
"cache": torch.zeros((1, 17, 4), dtype=torch.long),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_first_append_prepends_voice_and_persona_once() -> None:
|
||||
codec = _FakeCodec()
|
||||
runtime = _runtime(codec)
|
||||
|
||||
first = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
|
||||
second = runtime.prepare_append(_duplex_info(seq=2), prompt_len=18)
|
||||
|
||||
assert first.prefill_applied is True
|
||||
assert first.prompt_offset == 0
|
||||
assert first.user_codes.shape == (1, 8)
|
||||
assert first.inputs_embeds.shape == (18, 4)
|
||||
assert second.prefill_applied is False
|
||||
assert second.prompt_offset == 17
|
||||
assert second.user_codes.shape == (2, 8)
|
||||
assert second.inputs_embeds.shape == (1, 4)
|
||||
assert codec.encode_calls == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("append_count", "expected_tokens", "expected_provided"),
|
||||
[
|
||||
(
|
||||
1,
|
||||
[*SILENCE_TOKENS, 1, *SINE_TOKENS[1:]],
|
||||
[False, *([True] * 15)],
|
||||
),
|
||||
(
|
||||
2,
|
||||
[*SILENCE_TOKENS, 2, *([1] * 7)],
|
||||
[*([False] * 8), *([True] * 8)],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_prepare_append_preserves_native_depformer_teacher_forcing(
|
||||
append_count: int,
|
||||
expected_tokens: list[int],
|
||||
expected_provided: list[bool],
|
||||
) -> None:
|
||||
runtime = _runtime(_FakeCodec())
|
||||
|
||||
result = None
|
||||
for seq in range(1, append_count + 1):
|
||||
result = runtime.prepare_append(_duplex_info(seq=seq), prompt_len=18)
|
||||
|
||||
assert result is not None
|
||||
assert result.info_update["pplex_depformer_audio_tokens"].tolist() == expected_tokens
|
||||
assert result.info_update["pplex_depformer_audio_provided"].tolist() == expected_provided
|
||||
|
||||
|
||||
def test_frame_embed_uses_previous_effective_agent_frame() -> None:
|
||||
fake_talker = SimpleNamespace(
|
||||
config=SimpleNamespace(
|
||||
num_audio_codebooks=16,
|
||||
audio_vocab_size=2048,
|
||||
),
|
||||
input_embeddings=lambda stack: stack,
|
||||
)
|
||||
last_agent = torch.arange(10, 18)
|
||||
|
||||
embeds = PersonaPlexTalkerForConditionalGeneration._build_frame_embed(
|
||||
fake_talker,
|
||||
torch.tensor([3]),
|
||||
last_agent,
|
||||
torch.arange(8),
|
||||
torch.device("cpu"),
|
||||
)
|
||||
|
||||
assert embeds.shape == (1, 17)
|
||||
assert torch.equal(embeds[0, 1:9], last_agent)
|
||||
|
||||
|
||||
def test_repeated_append_identity_does_not_advance_codec() -> None:
|
||||
codec = _FakeCodec()
|
||||
runtime = _runtime(codec)
|
||||
|
||||
info = _duplex_info(seq=1)
|
||||
first = runtime.prepare_append(info, prompt_len=18, request_id="req")
|
||||
retry = runtime.prepare_append(info, prompt_len=18, request_id="req")
|
||||
|
||||
assert codec.encode_calls == 1
|
||||
assert torch.equal(first.inputs_embeds, retry.inputs_embeds)
|
||||
assert torch.equal(first.user_codes, retry.user_codes)
|
||||
|
||||
|
||||
def test_next_append_uses_prior_sample_and_causally_delayed_user_frame() -> None:
|
||||
runtime = _runtime(_FakeCodec())
|
||||
|
||||
runtime.prepare_append(_duplex_info(seq=1), prompt_len=18, request_id="req")
|
||||
first_agent = torch.arange(8, dtype=torch.long)
|
||||
runtime.record_sample(request_id="req", text_token=torch.tensor(101), agent_codes=first_agent)
|
||||
runtime.prepare_append(_duplex_info(seq=2), prompt_len=19, request_id="req")
|
||||
second_agent = torch.arange(10, 18, dtype=torch.long)
|
||||
runtime.record_sample(request_id="req", text_token=torch.tensor(102), agent_codes=second_agent)
|
||||
runtime.prepare_append(_duplex_info(seq=3), prompt_len=20, request_id="req")
|
||||
|
||||
silence = torch.tensor(SILENCE_TOKENS, dtype=torch.long)
|
||||
sine = torch.tensor(SINE_TOKENS, dtype=torch.long)
|
||||
first_call = runtime.stage_model.frame_calls[0]
|
||||
assert all(torch.equal(first_call[key], silence) for key in ("last_agent", "prev_agent"))
|
||||
assert all(torch.equal(first_call[key], sine) for key in ("user_d0", "user_d1"))
|
||||
|
||||
second_call = runtime.stage_model.frame_calls[1]
|
||||
assert second_call["text_token"].tolist() == [101]
|
||||
expected_first_effective = torch.cat([first_agent[:1], silence[1:]])
|
||||
assert torch.equal(second_call["last_agent"], expected_first_effective)
|
||||
assert torch.equal(second_call["prev_agent"], expected_first_effective)
|
||||
assert torch.equal(second_call["user_d0"], torch.full((8,), 1, dtype=torch.long))
|
||||
assert torch.equal(second_call["user_d1"], sine)
|
||||
|
||||
third_call = runtime.stage_model.frame_calls[2]
|
||||
assert third_call["text_token"].tolist() == [102]
|
||||
assert torch.equal(third_call["last_agent"], second_agent)
|
||||
assert torch.equal(third_call["prev_agent"], second_agent)
|
||||
assert torch.equal(third_call["user_d0"], torch.full((8,), 2, dtype=torch.long))
|
||||
assert torch.equal(third_call["user_d1"], torch.full((8,), 1, dtype=torch.long))
|
||||
|
||||
|
||||
def test_decoded_pcm_is_writable_for_torch_zero_copy() -> None:
|
||||
pcm = PersonaPlexStage0DuplexRuntime._decode_pcm(_duplex_info(seq=1)["payload"])
|
||||
|
||||
assert pcm.flags.writeable
|
||||
|
||||
|
||||
def _prepare_two_sessions(
|
||||
runtime: PersonaPlexStage0DuplexRuntime,
|
||||
) -> tuple[object, object, object, object]:
|
||||
first_1 = runtime.prepare_append(_duplex_info(seq=1), prompt_len=18)
|
||||
second_1 = runtime.prepare_append(
|
||||
_duplex_info(seq=1, session_id="other", incarnation=2),
|
||||
prompt_len=18,
|
||||
)
|
||||
first_2 = runtime.prepare_append(_duplex_info(seq=2), prompt_len=19)
|
||||
second_2 = runtime.prepare_append(
|
||||
_duplex_info(seq=2, session_id="other", incarnation=2),
|
||||
prompt_len=19,
|
||||
)
|
||||
return first_1, second_1, first_2, second_2
|
||||
|
||||
|
||||
def test_live_sessions_keep_independent_streaming_encoders() -> None:
|
||||
first_codec = _FakeCodec()
|
||||
second_codec = _FakeCodec()
|
||||
runtime = _runtime(first_codec, second_codec)
|
||||
|
||||
first_1, second_1, first_2, second_2 = _prepare_two_sessions(runtime)
|
||||
|
||||
assert first_1.user_codes[:, 0].tolist() == [1]
|
||||
assert second_1.user_codes[:, 0].tolist() == [1]
|
||||
assert first_2.user_codes[:, 0].tolist() == [1, 2]
|
||||
assert second_2.user_codes[:, 0].tolist() == [1, 2]
|
||||
|
||||
|
||||
def test_stage0_session_capacity_fails_before_codec_state_is_shared() -> None:
|
||||
runtime = _runtime(_FakeCodec(), _FakeCodec())
|
||||
_prepare_two_sessions(runtime)
|
||||
|
||||
with pytest.raises(RuntimeError, match="capacity 2"):
|
||||
runtime.prepare_append(
|
||||
_duplex_info(seq=1, session_id="third", incarnation=3),
|
||||
prompt_len=18,
|
||||
)
|
||||
|
||||
|
||||
def test_close_session_resets_and_reuses_released_codec() -> None:
|
||||
first_codec = _FakeCodec()
|
||||
second_codec = _FakeCodec()
|
||||
runtime = _runtime(first_codec, second_codec)
|
||||
_prepare_two_sessions(runtime)
|
||||
|
||||
runtime.close_session("session", 1)
|
||||
|
||||
assert first_codec.reset_calls == 1
|
||||
assert second_codec.reset_calls == 0
|
||||
replacement = runtime.prepare_append(
|
||||
_duplex_info(seq=1, session_id="replacement", incarnation=4),
|
||||
prompt_len=18,
|
||||
)
|
||||
assert replacement.user_codes[:, 0].tolist() == [3]
|
||||
@@ -0,0 +1,316 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.model_executor.models.personaplex import (
|
||||
personaplex_code2wav,
|
||||
personaplex_mimi,
|
||||
)
|
||||
from vllm_omni.model_executor.models.personaplex.personaplex_code2wav import (
|
||||
PersonaPlexCode2Wav,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _FakeStreamingMimi(nn.Module):
|
||||
def __init__(self, stream_value: int = 0) -> None:
|
||||
super().__init__()
|
||||
self.stream_value = stream_value
|
||||
self.decode_frame_calls = 0
|
||||
self.reset_calls = 0
|
||||
|
||||
def decode(self, codes: torch.Tensor, return_dict: bool = True):
|
||||
assert return_dict
|
||||
frames = int(codes.shape[-1])
|
||||
return SimpleNamespace(audio_values=torch.arange(frames * 4, dtype=torch.float32).reshape(1, 1, -1))
|
||||
|
||||
def streaming_init(self, batch_size: int) -> None:
|
||||
assert batch_size == 1
|
||||
|
||||
def decode_frame(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
assert codes.shape == (1, 2)
|
||||
self.decode_frame_calls += 1
|
||||
value = self.stream_value + self.decode_frame_calls
|
||||
return torch.full((1, 4), float(value), dtype=torch.float32)
|
||||
|
||||
def reset_streaming(self) -> None:
|
||||
self.reset_calls += 1
|
||||
|
||||
|
||||
class _FakeChunkedMimi(_FakeStreamingMimi):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.decode_frames_calls = 0
|
||||
self.decode_frames_shapes: list[tuple[int, ...]] = []
|
||||
|
||||
def decode_frames(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
self.decode_frames_calls += 1
|
||||
self.decode_frames_shapes.append(tuple(codes.shape))
|
||||
return torch.arange(codes.shape[-1] * 4, dtype=torch.float32).reshape(1, -1)
|
||||
|
||||
|
||||
class _FakeMimiModel(nn.Module):
|
||||
def __init__(self, _config) -> None:
|
||||
super().__init__()
|
||||
self.encoder_transformer = nn.Linear(1, 1, bias=False)
|
||||
self.decoder_transformer = nn.Linear(1, 1, bias=False)
|
||||
self.quantizer = nn.Linear(1, 1, bias=False)
|
||||
self.encoder = SimpleNamespace(layers=[])
|
||||
self.downsample = SimpleNamespace(conv=nn.Conv1d(1, 1, 1))
|
||||
self.upsample = SimpleNamespace(conv=nn.ConvTranspose1d(1, 1, 1))
|
||||
self.decoder = SimpleNamespace(layers=[])
|
||||
|
||||
def load_state_dict(self, _state_dict, strict: bool = True):
|
||||
assert not strict
|
||||
return SimpleNamespace(
|
||||
missing_keys=[
|
||||
"encoder_transformer.weight",
|
||||
"decoder_transformer.weight",
|
||||
],
|
||||
unexpected_keys=[],
|
||||
)
|
||||
|
||||
|
||||
class _FakeMimiTransformer(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.zeros(1))
|
||||
|
||||
def load_weights(self, _state_dict, _prefix: str) -> int:
|
||||
return 80
|
||||
|
||||
|
||||
def test_moshi_checkpoint_mapping_drops_replaced_transformer_weights(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
mapper = getattr(personaplex_mimi, "_map_moshi_codec_weights", None)
|
||||
assert callable(mapper), "PersonaPlex Mimi must load its bundled codec weights without Hugging Face"
|
||||
|
||||
source = {
|
||||
"encoder.model.3.conv.conv.weight": torch.ones(1),
|
||||
"downsample.conv.conv.conv.weight": torch.ones(1),
|
||||
"quantizer.rvq_first.vq.layers.0._codebook.embedding_sum": torch.ones(1),
|
||||
"quantizer.rvq_rest.vq.layers.3._codebook._initialized": torch.ones(1),
|
||||
"encoder_transformer.transformer.layers.0.norm1.weight": torch.ones(1),
|
||||
}
|
||||
|
||||
mapped = mapper(source)
|
||||
|
||||
assert set(mapped) == {
|
||||
"encoder.layers.3.conv.weight",
|
||||
"downsample.conv.weight",
|
||||
"quantizer.semantic_residual_vector_quantizer.layers.0.codebook.embed_sum",
|
||||
"quantizer.acoustic_residual_vector_quantizer.layers.3.codebook.initialized",
|
||||
}
|
||||
assert mapped["encoder.layers.3.conv.weight"] is source["encoder.model.3.conv.conv.weight"]
|
||||
monkeypatch.setattr("transformers.MimiConfig", lambda: object())
|
||||
monkeypatch.setattr("transformers.MimiModel", _FakeMimiModel)
|
||||
monkeypatch.setattr("safetensors.torch.load_file", lambda *_args, **_kwargs: {})
|
||||
monkeypatch.setattr(personaplex_mimi, "_MimiStreamingTransformer", _FakeMimiTransformer)
|
||||
|
||||
codec = personaplex_mimi.PersonaPlexMimiCodec(
|
||||
checkpoint="/unused",
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
assert not hasattr(codec.model, "encoder_transformer")
|
||||
assert not hasattr(codec.model, "decoder_transformer")
|
||||
assert not any(
|
||||
name.startswith(("model.encoder_transformer.", "model.decoder_transformer."))
|
||||
for name, _ in codec.named_parameters()
|
||||
)
|
||||
|
||||
|
||||
def test_mimi_full_stream_reset_reuses_all_state_storage_and_clears_offsets() -> None:
|
||||
codec = personaplex_mimi.PersonaPlexMimiCodec.__new__(personaplex_mimi.PersonaPlexMimiCodec)
|
||||
nn.Module.__init__(codec)
|
||||
codec.device = torch.device("cpu")
|
||||
codec.dtype = torch.float32
|
||||
codec._enc_stages = [("conv", personaplex_mimi._StreamConv1d(nn.Conv1d(1, 1, 3)))]
|
||||
codec._dec_stages = []
|
||||
codec._downsample = personaplex_mimi._StreamConv1d(nn.Conv1d(1, 1, 3))
|
||||
codec._upsample = personaplex_mimi._StreamConvTr1d(nn.ConvTranspose1d(1, 1, 3))
|
||||
codec.encoder_transformer = personaplex_mimi._MimiStreamingTransformer(
|
||||
num_layers=1,
|
||||
dim=4,
|
||||
num_heads=1,
|
||||
context=3,
|
||||
)
|
||||
codec.decoder_transformer = personaplex_mimi._MimiStreamingTransformer(
|
||||
num_layers=1,
|
||||
dim=4,
|
||||
num_heads=1,
|
||||
context=3,
|
||||
)
|
||||
codec.streaming_init(batch_size=2)
|
||||
|
||||
conv_states = list(codec._conv_states())
|
||||
conv_buffers = [
|
||||
state.prev if isinstance(state, personaplex_mimi._StreamConv1d) else state.partial for state in conv_states
|
||||
]
|
||||
transformers = [codec.encoder_transformer, codec.decoder_transformer]
|
||||
kv_states = [kv for transformer in transformers for kv in transformer._kv]
|
||||
state_tensors = [
|
||||
*conv_buffers,
|
||||
*(kv.cache for kv in kv_states),
|
||||
*(kv.end_offset for kv in kv_states),
|
||||
*(kv.start_offset for kv in kv_states),
|
||||
*(transformer._offset for transformer in transformers),
|
||||
]
|
||||
for tensor in state_tensors:
|
||||
tensor.fill_(1)
|
||||
for state in conv_states:
|
||||
state._fresh.zero_()
|
||||
storage = [tensor.data_ptr() for tensor in state_tensors]
|
||||
|
||||
codec.reset_streaming()
|
||||
|
||||
assert storage == [tensor.data_ptr() for tensor in state_tensors]
|
||||
cleared = [
|
||||
*conv_buffers,
|
||||
*(kv.end_offset for kv in kv_states),
|
||||
*(kv.start_offset for kv in kv_states),
|
||||
*(transformer._offset for transformer in transformers),
|
||||
]
|
||||
assert all(not tensor.any() for tensor in cleared)
|
||||
assert all(state._fresh.all() for state in conv_states)
|
||||
|
||||
|
||||
def _model(*, max_sessions: int = 1) -> tuple[PersonaPlexCode2Wav, _FakeStreamingMimi]:
|
||||
mimi_config = SimpleNamespace(num_codebooks=2, sample_rate=24000, samples_per_frame=4, mimi_name=None)
|
||||
config = SimpleNamespace(mimi_config=mimi_config, mimi_name=None)
|
||||
vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
model="/unused",
|
||||
hf_config=config,
|
||||
duplex_max_sessions=max_sessions,
|
||||
),
|
||||
device_config=SimpleNamespace(device="cpu"),
|
||||
)
|
||||
model = PersonaPlexCode2Wav(vllm_config=vllm_config)
|
||||
mimi = _FakeStreamingMimi()
|
||||
model.mimi = mimi
|
||||
model._mimi_device = torch.device("cpu")
|
||||
return model, mimi
|
||||
|
||||
|
||||
def _codes(frames: int, *, start: int = 0) -> torch.Tensor:
|
||||
return torch.stack([torch.arange(start + offset, start + offset + frames) for offset in (0, 100)]).reshape(-1)
|
||||
|
||||
|
||||
def _audio(output) -> torch.Tensor:
|
||||
return output.multimodal_outputs["model_outputs"][0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("second_codes", [_codes(3), _codes(1, start=100)])
|
||||
def test_resumable_codes_emit_only_new_pcm(second_codes: torch.Tensor) -> None:
|
||||
model, mimi = _model()
|
||||
|
||||
first = model(input_ids=_codes(2), request_ids=["req"])
|
||||
second = model(input_ids=second_codes, request_ids=["req"])
|
||||
|
||||
assert _audio(first).numel() == 8
|
||||
assert _audio(second).numel() == 4
|
||||
assert mimi.decode_frame_calls == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("frames", "expected_shapes"),
|
||||
[
|
||||
(5, [(1, 2, 5)]),
|
||||
(12, [(1, 2, 5), (1, 2, 5), (1, 2, 2)]),
|
||||
],
|
||||
)
|
||||
def test_streaming_decode_uses_bounded_multi_frame_calls(
|
||||
frames: int,
|
||||
expected_shapes: list[tuple[int, ...]],
|
||||
) -> None:
|
||||
model, _ = _model()
|
||||
mimi = _FakeChunkedMimi()
|
||||
model.mimi = mimi
|
||||
|
||||
output = model(input_ids=_codes(frames), request_ids=["req"])
|
||||
|
||||
assert _audio(output).numel() == frames * 4
|
||||
assert mimi.decode_frames_shapes == expected_shapes
|
||||
assert mimi.decode_frame_calls == 0
|
||||
|
||||
|
||||
def test_profile_inputs_skip_decode_while_malformed_online_input_warns(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
model, mimi = _model()
|
||||
runtime_info = [{"meta": {"personaplex_dummy_profile": True}}]
|
||||
warnings: list[tuple] = []
|
||||
monkeypatch.setattr(personaplex_code2wav.logger, "warning", lambda *args: warnings.append(args))
|
||||
|
||||
output = model(
|
||||
input_ids=torch.arange(3),
|
||||
runtime_additional_information=runtime_info,
|
||||
)
|
||||
|
||||
assert _audio(output).numel() == 0
|
||||
assert mimi.decode_frame_calls == 0
|
||||
assert warnings == []
|
||||
|
||||
output = model(input_ids=torch.arange(3))
|
||||
|
||||
assert _audio(output).numel() == 0
|
||||
assert mimi.decode_frame_calls == 0
|
||||
assert len(warnings) == 1
|
||||
assert "not divisible by" in warnings[0][0]
|
||||
assert model.get_dummy_runtime_additional_information(2) == [
|
||||
{"meta": {"personaplex_dummy_profile": True}},
|
||||
{"meta": {"personaplex_dummy_profile": True}},
|
||||
]
|
||||
|
||||
|
||||
def test_request_id_falls_back_to_runtime_information() -> None:
|
||||
model, _ = _model()
|
||||
info = [{"request_id": "runtime-req"}]
|
||||
|
||||
model(input_ids=_codes(2), runtime_additional_information=info)
|
||||
second = model(input_ids=_codes(3), runtime_additional_information=info)
|
||||
|
||||
assert _audio(second).numel() == 4
|
||||
|
||||
|
||||
def test_decoder_slot_lifecycle_isolated_capacity_and_reuse() -> None:
|
||||
model, _ = _model(max_sessions=2)
|
||||
first = _FakeStreamingMimi(stream_value=10)
|
||||
second = _FakeStreamingMimi(stream_value=20)
|
||||
model._set_mimi_codecs([first, second])
|
||||
assert model._max_codec_sessions == 2
|
||||
|
||||
initial = model(
|
||||
input_ids=torch.cat([_codes(1), _codes(1)]),
|
||||
request_ids=["first", "second"],
|
||||
seq_token_counts=[2, 2],
|
||||
)
|
||||
continued = model(
|
||||
input_ids=torch.cat([_codes(1, start=10), _codes(1, start=20)]),
|
||||
request_ids=["first", "second"],
|
||||
seq_token_counts=[2, 2],
|
||||
)
|
||||
|
||||
initial_audio = initial.multimodal_outputs["model_outputs"]
|
||||
continued_audio = continued.multimodal_outputs["model_outputs"]
|
||||
assert initial_audio[0].tolist() == [11.0] * 4
|
||||
assert initial_audio[1].tolist() == [21.0] * 4
|
||||
assert continued_audio[0].tolist() == [12.0] * 4
|
||||
assert continued_audio[1].tolist() == [22.0] * 4
|
||||
assert model._request_codec_slots == {"first": 0, "second": 1}
|
||||
with pytest.raises(RuntimeError, match="decoder capacity 2 is exhausted"):
|
||||
model._codec_for_request("overflow")
|
||||
|
||||
model.on_requests_finished(["first"])
|
||||
assert first.reset_calls == 1
|
||||
assert second.reset_calls == 0
|
||||
replacement = model(input_ids=_codes(1), request_ids=["replacement"])
|
||||
assert _audio(replacement).tolist() == [13.0] * 4
|
||||
@@ -0,0 +1,364 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from tests.e2e.online_serving import personaplex_realtime_duplex as e2e_driver
|
||||
from vllm_omni.config.stage_config import (
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import (
|
||||
DuplexInputMode,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
from vllm_omni.experimental.fullduplex.personaplex.data_plane import (
|
||||
PersonaPlexDataPlaneContext,
|
||||
PersonaPlexDataPlaneSession,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.personaplex.input import (
|
||||
PersonaPlexPcmAppendBuffer,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.personaplex.runtime_extension import (
|
||||
PersonaPlexDuplexRuntimeExtension,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.personaplex.serving_adapter import (
|
||||
PersonaPlexServingRuntimeAdapter,
|
||||
)
|
||||
from vllm_omni.model_executor.models.personaplex.pipeline import (
|
||||
PERSONAPLEX_PIPELINE,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _pcm_payload(samples: np.ndarray, *, sample_rate_hz: int = 24000) -> dict[str, object]:
|
||||
samples = np.ascontiguousarray(samples, dtype="<f4")
|
||||
return {
|
||||
"type": "audio",
|
||||
"format": "pcm_f32le",
|
||||
"sample_rate_hz": sample_rate_hz,
|
||||
"audio": base64.b64encode(samples.tobytes()).decode("ascii"),
|
||||
}
|
||||
|
||||
|
||||
def _audio_client(
|
||||
frame_count: int = 10,
|
||||
*,
|
||||
voiced_frames: int = 0,
|
||||
) -> SimpleNamespace:
|
||||
frames = np.full((frame_count, e2e_driver.FRAME_SAMPLES), 7, dtype="<i2")
|
||||
frames[:voiced_frames] = 1000
|
||||
raw = frames.tobytes()
|
||||
event = {
|
||||
"type": "response.audio.delta",
|
||||
"response_id": "response-1",
|
||||
"sample_rate_hz": e2e_driver.SAMPLE_RATE_HZ,
|
||||
"metadata": {
|
||||
"vllm_omni": {
|
||||
"runtime_impl": "scheduler_data_plane",
|
||||
"uses_model_runner_scheduler": True,
|
||||
"runner_kv_backed": True,
|
||||
}
|
||||
},
|
||||
}
|
||||
events = SimpleNamespace(
|
||||
events=[event],
|
||||
response_audio={"response-1": [raw]},
|
||||
audio_bytes=lambda: raw,
|
||||
)
|
||||
return SimpleNamespace(events=events)
|
||||
|
||||
|
||||
def test_realtime_audio_frame_stats_separate_voiced_and_silent_frames() -> None:
|
||||
silent = np.zeros(e2e_driver.FRAME_SAMPLES, dtype="<i2")
|
||||
voiced = np.full(e2e_driver.FRAME_SAMPLES, 1000, dtype="<i2")
|
||||
|
||||
stats = e2e_driver._audio_frame_stats(
|
||||
silent.tobytes() + voiced.tobytes(),
|
||||
input_frames=3,
|
||||
voiced_frame_rms_threshold=1e-4,
|
||||
)
|
||||
|
||||
assert stats == {
|
||||
"output_frames": 2,
|
||||
"voiced_frames": 1,
|
||||
"silent_frames": 1,
|
||||
"frame_deficit": 1,
|
||||
"frame_coverage_ratio": pytest.approx(2 / 3),
|
||||
}
|
||||
|
||||
|
||||
def test_realtime_audio_frame_stats_reject_partial_codec_frame() -> None:
|
||||
partial_frame = np.zeros(e2e_driver.FRAME_SAMPLES - 1, dtype="<i2")
|
||||
|
||||
with pytest.raises(AssertionError, match="whole 80 ms codec frames"):
|
||||
e2e_driver._audio_frame_stats(
|
||||
partial_frame.tobytes(),
|
||||
input_frames=1,
|
||||
voiced_frame_rms_threshold=1e-4,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frame_count", [10, 40])
|
||||
def test_e2e_driver_rejects_inaudible_sessions(frame_count: int) -> None:
|
||||
args = SimpleNamespace(
|
||||
max_frame_deficit=4,
|
||||
voiced_frame_rms_threshold=1e-3,
|
||||
min_voiced_frames=5,
|
||||
)
|
||||
|
||||
with pytest.raises(AssertionError, match="audible speech"):
|
||||
e2e_driver._session_result(
|
||||
_audio_client(frame_count),
|
||||
input_frames=frame_count,
|
||||
args=args,
|
||||
minimum_chunks=1,
|
||||
)
|
||||
|
||||
|
||||
def test_e2e_driver_uses_absolute_audible_floor() -> None:
|
||||
args = SimpleNamespace(
|
||||
max_frame_deficit=4,
|
||||
voiced_frame_rms_threshold=1e-3,
|
||||
min_voiced_frames=5,
|
||||
)
|
||||
|
||||
e2e_driver._session_result(
|
||||
_audio_client(200, voiced_frames=5),
|
||||
input_frames=200,
|
||||
args=args,
|
||||
minimum_chunks=1,
|
||||
)
|
||||
|
||||
|
||||
def test_e2e_driver_defaults_replacement_to_full_input() -> None:
|
||||
args = e2e_driver.parse_args(
|
||||
[
|
||||
"--model",
|
||||
"/model",
|
||||
"--input-wav",
|
||||
"/input.wav",
|
||||
]
|
||||
)
|
||||
|
||||
assert args.replacement_frames == 0
|
||||
assert args.min_voiced_frames == 5
|
||||
|
||||
|
||||
def test_e2e_driver_records_and_validates_input_sha256(tmp_path: Path) -> None:
|
||||
input_wav = tmp_path / "input.wav"
|
||||
input_wav.write_bytes(b"fixture-wav")
|
||||
|
||||
identity = e2e_driver._input_identity(
|
||||
input_wav,
|
||||
expected_sha256="bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
|
||||
)
|
||||
|
||||
assert identity == {
|
||||
"path": str(input_wav.resolve()),
|
||||
"sha256": "bade4b3c163edde390ff391207d34a887257d8a9cc3b621cc8c618b6e6761304",
|
||||
}
|
||||
|
||||
|
||||
def test_e2e_driver_rejects_unexpected_input_sha256(tmp_path: Path) -> None:
|
||||
input_wav = tmp_path / "input.wav"
|
||||
input_wav.write_bytes(b"fixture-wav")
|
||||
|
||||
with pytest.raises(ValueError, match="input WAV SHA-256 mismatch"):
|
||||
e2e_driver._input_identity(input_wav, expected_sha256="0" * 64)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("attribute", "expected"),
|
||||
[
|
||||
(
|
||||
"duplex_runtime_extension",
|
||||
"vllm_omni.experimental.fullduplex.personaplex.runtime_extension.PersonaPlexDuplexRuntimeExtension",
|
||||
),
|
||||
(
|
||||
"duplex_serving_adapter",
|
||||
"vllm_omni.experimental.fullduplex.personaplex.serving_adapter.PersonaPlexServingRuntimeAdapter",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_personaplex_pipeline_registers_unified_duplex_components(
|
||||
attribute: str,
|
||||
expected: str,
|
||||
) -> None:
|
||||
assert PERSONAPLEX_PIPELINE.duplex_control_enabled is True
|
||||
assert getattr(PERSONAPLEX_PIPELINE, attribute) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("engine_arg", "expected"),
|
||||
[
|
||||
("skip_tokenizer_init", True),
|
||||
("enable_prefix_caching", False),
|
||||
],
|
||||
)
|
||||
def test_personaplex_stage_engine_args(
|
||||
engine_arg: str,
|
||||
expected: bool,
|
||||
) -> None:
|
||||
deploy_path = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
|
||||
stages = merge_pipeline_deploy(
|
||||
PERSONAPLEX_PIPELINE,
|
||||
load_deploy_config(deploy_path),
|
||||
)
|
||||
|
||||
assert [stage.yaml_engine_args.get(engine_arg) for stage in stages] == [
|
||||
expected,
|
||||
expected,
|
||||
]
|
||||
|
||||
|
||||
def test_personaplex_duplex_capacity_is_propagated_to_all_model_stages() -> None:
|
||||
deploy_path = Path(__file__).parents[5] / "vllm_omni" / "deploy" / "personaplex.yaml"
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
stages = merge_pipeline_deploy(PERSONAPLEX_PIPELINE, deploy)
|
||||
|
||||
assert deploy.duplex_session.max_sessions == 2
|
||||
assert [stage.yaml_engine_args.get("duplex_max_sessions") for stage in stages] == [2, 2]
|
||||
assert "personaplex_codec_max_sessions" not in deploy.connectors["connector_of_shared_memory"]["extra"]
|
||||
|
||||
|
||||
def test_personaplex_capabilities_are_honest() -> None:
|
||||
single = PersonaPlexServingRuntimeAdapter(lambda *_: None).capabilities(max_sessions=1)
|
||||
multi = PersonaPlexServingRuntimeAdapter(lambda *_: None).capabilities(max_sessions=2)
|
||||
|
||||
assert multi.implementation_level == "model_native_duplex"
|
||||
assert multi.input_modes == ["append_audio_chunk"]
|
||||
assert multi.chunk_period_ms == 80
|
||||
assert single.supports_multi_session is False
|
||||
assert multi.supports_multi_session is True
|
||||
assert multi.supports_multi_session_same_replica is True
|
||||
assert multi.supports_barge_in is False
|
||||
|
||||
|
||||
def test_pcm_buffer_emits_one_80ms_frame_transactionally() -> None:
|
||||
buffer = PersonaPlexPcmAppendBuffer()
|
||||
reservation = buffer.prepare_append(
|
||||
_pcm_payload(np.arange(1920, dtype=np.float32)),
|
||||
operation_id="op-1",
|
||||
chunk_period_ms=80,
|
||||
allow_emit=True,
|
||||
)
|
||||
|
||||
assert reservation is not None
|
||||
assert reservation.byte_count == 1920 * 4
|
||||
assert buffer.pending_byte_count == 0
|
||||
reservation.rollback()
|
||||
assert buffer.pending_byte_count == 1920 * 4
|
||||
|
||||
|
||||
def test_pcm_buffer_rejects_changed_sample_rate() -> None:
|
||||
buffer = PersonaPlexPcmAppendBuffer()
|
||||
assert (
|
||||
buffer.prepare_append(
|
||||
_pcm_payload(np.zeros(100, np.float32)),
|
||||
operation_id="op-1",
|
||||
chunk_period_ms=80,
|
||||
allow_emit=False,
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="sample_rate"):
|
||||
buffer.prepare_append(
|
||||
_pcm_payload(np.zeros(100, np.float32), sample_rate_hz=16000),
|
||||
operation_id="op-2",
|
||||
chunk_period_ms=80,
|
||||
allow_emit=False,
|
||||
)
|
||||
|
||||
|
||||
def test_pcm_buffer_rejects_non_finite_samples() -> None:
|
||||
buffer = PersonaPlexPcmAppendBuffer()
|
||||
|
||||
with pytest.raises(ValueError, match="finite"):
|
||||
buffer.prepare_append(
|
||||
_pcm_payload(np.array([0.0, np.nan], np.float32)),
|
||||
operation_id="op-1",
|
||||
chunk_period_ms=80,
|
||||
allow_emit=False,
|
||||
)
|
||||
|
||||
|
||||
def test_runtime_extension_builds_one_frame_scheduler_append() -> None:
|
||||
extension = PersonaPlexDuplexRuntimeExtension()
|
||||
defaults = (
|
||||
SamplingParams(temperature=0.8, top_k=10, max_tokens=10),
|
||||
SamplingParams(max_tokens=1024),
|
||||
)
|
||||
|
||||
configured = extension.configure_sampling_params(runtime_config={}, defaults=defaults)
|
||||
assert configured[0].temperature == 0.0
|
||||
assert configured[0].top_k == 1
|
||||
assert configured[0].max_tokens == 1
|
||||
assert configured[1] is defaults[1]
|
||||
|
||||
plan = extension.plan_append(
|
||||
request_id="req",
|
||||
fence=DuplexFence("session", incarnation=2, epoch=3),
|
||||
session_config={"instructions": "Be concise."},
|
||||
runtime_config={"personaplex_prefill_slots": 4},
|
||||
seq=1,
|
||||
turn_seq=1,
|
||||
mode=DuplexInputMode.APPEND_AUDIO_CHUNK,
|
||||
payload=_pcm_payload(np.zeros(1920, np.float32)),
|
||||
final=False,
|
||||
sampling_params=configured[0],
|
||||
)
|
||||
assert len(plan.prompt["prompt_token_ids"]) == 5
|
||||
duplex = plan.prompt["model_intermediate_buffer"]["duplex"]
|
||||
assert duplex["session_id"] == "session"
|
||||
assert duplex["incarnation"] == 2
|
||||
assert duplex["epoch"] == 3
|
||||
assert duplex["seq"] == 1
|
||||
assert duplex["data_plane"] is True
|
||||
|
||||
|
||||
def test_projector_emits_only_cumulative_audio_and_text_suffixes() -> None:
|
||||
encoded_sizes: list[int] = []
|
||||
|
||||
def encode_audio(audio, _sample_rate, _response_format, _speed):
|
||||
size = int(np.asarray(audio, dtype=np.float32).size)
|
||||
encoded_sizes.append(size)
|
||||
return f"audio-{size}"
|
||||
|
||||
projector = PersonaPlexDataPlaneSession(encode_audio)
|
||||
projector.begin_request("req")
|
||||
context = PersonaPlexDataPlaneContext(response_format="wav")
|
||||
|
||||
first = SimpleNamespace(
|
||||
request_id="req",
|
||||
outputs=[SimpleNamespace(text="he", multimodal_output={})],
|
||||
multimodal_output={"model_outputs": np.arange(4, dtype=np.float32), "sr": 24000},
|
||||
finished=False,
|
||||
)
|
||||
second = SimpleNamespace(
|
||||
request_id="req",
|
||||
outputs=[SimpleNamespace(text="hello", multimodal_output={})],
|
||||
multimodal_output={"model_outputs": np.arange(6, dtype=np.float32), "sr": 24000},
|
||||
finished=False,
|
||||
)
|
||||
|
||||
projected = list(
|
||||
projector.project(
|
||||
{"data_plane_outputs": [first, second]},
|
||||
context=context,
|
||||
)
|
||||
)
|
||||
|
||||
assert encoded_sizes == [4, 2]
|
||||
assert [item["text"] for item in projected] == ["he", "llo"]
|
||||
assert [item["audio_data"] for item in projected] == ["audio-4", "audio-2"]
|
||||
assert all(item["sample_rate_hz"] == 24000 for item in projected)
|
||||
assert all(item["end_of_turn"] is False for item in projected)
|
||||
@@ -0,0 +1,122 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Contract tests for the fullduplex core (PR #3907 alignment).
|
||||
|
||||
Covers the DuplexFence identity value, the session identity rules (epoch only
|
||||
advances on barge-in, response_index once per begin_response), stale-output
|
||||
dropping in the turn lifecycle. The lockstep (single-eternal-response)
|
||||
lifecycle is model-owned (see
|
||||
tests/e2e/features/fullduplex/test_personaplex_adapter.py).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.core import protocol as ev
|
||||
from vllm_omni.experimental.fullduplex.core.adapter import (
|
||||
DuplexAdapter,
|
||||
DuplexCapability,
|
||||
OutputChunk,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime
|
||||
from vllm_omni.experimental.fullduplex.core.session import (
|
||||
DuplexSession,
|
||||
DuplexSessionConfig,
|
||||
DuplexState,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_fence_is_frozen_value_identity():
|
||||
f = DuplexFence(session_id="s1")
|
||||
assert (f.epoch, f.turn_id, f.response_seq) == (0, 0, 0)
|
||||
with pytest.raises(dataclasses.FrozenInstanceError):
|
||||
f.epoch = 1 # type: ignore[misc]
|
||||
assert DuplexFence("s1", 1, 2, 3) == DuplexFence("s1", 1, 2, 3)
|
||||
assert DuplexFence("s1", 1, 2, 3) != DuplexFence("s1", 2, 2, 3)
|
||||
|
||||
|
||||
def test_session_identity_rules():
|
||||
s = DuplexSession(session_id="test", config=DuplexSessionConfig())
|
||||
r1, e1 = s.begin_response()
|
||||
r2, e2 = s.begin_response()
|
||||
assert r2 == r1 + 1, "response index advances exactly once per begin"
|
||||
assert e1 == e2 == 0, "epoch does not move on response begin"
|
||||
s.end_response()
|
||||
assert s.epoch == 0, "epoch does not move on response end"
|
||||
new_epoch = s.barge_in()
|
||||
assert new_epoch == 1, "epoch advances only on barge-in"
|
||||
assert s.is_stale(e1) and not s.is_stale(new_epoch)
|
||||
|
||||
|
||||
class _TurnAdapter(DuplexAdapter):
|
||||
"""Emits chunks until cancelled; counts respond() lifecycles."""
|
||||
|
||||
def __init__(self):
|
||||
self.responses = 0
|
||||
self.barge_ins = 0
|
||||
|
||||
def capabilities(self) -> DuplexCapability:
|
||||
return DuplexCapability(
|
||||
input_modalities=frozenset({"audio"}),
|
||||
output_modalities=frozenset({"audio"}),
|
||||
proactive=True,
|
||||
)
|
||||
|
||||
async def on_input(self, session, modality, data) -> None:
|
||||
pass
|
||||
|
||||
def should_respond(self, session) -> bool:
|
||||
return True
|
||||
|
||||
async def on_barge_in(self, session) -> None:
|
||||
self.barge_ins += 1
|
||||
|
||||
async def respond(self, session):
|
||||
self.responses += 1
|
||||
for i in range(100):
|
||||
await asyncio.sleep(0.001)
|
||||
yield OutputChunk("audio", i)
|
||||
|
||||
|
||||
async def _drive(runtime, events):
|
||||
async def gen():
|
||||
for e in events:
|
||||
await asyncio.sleep(0.005)
|
||||
yield e
|
||||
|
||||
out = []
|
||||
|
||||
async def emit(e):
|
||||
out.append(e)
|
||||
|
||||
await runtime.run(gen(), emit)
|
||||
return out
|
||||
|
||||
|
||||
def test_turn_mode_supersede_and_stale_drop():
|
||||
adapter = _TurnAdapter()
|
||||
session = DuplexSession(session_id="test", config=DuplexSessionConfig(proactive=True))
|
||||
runtime = DuplexRuntime(session, adapter)
|
||||
out = asyncio.run(
|
||||
_drive(
|
||||
runtime,
|
||||
[
|
||||
{"type": ev.RESPONSE_CREATE},
|
||||
{"type": ev.RESPONSE_CREATE}, # supersedes the first
|
||||
{"type": ev.CLOSE},
|
||||
],
|
||||
)
|
||||
)
|
||||
assert adapter.responses == 2
|
||||
assert adapter.barge_ins == 1, "superseding an in-flight response barges in"
|
||||
created = [e for e in out if e["type"] == ev.RESPONSE_CREATED]
|
||||
done = [e for e in out if e["type"] == ev.RESPONSE_DONE]
|
||||
assert len(created) == 2
|
||||
assert len(done) == 1, "stale first response must not emit done"
|
||||
assert done[0]["response_index"] == created[1]["response_index"]
|
||||
assert session.state is DuplexState.CLOSED
|
||||
@@ -0,0 +1,262 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""PersonaPlex adapter driven through the model-owned lockstep runtime.
|
||||
|
||||
PersonaPlexDuplexRuntime starts ONE response for the whole session, streams one
|
||||
agent frame per user frame, and drains on close (core stays verbatim upstream;
|
||||
the lockstep lifecycle is model policy) — all with a GPU-free stub stepper.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.core import protocol as ev
|
||||
from vllm_omni.experimental.fullduplex.core.session import (
|
||||
DuplexSession,
|
||||
DuplexSessionConfig,
|
||||
DuplexState,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.personaplex.adapter import (
|
||||
PersonaPlexDuplexAdapter,
|
||||
PersonaPlexDuplexRuntime,
|
||||
)
|
||||
from vllm_omni.experimental.fullduplex.personaplex.config import FRAME_SIZE
|
||||
from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class StubStepper:
|
||||
"""Deterministic FrameStepper: frame i -> audio filled with i, text 't{i}' on odds."""
|
||||
|
||||
sample_rate = 24000
|
||||
frame_size = FRAME_SIZE
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.opened: tuple[str | None, str | None] | None = None
|
||||
self.steps = 0
|
||||
|
||||
def open_session(self, voice_prompt=None, persona=None) -> None:
|
||||
self.opened = (voice_prompt, persona)
|
||||
self.steps = 0
|
||||
|
||||
def step(self, user_pcm):
|
||||
assert user_pcm.shape[0] == self.frame_size
|
||||
self.steps += 1
|
||||
i = self.steps
|
||||
return FrameOutput(
|
||||
audio=np.full(self.frame_size, float(i), dtype=np.float32),
|
||||
text=(f"t{i}" if i % 2 else None),
|
||||
)
|
||||
|
||||
|
||||
class RecordingAdapter(PersonaPlexDuplexAdapter):
|
||||
def __init__(self, stepper: StubStepper) -> None:
|
||||
super().__init__(stepper)
|
||||
self.closed = False
|
||||
|
||||
async def on_close(self, session: DuplexSession) -> None:
|
||||
self.closed = True
|
||||
await super().on_close(session)
|
||||
|
||||
|
||||
def _collector():
|
||||
out: list[dict] = []
|
||||
|
||||
async def emit(event: dict) -> None:
|
||||
out.append(event)
|
||||
|
||||
return out, emit
|
||||
|
||||
|
||||
async def _feed(events):
|
||||
for e in events:
|
||||
yield e
|
||||
|
||||
|
||||
async def _raising_feed(frame):
|
||||
yield {"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}
|
||||
await asyncio.sleep(0)
|
||||
raise RuntimeError("input failed")
|
||||
|
||||
|
||||
def _respond_tasks() -> list[asyncio.Task]:
|
||||
current = asyncio.current_task()
|
||||
return [
|
||||
task
|
||||
for task in asyncio.all_tasks()
|
||||
if task is not current and "_respond_once" in getattr(task.get_coro(), "__qualname__", "")
|
||||
]
|
||||
|
||||
|
||||
def _session():
|
||||
cfg = DuplexSessionConfig(
|
||||
input_modalities=("audio",),
|
||||
output_modalities=("audio", "text"),
|
||||
)
|
||||
return DuplexSession("s", cfg)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lockstep_one_agent_frame_per_user_frame():
|
||||
stub = StubStepper()
|
||||
adapter = PersonaPlexDuplexAdapter(stub, voice_prompt="NATF2.pt", persona="be terse")
|
||||
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||
out, emit = _collector()
|
||||
|
||||
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||
events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame} for _ in range(3)]
|
||||
events.append({"type": ev.CLOSE})
|
||||
await rt.run(_feed(events), emit)
|
||||
|
||||
types = [e["type"] for e in out]
|
||||
# exactly one response for the whole session, drained cleanly on close
|
||||
assert types.count(ev.RESPONSE_CREATED) == 1
|
||||
assert types.count(ev.RESPONSE_DONE) == 1
|
||||
audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
|
||||
text = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "text"]
|
||||
assert len(audio) == 3 # one agent frame per user frame
|
||||
assert [int(e["data"][0]) for e in audio] == [1, 2, 3] # lockstep order
|
||||
assert text == ["t1", "t3"]
|
||||
assert stub.steps == 3
|
||||
assert stub.opened == ("NATF2.pt", "be terse")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unaligned_chunks_are_reframed_to_80ms():
|
||||
stub = StubStepper()
|
||||
adapter = PersonaPlexDuplexAdapter(stub)
|
||||
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||
out, emit = _collector()
|
||||
|
||||
# 2.5 frames in one chunk + 0.5 frame in another => 3 whole frames
|
||||
events = [
|
||||
{"type": ev.INPUT_APPEND, "modality": "audio", "data": np.zeros(FRAME_SIZE * 5 // 2, dtype=np.float32)},
|
||||
{"type": ev.INPUT_APPEND, "modality": "audio", "data": np.zeros(FRAME_SIZE // 2, dtype=np.float32)},
|
||||
{"type": ev.CLOSE},
|
||||
]
|
||||
await rt.run(_feed(events), emit)
|
||||
audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
|
||||
assert len(audio) == 3
|
||||
assert stub.steps == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_without_input_is_clean():
|
||||
adapter = PersonaPlexDuplexAdapter(StubStepper())
|
||||
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||
out, emit = _collector()
|
||||
await rt.run(_feed([{"type": ev.CLOSE}]), emit)
|
||||
# no input -> response never started -> no created/done, no error
|
||||
assert not any(e["type"] == ev.ERROR for e in out)
|
||||
assert not any(e["type"] == ev.RESPONSE_CREATED for e in out)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_iterator_eof_without_close_drains_and_stops():
|
||||
# Transport drop: the input iterator ends WITHOUT an explicit CLOSE event.
|
||||
# The runtime must still signal on_close so the eternal response drains
|
||||
# instead of blocking on its inbox forever.
|
||||
stub = StubStepper()
|
||||
adapter = PersonaPlexDuplexAdapter(stub)
|
||||
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||
out, emit = _collector()
|
||||
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||
events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}]
|
||||
await asyncio.wait_for(rt.run(_feed(events), emit), timeout=5.0)
|
||||
types = [e["type"] for e in out]
|
||||
assert types.count(ev.RESPONSE_DONE) == 1
|
||||
assert stub.steps == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_input_exception_closes_response_and_runtime() -> None:
|
||||
stub = StubStepper()
|
||||
adapter = RecordingAdapter(stub)
|
||||
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||
out, emit = _collector()
|
||||
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||
|
||||
with pytest.raises(RuntimeError, match="input failed"):
|
||||
await rt.run(_raising_feed(frame), emit)
|
||||
|
||||
leaked = _respond_tasks()
|
||||
for task in leaked:
|
||||
task.cancel()
|
||||
if leaked:
|
||||
await asyncio.gather(*leaked, return_exceptions=True)
|
||||
|
||||
assert adapter.closed is True
|
||||
assert rt.session.state is DuplexState.CLOSED
|
||||
assert not leaked
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_session_config_gets_lockstep_lifecycle():
|
||||
# The lifecycle comes from the model-owned runtime class, not a session
|
||||
# flag: a plain default-constructed session config runs lockstep.
|
||||
stub = StubStepper()
|
||||
adapter = PersonaPlexDuplexAdapter(stub)
|
||||
session = DuplexSession(
|
||||
"s",
|
||||
DuplexSessionConfig(input_modalities=("audio",), output_modalities=("audio", "text")),
|
||||
)
|
||||
rt = PersonaPlexDuplexRuntime(session, adapter)
|
||||
out, emit = _collector()
|
||||
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||
events = [{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame}, {"type": ev.CLOSE}]
|
||||
await rt.run(_feed(events), emit)
|
||||
types = [e["type"] for e in out]
|
||||
assert types.count(ev.RESPONSE_CREATED) == 1
|
||||
assert types.count(ev.RESPONSE_DONE) == 1
|
||||
assert stub.steps == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_tail_is_padded_and_stepped_on_close():
|
||||
# A trailing sub-frame chunk is zero-padded and stepped at close (same
|
||||
# contract as PersonaPlexSession.flush) rather than silently dropped.
|
||||
stub = StubStepper()
|
||||
adapter = PersonaPlexDuplexAdapter(stub)
|
||||
rt = PersonaPlexDuplexRuntime(_session(), adapter)
|
||||
out, emit = _collector()
|
||||
events = [
|
||||
{"type": ev.INPUT_APPEND, "modality": "audio", "data": np.ones(FRAME_SIZE // 2, dtype=np.float32)},
|
||||
{"type": ev.CLOSE},
|
||||
]
|
||||
await rt.run(_feed(events), emit)
|
||||
audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
|
||||
assert len(audio) == 1
|
||||
assert stub.steps == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_eternal_response_across_barge_in():
|
||||
# Moved from the core contract tests when the lockstep lifecycle became
|
||||
# model-owned: ONE response for the whole session, RESPONSE_CANCEL bumps the
|
||||
# epoch fence (playback cursor) but the stream keeps flowing, and close
|
||||
# drains via on_close.
|
||||
stub = StubStepper()
|
||||
adapter = PersonaPlexDuplexAdapter(stub)
|
||||
session = _session()
|
||||
rt = PersonaPlexDuplexRuntime(session, adapter)
|
||||
out, emit = _collector()
|
||||
|
||||
frame = np.zeros(FRAME_SIZE, dtype=np.float32)
|
||||
events = [
|
||||
{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame},
|
||||
{"type": ev.RESPONSE_CANCEL}, # barge-in: epoch bumps, stream continues
|
||||
{"type": ev.INPUT_APPEND, "modality": "audio", "data": frame},
|
||||
{"type": ev.CLOSE},
|
||||
]
|
||||
await rt.run(_feed(events), emit)
|
||||
|
||||
types = [e["type"] for e in out]
|
||||
assert types.count(ev.RESPONSE_CREATED) == 1, "exactly one eternal response"
|
||||
assert types.count(ev.RESPONSE_DONE) == 1
|
||||
audio = [e for e in out if e["type"] == ev.RESPONSE_DELTA and e["modality"] == "audio"]
|
||||
assert len(audio) == 2, "stream keeps flowing across barge-in"
|
||||
assert session.epoch == 1, "cancel still advanced the epoch fence"
|
||||
assert stub.steps == 2
|
||||
@@ -0,0 +1,178 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""GPU-free tests for the batched PersonaPlex slot manager (stub engine)."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput, PrefillStep
|
||||
from vllm_omni.experimental.fullduplex.personaplex.serving.batched import (
|
||||
BatchedSessionManager,
|
||||
SlotPhase,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
FRAME = 4
|
||||
|
||||
|
||||
class StubEngine:
|
||||
"""Records the slot-API call sequence; echoes each row's pcm as its output."""
|
||||
|
||||
frame_size = FRAME
|
||||
|
||||
def __init__(self, batch_size: int = 4, prefill_len: int = 2) -> None:
|
||||
self.config = SimpleNamespace(batch_size=batch_size)
|
||||
self.prefill_len = prefill_len
|
||||
self.recycled: list[int] = []
|
||||
self.voice_ended: list[int] = []
|
||||
self.prefill_done: list[int] = []
|
||||
self.steps: list[tuple[np.ndarray, dict | None]] = []
|
||||
|
||||
def prefill_steps(self, persona=None):
|
||||
steps = [PrefillStep("sacrifice"), PrefillStep("voice", embedding=None), PrefillStep("voice_end")]
|
||||
steps += [PrefillStep("tokens", moshi_tokens=None, text_token=3, user_sine=True)] * self.prefill_len
|
||||
return steps
|
||||
|
||||
def recycle_slot(self, b):
|
||||
self.recycled.append(b)
|
||||
|
||||
def voice_end(self, b):
|
||||
self.voice_ended.append(b)
|
||||
|
||||
def end_prefill(self, b):
|
||||
self.prefill_done.append(b)
|
||||
|
||||
def step_batch(self, pcm, prefill=None):
|
||||
self.steps.append((pcm.copy(), dict(prefill) if prefill else None))
|
||||
return [FrameOutput(audio=pcm[b].copy(), text=None) for b in range(self.config.batch_size)]
|
||||
|
||||
|
||||
def test_all_idle_does_not_step():
|
||||
eng = StubEngine()
|
||||
mgr = BatchedSessionManager(eng)
|
||||
assert mgr.tick() is False
|
||||
assert eng.steps == []
|
||||
|
||||
|
||||
def test_pristine_slot_is_live_immediately():
|
||||
eng = StubEngine()
|
||||
mgr = BatchedSessionManager(eng)
|
||||
slot, ready_now = mgr.acquire()
|
||||
assert ready_now is True
|
||||
assert slot.phase is SlotPhase.LIVE
|
||||
assert eng.recycled == []
|
||||
|
||||
|
||||
def test_custom_persona_prefills_and_recycles_on_tick_thread():
|
||||
eng = StubEngine(prefill_len=2)
|
||||
mgr = BatchedSessionManager(eng)
|
||||
slot, ready_now = mgr.acquire(persona="custom")
|
||||
assert ready_now is False
|
||||
assert slot.phase is SlotPhase.PREFILL
|
||||
assert eng.recycled == [] # deferred: reset must run between engine steps
|
||||
|
||||
ready = []
|
||||
slot.on_ready = lambda: ready.append(True)
|
||||
|
||||
mgr.tick() # recycle + sacrifice
|
||||
assert eng.recycled == [slot.index]
|
||||
assert eng.steps[-1][1][slot.index].kind == "sacrifice"
|
||||
|
||||
mgr.tick() # voice frame
|
||||
assert eng.steps[-1][1][slot.index].kind == "voice"
|
||||
assert eng.voice_ended == []
|
||||
|
||||
mgr.tick() # voice_end executed BETWEEN ticks, then first tokens frame
|
||||
assert eng.voice_ended == [slot.index]
|
||||
assert eng.steps[-1][1][slot.index].kind == "tokens"
|
||||
|
||||
mgr.tick() # second tokens frame
|
||||
assert not ready
|
||||
mgr.tick() # prefill exhausted -> live
|
||||
assert eng.prefill_done == [slot.index]
|
||||
assert slot.phase is SlotPhase.LIVE
|
||||
assert ready == [True]
|
||||
|
||||
|
||||
def test_pristine_slot_after_ticking_requires_recycle():
|
||||
# step_batch advances EVERY row each tick, so a pristine idle row that has
|
||||
# lived through ticks has drifted from its boot prefill and must be
|
||||
# recycled before a client may join it.
|
||||
eng = StubEngine()
|
||||
mgr = BatchedSessionManager(eng)
|
||||
first, ready_now = mgr.acquire()
|
||||
assert ready_now is True # claimed before any tick: truly pristine
|
||||
mgr.tick()
|
||||
second, ready_now2 = mgr.acquire()
|
||||
assert second.index != first.index
|
||||
assert ready_now2 is False
|
||||
assert second.phase is SlotPhase.PREFILL
|
||||
|
||||
|
||||
def test_used_slot_requires_recycle_even_with_default_persona():
|
||||
eng = StubEngine()
|
||||
mgr = BatchedSessionManager(eng)
|
||||
slot, ready_now = mgr.acquire()
|
||||
assert ready_now
|
||||
mgr.release(slot)
|
||||
slot2, ready_now2 = mgr.acquire()
|
||||
assert slot2.index == slot.index
|
||||
assert ready_now2 is False
|
||||
assert slot2.phase is SlotPhase.PREFILL
|
||||
|
||||
|
||||
def test_acquire_returns_none_when_full():
|
||||
eng = StubEngine(batch_size=2)
|
||||
mgr = BatchedSessionManager(eng)
|
||||
assert mgr.acquire() is not None
|
||||
assert mgr.acquire() is not None
|
||||
assert mgr.acquire() is None
|
||||
|
||||
|
||||
def test_feed_frames_epoch_guard_and_backlog_cap():
|
||||
eng = StubEngine()
|
||||
mgr = BatchedSessionManager(eng, max_buffered_frames=3)
|
||||
slot, _ = mgr.acquire()
|
||||
mgr.feed(slot, np.arange(FRAME * 2, dtype=np.float32), slot.epoch)
|
||||
assert len(slot.inq) == 2
|
||||
mgr.feed(slot, np.zeros(3, dtype=np.float32), slot.epoch) # partial tail buffered
|
||||
assert len(slot.inq) == 2 and slot.pending.shape[0] == 3
|
||||
mgr.feed(slot, np.zeros(FRAME * 5, dtype=np.float32), slot.epoch)
|
||||
assert len(slot.inq) == 3 # oldest dropped beyond the cap
|
||||
mgr.feed(slot, np.zeros(FRAME, dtype=np.float32), slot.epoch - 1) # stale epoch ignored
|
||||
assert len(slot.inq) == 3
|
||||
|
||||
|
||||
def test_live_slot_consumes_input_and_emits_output():
|
||||
eng = StubEngine()
|
||||
mgr = BatchedSessionManager(eng)
|
||||
slot, _ = mgr.acquire()
|
||||
got = []
|
||||
slot.on_output = got.append
|
||||
frame = np.full(FRAME, 0.5, dtype=np.float32)
|
||||
mgr.feed(slot, frame, slot.epoch)
|
||||
mgr.tick()
|
||||
assert len(got) == 1
|
||||
np.testing.assert_array_equal(got[0].audio, frame)
|
||||
# queue drained: next tick feeds zeros
|
||||
mgr.tick()
|
||||
np.testing.assert_array_equal(got[1].audio, np.zeros(FRAME, dtype=np.float32))
|
||||
|
||||
|
||||
def test_release_mid_flight_stops_emission():
|
||||
eng = StubEngine()
|
||||
mgr = BatchedSessionManager(eng)
|
||||
slot, _ = mgr.acquire()
|
||||
got = []
|
||||
slot.on_output = got.append
|
||||
mgr.tick()
|
||||
assert len(got) == 1
|
||||
mgr.release(slot)
|
||||
# another slot keeps the batch non-idle
|
||||
other, _ = mgr.acquire() # same slot index, now PREFILL (used)
|
||||
assert other.phase is SlotPhase.PREFILL
|
||||
mgr.tick()
|
||||
assert len(got) == 1 # no further output for the released epoch
|
||||
@@ -0,0 +1,240 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""GPU e2e for PersonaPlex elastic batching (bit-exact isolation + recycle).
|
||||
|
||||
Greedy decoding makes every cross-slot leak observable as a bit difference:
|
||||
|
||||
1. Isolation (bit-exact): recycling one slot mid-flight leaves every other
|
||||
slot's output bit-identical to a run that never recycled.
|
||||
2. Replay exactness: a recycled slot's opening live ticks match a pristine boot
|
||||
at numerical-noise level (a broken KV mask / staging roll garbles frame 0).
|
||||
3. Coherence (behavioral): a recycled slot fed a real spoken question produces
|
||||
a real reply in its inner-monologue text stream.
|
||||
|
||||
Run (moshi-free native stack; ~30 GB GPU + HF access to the gated repo; set
|
||||
PPLEX_QUESTION_WAV to a spoken-question wav to enable the coherence test):
|
||||
pytest tests/e2e/features/fullduplex/test_personaplex_elastic.py -v -s
|
||||
|
||||
Bench (B sweep, ms/tick vs the 80 ms budget):
|
||||
python tests/e2e/features/fullduplex/test_personaplex_elastic.py --bench 2 4 8 16
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
# slow + H100 + omni routes this module into the weekly Buildkite E2E group
|
||||
# ("Omni · H100" in .buildkite/cuda/test-weekly.yml: -m "slow and H100 and
|
||||
# omni" --run-level full_model); the ready/merge CPU sweeps exclude it by mark.
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.omni, pytest.mark.slow, pytest.mark.H100]
|
||||
|
||||
if not torch.cuda.is_available(): # noqa: E402
|
||||
pytest.skip("needs CUDA", allow_module_level=True)
|
||||
|
||||
from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig # noqa: E402
|
||||
from vllm_omni.experimental.fullduplex.personaplex.runtime import ( # noqa: E402
|
||||
PersonaPlexEngine,
|
||||
)
|
||||
|
||||
B = 4
|
||||
LIVE_TICKS = 40 # live frames compared per claim
|
||||
RECYCLE_AT = 10
|
||||
|
||||
|
||||
def _make_engine() -> PersonaPlexEngine:
|
||||
cfg = PersonaPlexConfig(batch_size=B) # native stepper is greedy-only
|
||||
try:
|
||||
return PersonaPlexEngine(cfg).load()
|
||||
except Exception as exc: # gated-repo access is environment-dependent
|
||||
name = type(exc).__name__
|
||||
if "Gated" in name or "401" in str(exc) or "403" in str(exc):
|
||||
pytest.skip(f"no access to the gated nvidia/personaplex-7b-v1 repo: {name}")
|
||||
raise
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def engine() -> PersonaPlexEngine:
|
||||
return _make_engine()
|
||||
|
||||
|
||||
def _slot_stream(slot: int, tick: int, frame: int) -> np.ndarray:
|
||||
"""Deterministic per-slot input: distinct low-amplitude tones."""
|
||||
sr = 24000
|
||||
freq = 180.0 + 90.0 * slot
|
||||
t0 = tick * frame
|
||||
t = (np.arange(frame) + t0) / sr
|
||||
return (0.12 * np.sin(2 * np.pi * freq * t)).astype(np.float32)
|
||||
|
||||
|
||||
def _run(engine, total_ticks: int, recycle: bool):
|
||||
"""Run the lockstep loop; optionally recycle slot 1 at RECYCLE_AT.
|
||||
|
||||
Returns per-slot lists of (audio | None) per live tick index, where the
|
||||
recycled slot's list restarts its input stream at its post-prefill tick.
|
||||
"""
|
||||
frame = engine.frame_size
|
||||
outs: list[list[np.ndarray | None]] = [[] for _ in range(B)]
|
||||
stream_tick = [0] * B # per-slot conversation clock driving its input
|
||||
prefill = None
|
||||
tick = 0
|
||||
while tick < total_ticks:
|
||||
if recycle and tick == RECYCLE_AT and prefill is None:
|
||||
engine.recycle_slot(1)
|
||||
prefill = list(engine.prefill_steps(None))
|
||||
pcm = np.zeros((B, frame), dtype=np.float32)
|
||||
step_for_1 = None
|
||||
if prefill:
|
||||
while prefill and prefill[0].kind == "voice_end":
|
||||
engine.voice_end(1)
|
||||
prefill.pop(0)
|
||||
if prefill:
|
||||
step_for_1 = prefill.pop(0)
|
||||
else:
|
||||
engine.end_prefill(1)
|
||||
prefill = None
|
||||
stream_tick[1] = 0 # restart slot 1's input stream from its own t=0
|
||||
for b in range(B):
|
||||
if b == 1 and step_for_1 is not None:
|
||||
continue # engine overrides the user channel for prefill rows
|
||||
pcm[b] = _slot_stream(b, stream_tick[b], frame)
|
||||
stream_tick[b] += 1
|
||||
frame_outs = engine.step_batch(pcm, {1: step_for_1} if step_for_1 is not None else None)
|
||||
for b in range(B):
|
||||
if b == 1 and step_for_1 is not None:
|
||||
continue # prefill tick: nothing conversational to record
|
||||
frame_out = frame_outs[b]
|
||||
outs[b].append(None if frame_out is None else frame_out.audio)
|
||||
tick += 1
|
||||
return outs
|
||||
|
||||
|
||||
def test_recycle_isolation_and_fresh_equivalence(engine):
|
||||
engine.open_batch()
|
||||
prefill_len = sum(1 for s in engine.prefill_steps(None) if s.kind != "voice_end")
|
||||
total = RECYCLE_AT + prefill_len + LIVE_TICKS
|
||||
|
||||
base = _run(engine, total, recycle=False)
|
||||
|
||||
engine.open_batch() # deterministic re-boot (greedy, fixed voice/persona)
|
||||
rec = _run(engine, total, recycle=True)
|
||||
|
||||
# Claim 1 -- isolation: untouched slots bit-identical to the no-recycle run.
|
||||
for b in (0, 2, 3):
|
||||
assert len(base[b]) == total and len(rec[b]) == total
|
||||
for i, (x, y) in enumerate(zip(base[b], rec[b])):
|
||||
assert (x is None) == (y is None), f"slot {b} tick {i} validity diverged"
|
||||
if x is not None:
|
||||
assert np.array_equal(x, y), f"slot {b} tick {i} audio diverged (recycle leaked)"
|
||||
|
||||
# Claim 2 -- replay exactness: the recycled slot's opening live ticks match a
|
||||
# pristine boot at numerical-noise level. A broken KV mask / staging-ring
|
||||
# roll / Mimi reset garbles the audio from the very first frame (diffs of
|
||||
# 0.1-1.0); a correct replay differs only by decoder-state rounding (~1e-4).
|
||||
#
|
||||
# Full-length bitwise identity is NOT expected, for two structural reasons:
|
||||
# the replayed system prompt sits at a different ABSOLUTE RoPE offset (bf16
|
||||
# rounds differently there), and the recycled Mimi ENCODER cannot restore a
|
||||
# fresh stream's start padding (a batch-shared scalar), so from the first
|
||||
# non-silent user frame it emits different-but-equally-valid codes for the
|
||||
# same audio. Once any near-tie flips a token, the two equally valid
|
||||
# trajectories legitimately diverge; coherence of the recycled slot is
|
||||
# asserted behaviorally in test_recycled_slot_converses.
|
||||
fresh = [a for a in base[1] if a is not None][:LIVE_TICKS]
|
||||
reborn = [a for a in rec[1][RECYCLE_AT:] if a is not None][:LIVE_TICKS]
|
||||
assert len(reborn) >= LIVE_TICKS // 2, "recycled slot produced too few live frames"
|
||||
diffs = [float(np.max(np.abs(x - y))) for x, y in zip(fresh, reborn)]
|
||||
first_flip = next((i for i, d in enumerate(diffs) if d > 5e-2), len(diffs))
|
||||
print(
|
||||
f"[elastic] replay exactness: {first_flip}/{len(diffs)} opening ticks at noise level "
|
||||
f"(head diffs: {[f'{d:.1e}' for d in diffs[:6]]})"
|
||||
)
|
||||
assert first_flip >= 4, (
|
||||
f"recycled slot diverged from fresh behavior at tick {first_flip} "
|
||||
f"(diffs {diffs[: first_flip + 2]}); the system-prompt replay is broken"
|
||||
)
|
||||
assert max(diffs[:4]) < 2e-2, f"reborn audio not at noise level: {max(diffs[:4])}"
|
||||
|
||||
|
||||
def test_recycled_slot_converses(engine):
|
||||
"""Behavioral guarantee: a recycled slot holds a coherent conversation.
|
||||
|
||||
Feed a real spoken question into a slot that was recycled mid-flight (while
|
||||
other slots keep running); the model's own inner-monologue text stream must
|
||||
produce a real reply. This is the end-user property the encoder-state
|
||||
caveats above cannot break.
|
||||
"""
|
||||
import os
|
||||
|
||||
import sphn
|
||||
|
||||
wav_path = os.environ.get("PPLEX_QUESTION_WAV")
|
||||
if not wav_path:
|
||||
pytest.skip("set PPLEX_QUESTION_WAV to a spoken-question wav to run")
|
||||
wav, _ = sphn.read(wav_path, sample_rate=24000)
|
||||
wav = wav[0].astype(np.float32)
|
||||
frame = engine.frame_size
|
||||
pad = (-len(wav)) % frame
|
||||
wav = np.concatenate([wav, np.zeros(pad, dtype=np.float32)])
|
||||
q_frames = [wav[i : i + frame] for i in range(0, len(wav), frame)]
|
||||
|
||||
engine.open_batch()
|
||||
pcm = np.zeros((B, frame), dtype=np.float32)
|
||||
for _ in range(5): # everyone live briefly
|
||||
engine.step_batch(pcm)
|
||||
engine.recycle_slot(1)
|
||||
prefill = list(engine.prefill_steps(None))
|
||||
while prefill:
|
||||
while prefill and prefill[0].kind == "voice_end":
|
||||
engine.voice_end(1)
|
||||
prefill.pop(0)
|
||||
if not prefill:
|
||||
break
|
||||
engine.step_batch(pcm, {1: prefill.pop(0)})
|
||||
engine.end_prefill(1)
|
||||
|
||||
text = ""
|
||||
total = len(q_frames) + 120 # question + ~10 s to answer
|
||||
for i in range(total):
|
||||
pcm = np.zeros((B, frame), dtype=np.float32)
|
||||
if i < len(q_frames):
|
||||
pcm[1] = q_frames[i]
|
||||
outs = engine.step_batch(pcm)
|
||||
if outs[1] is not None and outs[1].text:
|
||||
text += outs[1].text
|
||||
alpha = sum(c.isalpha() for c in text)
|
||||
print(f"[elastic] recycled slot inner monologue: {text!r}")
|
||||
assert alpha >= 15, f"recycled slot produced no coherent reply: {text!r}"
|
||||
|
||||
|
||||
def _bench(batch_sizes: list[int]) -> None:
|
||||
import time
|
||||
|
||||
for bs in batch_sizes:
|
||||
cfg = PersonaPlexConfig(batch_size=max(bs, 2))
|
||||
eng = PersonaPlexEngine(cfg).load()
|
||||
eng.open_batch()
|
||||
frame = eng.frame_size
|
||||
pcm = np.zeros((cfg.batch_size, frame), dtype=np.float32)
|
||||
for _ in range(10): # warm
|
||||
eng.step_batch(pcm)
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.monotonic()
|
||||
n = 100
|
||||
for _ in range(n):
|
||||
eng.step_batch(pcm)
|
||||
torch.cuda.synchronize()
|
||||
ms = (time.monotonic() - t0) / n * 1000
|
||||
mem = torch.cuda.max_memory_allocated() / 2**30
|
||||
print(f"B={cfg.batch_size:3d} {ms:6.1f} ms/tick RTF/stream {ms / 80:.3f} {mem:.1f} GiB")
|
||||
del eng
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sizes = [int(x) for x in sys.argv[2:]] if len(sys.argv) > 2 and sys.argv[1] == "--bench" else [2, 4, 8]
|
||||
_bench(sizes)
|
||||
@@ -0,0 +1,101 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Single-session lease teardown: a cancelled engine call must be drained.
|
||||
|
||||
Cancelling an asyncio task blocked on a worker-thread engine call returns
|
||||
immediately while the thread keeps running (executor futures cannot be
|
||||
interrupted). The server therefore tracks the in-flight future on the session
|
||||
state and waits it out before releasing the capacity-one lease; otherwise a
|
||||
fast reconnect could call ``open()`` on the shared engine while the previous
|
||||
connection's final ``step()`` is still mutating it.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import sys
|
||||
import threading
|
||||
import types
|
||||
|
||||
import pytest
|
||||
|
||||
try:
|
||||
import sphn # noqa: F401
|
||||
except ImportError:
|
||||
# These lease tests never encode/decode Opus. Keep the optional serving
|
||||
# dependency from hiding dependency-free teardown regressions.
|
||||
sys.modules["sphn"] = types.ModuleType("sphn")
|
||||
|
||||
from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
|
||||
from vllm_omni.experimental.fullduplex.personaplex.serving.server import DuplexServer
|
||||
from vllm_omni.experimental.fullduplex.personaplex.session import PersonaPlexServingSessionState
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_engine_call_is_drained_before_lease_release():
|
||||
server = DuplexServer(PersonaPlexConfig()) # engine never loaded: stub calls only
|
||||
state = PersonaPlexServingSessionState()
|
||||
|
||||
started = threading.Event()
|
||||
release = threading.Event()
|
||||
finished = threading.Event()
|
||||
|
||||
def slow_engine_op():
|
||||
started.set()
|
||||
release.wait(timeout=5.0)
|
||||
finished.set()
|
||||
return []
|
||||
|
||||
task = asyncio.create_task(server._engine_call(state, slow_engine_op))
|
||||
await asyncio.to_thread(started.wait, 5.0)
|
||||
|
||||
# Cancel while the worker thread is mid-call: the await returns at once,
|
||||
# the thread keeps running, and its completion event stays tracked on the
|
||||
# state (the asyncio future itself reads as cancelled/done here, which is
|
||||
# exactly why the drain must not rely on it).
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
assert state.inflight is not None
|
||||
assert not state.inflight.is_set()
|
||||
|
||||
# The drain must genuinely wait for the worker, not just return.
|
||||
drain = asyncio.create_task(server._drain_inflight(state))
|
||||
await asyncio.sleep(0.05)
|
||||
assert not drain.done()
|
||||
|
||||
release.set()
|
||||
await asyncio.wait_for(drain, timeout=5.0)
|
||||
assert finished.is_set()
|
||||
assert state.inflight.is_set()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_engine_call_clears_inflight():
|
||||
server = DuplexServer(PersonaPlexConfig())
|
||||
state = PersonaPlexServingSessionState()
|
||||
result = await server._engine_call(state, lambda: "ok")
|
||||
assert result == "ok"
|
||||
assert state.inflight is None
|
||||
await server._drain_inflight(state) # no-op on a clean state
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_timeout_keeps_lease_until_worker_finishes():
|
||||
server = DuplexServer(PersonaPlexConfig())
|
||||
server._active = True
|
||||
state = PersonaPlexServingSessionState()
|
||||
state.inflight = threading.Event()
|
||||
|
||||
drained = await server._release_lease_after_drain(state, timeout=0.01)
|
||||
|
||||
assert drained is False
|
||||
assert server._active is True
|
||||
|
||||
state.inflight.set()
|
||||
for _ in range(100):
|
||||
if not server._active:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert server._active is False
|
||||
@@ -0,0 +1,97 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Lockstep-contract tests for the PersonaPlex session driver (no GPU / no moshi).
|
||||
|
||||
A stub :class:`FrameStepper` stands in for the real Moshi backend so we can assert
|
||||
the framing + 1-frame-in / 1-frame-out invariants without weights.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.fullduplex.personaplex.config import FRAME_SIZE, PersonaPlexConfig
|
||||
from vllm_omni.experimental.fullduplex.personaplex.engine import FrameOutput, FrameStepper
|
||||
from vllm_omni.experimental.fullduplex.personaplex.session import PersonaPlexSession
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class StubStepper:
|
||||
"""Deterministic FrameStepper: frame i -> audio filled with i, text 't{i}' on odds."""
|
||||
|
||||
sample_rate = 24000
|
||||
frame_size = FRAME_SIZE
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.opened: tuple[str | None, str | None] | None = None
|
||||
self.steps = 0
|
||||
|
||||
def open_session(self, voice_prompt=None, persona=None) -> None:
|
||||
self.opened = (voice_prompt, persona)
|
||||
self.steps = 0
|
||||
|
||||
def step(self, user_pcm):
|
||||
assert user_pcm.shape[0] == self.frame_size
|
||||
self.steps += 1
|
||||
i = self.steps
|
||||
return FrameOutput(
|
||||
audio=np.full(self.frame_size, float(i), dtype=np.float32),
|
||||
text=(f"t{i}" if i % 2 else None),
|
||||
)
|
||||
|
||||
|
||||
def test_stub_satisfies_protocol():
|
||||
assert isinstance(StubStepper(), FrameStepper)
|
||||
|
||||
|
||||
def test_feed_emits_one_output_per_complete_frame():
|
||||
s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
|
||||
s.open(voice_prompt="NATF2.pt", persona="be terse")
|
||||
out = s.feed(np.zeros(FRAME_SIZE * 3, dtype=np.float32))
|
||||
assert len(out) == 3
|
||||
# lockstep order preserved: frame i -> audio filled with i
|
||||
assert [int(o.audio[0]) for o in out] == [1, 2, 3]
|
||||
assert [o.text for o in out] == ["t1", None, "t3"]
|
||||
|
||||
|
||||
def test_partial_frames_are_buffered_until_complete():
|
||||
stub = StubStepper()
|
||||
s = PersonaPlexSession(stub, PersonaPlexConfig())
|
||||
s.open()
|
||||
assert s.feed(np.zeros(FRAME_SIZE // 2, dtype=np.float32)) == [] # nothing yet
|
||||
assert stub.steps == 0
|
||||
out = s.feed(np.zeros(FRAME_SIZE, dtype=np.float32)) # 1.5 frames total -> 1 frame
|
||||
assert len(out) == 1
|
||||
assert stub.steps == 1
|
||||
|
||||
|
||||
def test_flush_pads_the_tail():
|
||||
stub = StubStepper()
|
||||
s = PersonaPlexSession(stub, PersonaPlexConfig())
|
||||
s.open()
|
||||
s.feed(np.ones(FRAME_SIZE + 10, dtype=np.float32)) # 1 frame + 10-sample tail
|
||||
assert stub.steps == 1
|
||||
tail = s.flush()
|
||||
assert len(tail) == 1 and stub.steps == 2
|
||||
assert s.flush() == [] # nothing left
|
||||
|
||||
|
||||
def test_run_concatenates_audio_and_joins_text():
|
||||
s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
|
||||
s.open()
|
||||
audio, text = s.run(np.zeros(FRAME_SIZE * 4, dtype=np.float32))
|
||||
assert audio.shape[0] == FRAME_SIZE * 4
|
||||
assert text == "t1t3" # frames 1..4, text on odds
|
||||
|
||||
|
||||
def test_open_forwards_voice_and_persona():
|
||||
stub = StubStepper()
|
||||
s = PersonaPlexSession(stub, PersonaPlexConfig())
|
||||
s.open(voice_prompt="NATM1.pt", persona="pirate")
|
||||
assert stub.opened == ("NATM1.pt", "pirate")
|
||||
|
||||
|
||||
def test_feed_before_open_raises():
|
||||
s = PersonaPlexSession(StubStepper(), PersonaPlexConfig())
|
||||
with pytest.raises(RuntimeError):
|
||||
s.feed(np.zeros(FRAME_SIZE, dtype=np.float32))
|
||||
Loaded 100 of 190 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user